arikawa/gateway/pacemaker.go

package gateway

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

var ErrDead = errors.New("no heartbeat replied")

// Time is a UnixNano timestamp.
type Time = int64

type Pacemaker struct {
	// Heartrate is the received duration between heartbeats.
	Heartrate time.Duration

	// Time in nanoseconds, guarded by atomic read/writes.
	SentBeat Time
	EchoBeat Time

	// Any callback that returns an error will stop the pacer.
	Pace func() error

	stop  chan struct{}
	death chan error
}

func (p *Pacemaker) Echo() {
	// Swap our received heartbeats
	// p.LastBeat[0], p.LastBeat[1] = time.Now(), p.LastBeat[0]
	atomic.StoreInt64(&p.EchoBeat, time.Now().UnixNano())
}

// Dead, if true, will have Pace return an ErrDead.
func (p *Pacemaker) Dead() bool {
	/* Deprecated
	if p.LastBeat[0].IsZero() || p.LastBeat[1].IsZero() {
		return false
	}

	return p.LastBeat[0].Sub(p.LastBeat[1]) > p.Heartrate*2
	*/

	var (
		echo = atomic.LoadInt64(&p.EchoBeat)
		sent = atomic.LoadInt64(&p.SentBeat)
	)

	if echo == 0 || sent == 0 {
		return false
	}

	return sent-echo > int64(p.Heartrate)*2
}

func (p *Pacemaker) Stop() {
	if p.stop != nil {
		p.stop <- struct{}{}
		WSDebug("(*Pacemaker).stop was sent a stop signal.")
	} else {
		WSDebug("(*Pacemaker).stop is nil, skipping.")
	}
}

func (p *Pacemaker) start() error {
	tick := time.NewTicker(p.Heartrate)
	defer tick.Stop()

	// Echo at least once
	p.Echo()

	for {
		WSDebug("Pacemaker loop restarted.")

		if err := p.Pace(); err != nil {
			return err
		}

		WSDebug("Paced.")

		// Paced, save:
		atomic.StoreInt64(&p.SentBeat, time.Now().UnixNano())

		if p.Dead() {
			return ErrDead
		}

		select {
		case <-p.stop:
			WSDebug("Received stop signal.")
			return nil

		case <-tick.C:
			WSDebug("Ticked. Restarting.")
		}
	}
}

// StartAsync starts the pacemaker asynchronously. The WaitGroup is optional.
func (p *Pacemaker) StartAsync(wg *sync.WaitGroup) (death chan error) {
	p.death = make(chan error)
	p.stop = make(chan struct{})

	wg.Add(1)

	go func() {
		p.death <- p.start()
		// Debug.
		WSDebug("Pacemaker returned.")
		// Mark the stop channel as nil, so later Close() calls won't block forever.
		p.stop = nil
		// Mark the pacemaker loop as done.
		wg.Done()
	}()

	return p.death
}