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

     1  package deadline
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/ngicks/gokugen"
     9  	"github.com/ngicks/gokugen/common"
    10  )
    11  
    12  type DeadlineExeededErr struct {
    13  	scheduled time.Time
    14  	executed  time.Time
    15  }
    16  
    17  func (e DeadlineExeededErr) Error() string {
    18  	return fmt.Sprintf(
    19  		"deadline exeeded: scheduled at=%s, but executed at=%s",
    20  		e.scheduled.Format(time.RFC3339Nano),
    21  		e.executed.Format(time.RFC3339Nano),
    22  	)
    23  }
    24  func (e DeadlineExeededErr) ScheduledTime() time.Time {
    25  	return e.scheduled
    26  }
    27  
    28  func (e DeadlineExeededErr) ExecutedTime() time.Time {
    29  	return e.executed
    30  }
    31  
    32  type DeadlineMiddleware struct {
    33  	deadline   time.Duration
    34  	getNow     common.GetNower
    35  	shouldSkip func(ctx gokugen.SchedulerContext) bool
    36  }
    37  
    38  func New(deadline time.Duration, shouldSkip func(ctx gokugen.SchedulerContext) bool) *DeadlineMiddleware {
    39  	return &DeadlineMiddleware{
    40  		deadline:   deadline,
    41  		shouldSkip: shouldSkip,
    42  		getNow:     common.GetNowImpl{},
    43  	}
    44  }
    45  
    46  func (mw *DeadlineMiddleware) Middleware(handler gokugen.ScheduleHandlerFn) gokugen.ScheduleHandlerFn {
    47  	return func(ctx gokugen.SchedulerContext) (gokugen.Task, error) {
    48  		if mw.shouldSkip(ctx) {
    49  			return handler(ctx)
    50  		}
    51  		return handler(
    52  			gokugen.WrapContext(
    53  				ctx,
    54  				gokugen.WithWorkFnWrapper(wrapper(mw.deadline, mw.getNow)),
    55  			),
    56  		)
    57  	}
    58  }
    59  
    60  func wrapper(
    61  	deadline time.Duration,
    62  	getNow common.GetNower,
    63  ) func(self gokugen.SchedulerContext, workFn gokugen.WorkFn) gokugen.WorkFn {
    64  	return func(self gokugen.SchedulerContext, workFn gokugen.WorkFn) gokugen.WorkFn {
    65  		return func(taskCtx context.Context, scheduled time.Time) (any, error) {
    66  			now := getNow.GetNow()
    67  			if now.Sub(scheduled) > deadline {
    68  				return nil, DeadlineExeededErr{scheduled: scheduled, executed: now}
    69  			}
    70  			return workFn(taskCtx, scheduled)
    71  		}
    72  	}
    73  }