github.com/diamondburned/arikawa/v2@v2.1.0/internal/heart/heart.go (about)

     1  // Package heart implements a general purpose pacemaker.
     2  package heart
     3  
     4  import (
     5  	"context"
     6  	"sync/atomic"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  )
    11  
    12  // Debug is the default logger that Pacemaker uses.
    13  var Debug = func(v ...interface{}) {}
    14  
    15  var ErrDead = errors.New("no heartbeat replied")
    16  
    17  // AtomicTime is a thread-safe UnixNano timestamp guarded by atomic.
    18  type AtomicTime struct {
    19  	unixnano int64
    20  }
    21  
    22  func (t *AtomicTime) Get() int64 {
    23  	return atomic.LoadInt64(&t.unixnano)
    24  }
    25  
    26  func (t *AtomicTime) Set(time time.Time) {
    27  	atomic.StoreInt64(&t.unixnano, time.UnixNano())
    28  }
    29  
    30  func (t *AtomicTime) Time() time.Time {
    31  	return time.Unix(0, t.Get())
    32  }
    33  
    34  // AtomicDuration is a thread-safe Duration guarded by atomic.
    35  type AtomicDuration struct {
    36  	duration int64
    37  }
    38  
    39  func (d *AtomicDuration) Get() time.Duration {
    40  	return time.Duration(atomic.LoadInt64(&d.duration))
    41  }
    42  
    43  func (d *AtomicDuration) Set(dura time.Duration) {
    44  	atomic.StoreInt64(&d.duration, int64(dura))
    45  }
    46  
    47  // Pacemaker is the internal pacemaker state. All fields are not thread-safe
    48  // unless they're atomic.
    49  type Pacemaker struct {
    50  	// Heartrate is the received duration between heartbeats.
    51  	Heartrate AtomicDuration
    52  
    53  	ticker time.Ticker
    54  	Ticks  <-chan time.Time
    55  
    56  	// Time in nanoseconds, guarded by atomic read/writes.
    57  	SentBeat AtomicTime
    58  	EchoBeat AtomicTime
    59  
    60  	// Any callback that returns an error will stop the pacer.
    61  	Pacer func(context.Context) error
    62  }
    63  
    64  func NewPacemaker(heartrate time.Duration, pacer func(context.Context) error) Pacemaker {
    65  	p := Pacemaker{
    66  		Heartrate: AtomicDuration{int64(heartrate)},
    67  		Pacer:     pacer,
    68  		ticker:    *time.NewTicker(heartrate),
    69  	}
    70  	p.Ticks = p.ticker.C
    71  	// Reset states to its old position.
    72  	now := time.Now()
    73  	p.EchoBeat.Set(now)
    74  	p.SentBeat.Set(now)
    75  
    76  	return p
    77  }
    78  
    79  func (p *Pacemaker) Echo() {
    80  	// Swap our received heartbeats
    81  	p.EchoBeat.Set(time.Now())
    82  }
    83  
    84  // Dead, if true, will have Pace return an ErrDead.
    85  func (p *Pacemaker) Dead() bool {
    86  	var (
    87  		echo = p.EchoBeat.Get()
    88  		sent = p.SentBeat.Get()
    89  	)
    90  
    91  	if echo == 0 || sent == 0 {
    92  		return false
    93  	}
    94  
    95  	return sent-echo > int64(p.Heartrate.Get())*2
    96  }
    97  
    98  // SetHeartRate sets the ticker's heart rate.
    99  func (p *Pacemaker) SetPace(heartrate time.Duration) {
   100  	p.Heartrate.Set(heartrate)
   101  
   102  	// To uncomment when 1.16 releases and we drop support for 1.14.
   103  	// p.ticker.Reset(heartrate)
   104  
   105  	p.ticker.Stop()
   106  	p.ticker = *time.NewTicker(heartrate)
   107  	p.Ticks = p.ticker.C
   108  }
   109  
   110  // Stop stops the pacemaker, or it does nothing if the pacemaker is not started.
   111  func (p *Pacemaker) StopTicker() {
   112  	p.ticker.Stop()
   113  }
   114  
   115  // pace sends a heartbeat with the appropriate timeout for the context.
   116  func (p *Pacemaker) Pace() error {
   117  	ctx, cancel := context.WithTimeout(context.Background(), p.Heartrate.Get())
   118  	defer cancel()
   119  
   120  	return p.PaceCtx(ctx)
   121  }
   122  
   123  func (p *Pacemaker) PaceCtx(ctx context.Context) error {
   124  	if err := p.Pacer(ctx); err != nil {
   125  		return err
   126  	}
   127  
   128  	p.SentBeat.Set(time.Now())
   129  
   130  	if p.Dead() {
   131  		return ErrDead
   132  	}
   133  
   134  	return nil
   135  }