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

     1  package observe
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/ngicks/gokugen"
     8  )
     9  
    10  type ObserveMiddleware struct {
    11  	ctxObserver    func(ctx gokugen.SchedulerContext)
    12  	workFnObserver func(ret any, err error)
    13  }
    14  
    15  func New(ctxObserver func(ctx gokugen.SchedulerContext), workFnObserver func(ret any, err error)) *ObserveMiddleware {
    16  	if ctxObserver == nil {
    17  		ctxObserver = func(ctx gokugen.SchedulerContext) {}
    18  	}
    19  	if workFnObserver == nil {
    20  		workFnObserver = func(ret any, err error) {}
    21  	}
    22  
    23  	return &ObserveMiddleware{
    24  		ctxObserver:    ctxObserver,
    25  		workFnObserver: workFnObserver,
    26  	}
    27  }
    28  
    29  func (mw *ObserveMiddleware) Middleware(handler gokugen.ScheduleHandlerFn) gokugen.ScheduleHandlerFn {
    30  	return func(ctx gokugen.SchedulerContext) (gokugen.Task, error) {
    31  		mw.ctxObserver(ctx)
    32  		return handler(
    33  			gokugen.WrapContext(
    34  				ctx,
    35  				gokugen.WithWorkFnWrapper(
    36  					func(self gokugen.SchedulerContext, workFn gokugen.WorkFn) gokugen.WorkFn {
    37  						return func(taskCtx context.Context, scheduled time.Time) (ret any, err error) {
    38  							ret, err = workFn(taskCtx, scheduled)
    39  							mw.workFnObserver(ret, err)
    40  							return
    41  						}
    42  					},
    43  				),
    44  			),
    45  		)
    46  	}
    47  }