github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/scheduler/timer.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"runtime"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/asynkron/protoactor-go/actor"
     9  )
    10  
    11  type CancelFunc func()
    12  
    13  type Stopper interface {
    14  	Stop()
    15  }
    16  
    17  const (
    18  	stateInit = iota
    19  	stateReady
    20  	stateDone
    21  )
    22  
    23  func startTimer(delay, interval time.Duration, fn func()) CancelFunc {
    24  	var t *time.Timer
    25  	var state int32
    26  	t = time.AfterFunc(delay, func() {
    27  		for atomic.LoadInt32(&state) == stateInit {
    28  			runtime.Gosched()
    29  		}
    30  
    31  		if state == stateDone {
    32  			return
    33  		}
    34  
    35  		fn()
    36  		t.Reset(interval)
    37  	})
    38  
    39  	// ensures t != nil and is required to avoid data race in
    40  	// AfterFunc calling t.Reset
    41  	atomic.StoreInt32(&state, stateReady)
    42  
    43  	return func() {
    44  		if atomic.SwapInt32(&state, stateDone) != stateDone {
    45  			t.Stop()
    46  		}
    47  	}
    48  }
    49  
    50  // A scheduler utilizing timers to send messages in the future and at regular intervals.
    51  type TimerScheduler struct {
    52  	ctx actor.SenderContext
    53  }
    54  
    55  type timerOptionFunc func(*TimerScheduler)
    56  
    57  // WithContext configures the scheduler to use ctx rather than the default,
    58  // EmptyRootContext.
    59  func WithContext(ctx actor.SenderContext) timerOptionFunc {
    60  	return func(s *TimerScheduler) {
    61  		s.ctx = ctx
    62  	}
    63  }
    64  
    65  // NewTimerScheduler creates a new scheduler using the EmptyRootContext.
    66  // Additional options may be specified to override the default behavior.
    67  func NewTimerScheduler(sender actor.SenderContext, opts ...timerOptionFunc) *TimerScheduler {
    68  	s := &TimerScheduler{ctx: sender}
    69  	for _, opt := range opts {
    70  		opt(s)
    71  	}
    72  	return s
    73  }
    74  
    75  // SendOnce waits for the duration to elapse and then calls actor.SenderContext.Send to forward the message to pid.
    76  func (s *TimerScheduler) SendOnce(delay time.Duration, pid *actor.PID, message interface{}) CancelFunc {
    77  	t := time.AfterFunc(delay, func() {
    78  		s.ctx.Send(pid, message)
    79  	})
    80  
    81  	return func() { t.Stop() }
    82  }
    83  
    84  // SendRepeatedly waits for the initial duration to elapse and then calls Send to forward the message to pid
    85  // repeatedly for each interval.
    86  func (s *TimerScheduler) SendRepeatedly(initial, interval time.Duration, pid *actor.PID, message interface{}) CancelFunc {
    87  	return startTimer(initial, interval, func() {
    88  		s.ctx.Send(pid, message)
    89  	})
    90  }
    91  
    92  // RequestOnce waits for the duration to elapse and then calls actor.SenderContext.Request to forward the message to
    93  // pid.
    94  func (s *TimerScheduler) RequestOnce(delay time.Duration, pid *actor.PID, message interface{}) CancelFunc {
    95  	t := time.AfterFunc(delay, func() {
    96  		s.ctx.Request(pid, message)
    97  	})
    98  
    99  	return func() { t.Stop() }
   100  }
   101  
   102  // RequestRepeatedly waits for the initial duration to elapse and then calls Request to forward the message to pid
   103  // repeatedly for each interval.
   104  func (s *TimerScheduler) RequestRepeatedly(delay, interval time.Duration, pid *actor.PID, message interface{}) CancelFunc {
   105  	return startTimer(delay, interval, func() {
   106  		s.ctx.Request(pid, message)
   107  	})
   108  }