github.com/ngicks/gokugen@v0.0.5/context.go (about)

     1  package gokugen
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"time"
     8  )
     9  
    10  var (
    11  	ErrValueNotFound = errors.New("value not found")
    12  )
    13  
    14  type WorkFn = func(taskCtx context.Context, scheduled time.Time) (any, error)
    15  type WorkFnWParam = func(taskCtx context.Context, scheduled time.Time, param any) (any, error)
    16  
    17  // SchedulerContext is minimal set of data relevant to scheduling and middlewares.
    18  type SchedulerContext interface {
    19  	ScheduledTime() time.Time
    20  	Work() WorkFn
    21  	Value(key any) (any, error)
    22  }
    23  
    24  type (
    25  	paramKeyTy  string
    26  	taskIdKeyTy string
    27  	workIdKeyTy string
    28  )
    29  
    30  func (s paramKeyTy) String() string {
    31  	return "paramKeyTy"
    32  }
    33  func (s taskIdKeyTy) String() string {
    34  	return "taskIdKeyTy"
    35  }
    36  func (s workIdKeyTy) String() string {
    37  	return "workIdKeyTy"
    38  }
    39  
    40  var (
    41  	paramKey  *paramKeyTy  = new(paramKeyTy)
    42  	taskIdKey *taskIdKeyTy = new(taskIdKeyTy)
    43  	workIdKey *workIdKeyTy = new(workIdKeyTy)
    44  )
    45  
    46  // PlainContext is intended to be a base context of SchedulerContext.
    47  type PlainContext struct {
    48  	scheduledTime time.Time
    49  	workFn        WorkFn
    50  	values        map[any]any
    51  }
    52  
    53  // NewPlainContext creates a new PlainContext instance.
    54  // But recommendation here is to use BuildContext instead.
    55  func NewPlainContext(scheduledTime time.Time, workFn WorkFn, values map[any]any) SchedulerContext {
    56  	return &PlainContext{
    57  		scheduledTime: scheduledTime,
    58  		workFn:        workFn,
    59  		values:        values,
    60  	}
    61  }
    62  
    63  func (ctx *PlainContext) ScheduledTime() time.Time {
    64  	return ctx.scheduledTime
    65  }
    66  func (ctx *PlainContext) Work() WorkFn {
    67  	return ctx.workFn
    68  }
    69  func (ctx *PlainContext) Value(key any) (any, error) {
    70  	if ctx.values == nil {
    71  		return nil, nil
    72  	}
    73  	return ctx.values[key], nil
    74  }
    75  
    76  func WrapWithWorkId(parent SchedulerContext, workId string) SchedulerContext {
    77  	return &workIdCtx{
    78  		SchedulerContext: parent,
    79  		workId:           workId,
    80  	}
    81  }
    82  
    83  type workIdCtx struct {
    84  	SchedulerContext
    85  	workId string
    86  }
    87  
    88  func (ctx *workIdCtx) Value(key any) (any, error) {
    89  	if key == workIdKey {
    90  		return ctx.workId, nil
    91  	}
    92  	return ctx.SchedulerContext.Value(key)
    93  }
    94  
    95  func WrapWithParam(parent SchedulerContext, param any) SchedulerContext {
    96  	return &paramLoadableCtx{
    97  		SchedulerContext: parent,
    98  		paramLoader:      func() (any, error) { return param, nil },
    99  	}
   100  }
   101  
   102  func WrapWithParamLoader(parent SchedulerContext, loader func() (any, error)) SchedulerContext {
   103  	return &paramLoadableCtx{
   104  		SchedulerContext: parent,
   105  		paramLoader:      loader,
   106  	}
   107  }
   108  
   109  type paramLoadableCtx struct {
   110  	SchedulerContext
   111  	paramLoader func() (any, error)
   112  }
   113  
   114  func (ctx *paramLoadableCtx) Value(key any) (any, error) {
   115  	if key == paramKey {
   116  		return ctx.paramLoader()
   117  	}
   118  	return ctx.SchedulerContext.Value(key)
   119  }
   120  
   121  func WrapWithWorkFn(parent SchedulerContext, workFn WorkFn) SchedulerContext {
   122  	return &fnWrapperCtx{
   123  		SchedulerContext: parent,
   124  		wrapper: func(self SchedulerContext, _ WorkFn) WorkFn {
   125  			return workFn
   126  		},
   127  	}
   128  }
   129  
   130  type WorkFnWrapper = func(self SchedulerContext, workFn WorkFn) WorkFn
   131  
   132  func WrapWithWorkFnWrapper(parent SchedulerContext, wrapper WorkFnWrapper) SchedulerContext {
   133  	return &fnWrapperCtx{
   134  		SchedulerContext: parent,
   135  		wrapper:          wrapper,
   136  	}
   137  }
   138  
   139  type fnWrapperCtx struct {
   140  	SchedulerContext
   141  	wrapper func(self SchedulerContext, workFn WorkFn) WorkFn
   142  }
   143  
   144  func (ctx *fnWrapperCtx) Work() WorkFn {
   145  	return ctx.wrapper(ctx, ctx.SchedulerContext.Work())
   146  }
   147  
   148  func WrapWithTaskId(parent SchedulerContext, taskId string) SchedulerContext {
   149  	return &taskIdCtx{
   150  		SchedulerContext: parent,
   151  		taskId:           taskId,
   152  	}
   153  }
   154  
   155  type taskIdCtx struct {
   156  	SchedulerContext
   157  	taskId string
   158  }
   159  
   160  func (ctx *taskIdCtx) Value(key any) (any, error) {
   161  	if key == taskIdKey {
   162  		return ctx.taskId, nil
   163  	}
   164  	return ctx.SchedulerContext.Value(key)
   165  }
   166  
   167  func GetParam(ctx SchedulerContext) (any, error) {
   168  	return ctx.Value(paramKey)
   169  }
   170  
   171  // GetTaskId gets task id from ctx.
   172  // This may be heavy or cause error.
   173  // If task id is not set, GetTaskId returns a wrapped ErrValueNotFound.
   174  func GetTaskId(ctx SchedulerContext) (string, error) {
   175  	id, err := ctx.Value(taskIdKey)
   176  	if err != nil {
   177  		return "", err
   178  	}
   179  	if id == nil {
   180  		return "", fmt.Errorf("%w: key=%s", ErrValueNotFound, taskIdKey)
   181  	}
   182  	return id.(string), nil
   183  }
   184  
   185  // GetWorkId gets task id from ctx.
   186  // This may be heavy or cause error.
   187  // If work id is not set, GetWorkId returns a wrapped ErrValueNotFound.
   188  func GetWorkId(ctx SchedulerContext) (string, error) {
   189  	id, err := ctx.Value(workIdKey)
   190  	if err != nil {
   191  		return "", err
   192  	}
   193  	if id == nil {
   194  		return "", fmt.Errorf("%w: key=%s", ErrValueNotFound, workIdKey)
   195  	}
   196  	return id.(string), nil
   197  }