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

     1  package taskstorage
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/ngicks/gokugen"
    10  )
    11  
    12  var (
    13  	ErrOtherNodeWorkingOnTheTask = errors.New("other node is already working on the task")
    14  )
    15  
    16  type RepositoryUpdater interface {
    17  	Repository
    18  	StateUpdater
    19  }
    20  
    21  // MultiNodeTaskStorage is almost same as SingleNodeTaskStorage but has one additional behavior.
    22  // MultiNodeTaskStorage tries to mark tasks as Working state right before task is being worked on.
    23  // If task is already marked, it fails to do task.
    24  // Multiple nodes can be synced to same data storage through RepositoryUpdater interface.
    25  // And only one node will do task.
    26  type MultiNodeTaskStorage struct {
    27  	repo RepositoryUpdater
    28  	sns  *SingleNodeTaskStorage
    29  }
    30  
    31  func NewMultiNodeTaskStorage(
    32  	repo RepositoryUpdater,
    33  	shouldRestore func(TaskInfo) bool,
    34  	workRegistry WorkRegistry,
    35  ) *MultiNodeTaskStorage {
    36  	return &MultiNodeTaskStorage{
    37  		repo: repo,
    38  		sns:  NewSingleNodeTaskStorage(repo, shouldRestore, workRegistry, nil),
    39  	}
    40  }
    41  
    42  func (m *MultiNodeTaskStorage) markWorking(handler gokugen.ScheduleHandlerFn) gokugen.ScheduleHandlerFn {
    43  	return func(ctx gokugen.SchedulerContext) (gokugen.Task, error) {
    44  		return handler(gokugen.WrapContext(ctx, gokugen.WithWorkFnWrapper(buildWrapper(m.repo))))
    45  	}
    46  }
    47  
    48  func buildWrapper(repo RepositoryUpdater) gokugen.WorkFnWrapper {
    49  	return func(self gokugen.SchedulerContext, workFn gokugen.WorkFn) gokugen.WorkFn {
    50  		return func(ctx context.Context, scheduled time.Time) (any, error) {
    51  			taskId, err := gokugen.GetTaskId(self)
    52  			if err != nil {
    53  				return nil, err
    54  			}
    55  			swapped, err := repo.UpdateState(taskId, Initialized, Working)
    56  			if err != nil {
    57  				return nil, err
    58  			}
    59  			if !swapped {
    60  				return nil, fmt.Errorf("%w: task id = %s", ErrOtherNodeWorkingOnTheTask, taskId)
    61  			}
    62  			return workFn(ctx, scheduled)
    63  		}
    64  	}
    65  }
    66  
    67  func (m *MultiNodeTaskStorage) Middleware(freeParam bool) []gokugen.MiddlewareFunc {
    68  	return append(m.sns.Middleware(freeParam), m.markWorking)
    69  }
    70  
    71  func (m *MultiNodeTaskStorage) Sync(
    72  	schedule func(ctx gokugen.SchedulerContext) (gokugen.Task, error),
    73  ) (rescheduled map[string]gokugen.Task, schedulingErr map[string]error, err error) {
    74  	return m.sns.Sync(schedule)
    75  }