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

     1  package gokugen
     2  
     3  //go:generate mockgen -source scheduler.go -destination __mock/scheduler.go
     4  
     5  import (
     6  	"context"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/ngicks/gokugen/scheduler"
    11  )
    12  
    13  type Scheduler interface {
    14  	Schedule(task *scheduler.Task) (*scheduler.TaskController, error)
    15  }
    16  
    17  var _ Task = &scheduler.TaskController{}
    18  
    19  type Task interface {
    20  	Cancel() (cancelled bool)
    21  	CancelWithReason(err error) (cancelled bool)
    22  	GetScheduledTime() time.Time
    23  	IsCancelled() bool
    24  	IsDone() bool
    25  }
    26  
    27  type ScheduleHandlerFn = func(ctx SchedulerContext) (Task, error)
    28  
    29  type MiddlewareFunc = func(handler ScheduleHandlerFn) ScheduleHandlerFn
    30  
    31  type MiddlewareApplicator[T Scheduler] struct {
    32  	scheduler T
    33  	mwMu      sync.Mutex
    34  	mw        []MiddlewareFunc
    35  }
    36  
    37  func NewMiddlewareApplicator[T Scheduler](scheduler T) *MiddlewareApplicator[T] {
    38  	return &MiddlewareApplicator[T]{
    39  		scheduler: scheduler,
    40  		mw:        make([]MiddlewareFunc, 0),
    41  	}
    42  }
    43  
    44  // Schedule schedules ctx to inner scheduler with middlewares applied.
    45  // Middlewares will be called in first-in-first-applied order.
    46  func (s *MiddlewareApplicator[T]) Schedule(ctx SchedulerContext) (Task, error) {
    47  	schedule := s.apply(func(ctx SchedulerContext) (Task, error) {
    48  		return s.scheduler.Schedule(
    49  			scheduler.NewTask(
    50  				ctx.ScheduledTime(),
    51  				func(
    52  					taskCtx context.Context,
    53  					scheduled time.Time,
    54  				) {
    55  					ctx.Work()(taskCtx, scheduled)
    56  				},
    57  			),
    58  		)
    59  	})
    60  	return schedule(ctx)
    61  }
    62  
    63  // Scheduler is getter of inner sheculer.
    64  func (ma *MiddlewareApplicator[T]) Scheduler() T {
    65  	return ma.scheduler
    66  }
    67  
    68  // Use registers MiddlewareFunc.
    69  // First registered one will be invoked first.
    70  func (s *MiddlewareApplicator[T]) Use(mw ...MiddlewareFunc) {
    71  	s.mwMu.Lock()
    72  	defer s.mwMu.Unlock()
    73  
    74  	s.mw = append(s.mw, mw...)
    75  }
    76  
    77  func (s *MiddlewareApplicator[T]) apply(schedule ScheduleHandlerFn) ScheduleHandlerFn {
    78  	s.mwMu.Lock()
    79  	defer s.mwMu.Unlock()
    80  
    81  	var wrapped ScheduleHandlerFn
    82  	wrapped = schedule
    83  	for i := len(s.mw) - 1; i >= 0; i-- {
    84  		handlerFunc := s.mw[i]
    85  		if handlerFunc != nil {
    86  			wrapped = handlerFunc(wrapped)
    87  		}
    88  	}
    89  	return wrapped
    90  }