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

     1  package scheduler
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	atomicparam "github.com/ngicks/type-param-common/sync-param/atomic-param"
     9  )
    10  
    11  type WorkFn = func(ctx context.Context, scheduled time.Time)
    12  
    13  // Task is simple set of data, which is consist of
    14  // scheduledTime, work, internal calling state and context canceller.
    15  //
    16  // work will be called with contex.Contex which will be closed when task is cancelled or scheduler is torn down.
    17  type Task struct {
    18  	scheduledTime time.Time
    19  	work          WorkFn
    20  	isDone        uint32
    21  	isCancelled   uint32
    22  	cancelCtx     atomicparam.Value[*context.CancelFunc]
    23  }
    24  
    25  // NewTask creates a new Task instance.
    26  // scheduledTime is scheduled time when work should be invoked.
    27  // work is work of Task, this will be only called once.
    28  func NewTask(scheduledTime time.Time, work WorkFn) *Task {
    29  	cancelCtx := atomicparam.NewValue[*context.CancelFunc]()
    30  	cancelCtx.Store(nil)
    31  	return &Task{
    32  		scheduledTime: scheduledTime,
    33  		work:          work,
    34  		cancelCtx:     cancelCtx,
    35  	}
    36  }
    37  
    38  func (t *Task) Do(ctx context.Context) {
    39  	innerCtx, cancel := context.WithCancel(ctx)
    40  	defer cancel()
    41  	if t.IsCancelled() {
    42  		return
    43  	}
    44  	if atomic.CompareAndSwapUint32(&t.isDone, 0, 1) && t.work != nil {
    45  		t.cancelCtx.Store(&cancel)
    46  		// forget cancel immediately in case where Task is held long time after it is done.
    47  		defer t.cancelCtx.Store(nil)
    48  		// in case of race condition.
    49  		// Cancel might be called right between cancelFunc storage and above IsCancelled call.
    50  		if t.IsCancelled() {
    51  			return
    52  		}
    53  		select {
    54  		case <-ctx.Done():
    55  			// Fast path: ctx is already cancelled.
    56  			return
    57  		default:
    58  		}
    59  		t.work(innerCtx, t.scheduledTime)
    60  	}
    61  }
    62  
    63  func (t *Task) GetScheduledTime() time.Time {
    64  	return t.scheduledTime
    65  }
    66  
    67  func (t *Task) Cancel() (cancelled bool) {
    68  	cancelled = atomic.CompareAndSwapUint32(&t.isCancelled, 0, 1)
    69  	if !cancelled {
    70  		return
    71  	}
    72  	if cancel := t.cancelCtx.Load(); cancel != nil {
    73  		(*cancel)()
    74  	}
    75  	return
    76  }
    77  
    78  func (t *Task) CancelWithReason(err error) (cancelled bool) {
    79  	return t.Cancel()
    80  }
    81  
    82  func (t *Task) IsCancelled() bool {
    83  	return atomic.LoadUint32(&t.isCancelled) == 1
    84  }
    85  
    86  func (t *Task) IsDone() bool {
    87  	return atomic.LoadUint32(&t.isDone) == 1
    88  }
    89  
    90  // TaskController is a small wrapper around Task.
    91  // Simply it removes Do method from Task and
    92  // expose other methods as it delegates them to inner Task.
    93  type TaskController struct {
    94  	t *Task
    95  }
    96  
    97  func NewTaskController(t *Task) *TaskController {
    98  	return &TaskController{
    99  		t: t,
   100  	}
   101  }
   102  
   103  func (t *TaskController) GetScheduledTime() time.Time {
   104  	return t.t.GetScheduledTime()
   105  }
   106  
   107  func (t *TaskController) Cancel() (cancelled bool) {
   108  	return t.t.Cancel()
   109  }
   110  
   111  func (t *TaskController) CancelWithReason(err error) (cancelled bool) {
   112  	return t.Cancel()
   113  }
   114  
   115  func (t *TaskController) IsCancelled() bool {
   116  	return t.t.IsCancelled()
   117  }
   118  
   119  func (t *TaskController) IsDone() bool {
   120  	return t.t.IsDone()
   121  }