github.com/diamondburned/arikawa@v1.3.14/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  type Pacemaker struct {
    35  	// Heartrate is the received duration between heartbeats.
    36  	Heartrate time.Duration
    37  
    38  	ticker time.Ticker
    39  	Ticks  <-chan time.Time
    40  
    41  	// Time in nanoseconds, guarded by atomic read/writes.
    42  	SentBeat AtomicTime
    43  	EchoBeat AtomicTime
    44  
    45  	// Any callback that returns an error will stop the pacer.
    46  	Pacer func(context.Context) error
    47  }
    48  
    49  func NewPacemaker(heartrate time.Duration, pacer func(context.Context) error) Pacemaker {
    50  	p := Pacemaker{
    51  		Heartrate: heartrate,
    52  		Pacer:     pacer,
    53  		ticker:    *time.NewTicker(heartrate),
    54  	}
    55  	p.Ticks = p.ticker.C
    56  	// Reset states to its old position.
    57  	now := time.Now()
    58  	p.EchoBeat.Set(now)
    59  	p.SentBeat.Set(now)
    60  
    61  	return p
    62  }
    63  
    64  func (p *Pacemaker) Echo() {
    65  	// Swap our received heartbeats
    66  	p.EchoBeat.Set(time.Now())
    67  }
    68  
    69  // Dead, if true, will have Pace return an ErrDead.
    70  func (p *Pacemaker) Dead() bool {
    71  	var (
    72  		echo = p.EchoBeat.Get()
    73  		sent = p.SentBeat.Get()
    74  	)
    75  
    76  	if echo == 0 || sent == 0 {
    77  		return false
    78  	}
    79  
    80  	return sent-echo > int64(p.Heartrate)*2
    81  }
    82  
    83  // Stop stops the pacemaker, or it does nothing if the pacemaker is not started.
    84  func (p *Pacemaker) StopTicker() {
    85  	p.ticker.Stop()
    86  }
    87  
    88  // pace sends a heartbeat with the appropriate timeout for the context.
    89  func (p *Pacemaker) Pace() error {
    90  	ctx, cancel := context.WithTimeout(context.Background(), p.Heartrate)
    91  	defer cancel()
    92  
    93  	return p.PaceCtx(ctx)
    94  }
    95  
    96  func (p *Pacemaker) PaceCtx(ctx context.Context) error {
    97  	if err := p.Pacer(ctx); err != nil {
    98  		return err
    99  	}
   100  
   101  	p.SentBeat.Set(time.Now())
   102  
   103  	if p.Dead() {
   104  		return ErrDead
   105  	}
   106  
   107  	return nil
   108  }