github.com/diamondburned/arikawa@v1.3.14/utils/wsutil/heart.go (about)

     1  package wsutil
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/diamondburned/arikawa/internal/heart"
     8  	"github.com/pkg/errors"
     9  )
    10  
    11  type errBrokenConnection struct {
    12  	underneath error
    13  }
    14  
    15  // Error formats the broken connection error with the message "explicit
    16  // connection break."
    17  func (err errBrokenConnection) Error() string {
    18  	return "explicit connection break: " + err.underneath.Error()
    19  }
    20  
    21  // Unwrap returns the underlying error.
    22  func (err errBrokenConnection) Unwrap() error {
    23  	return err.underneath
    24  }
    25  
    26  // ErrBrokenConnection marks the given error as a broken connection error. This
    27  // error will cause the pacemaker loop to break and return the error. The error,
    28  // when stringified, will say "explicit connection break."
    29  func ErrBrokenConnection(err error) error {
    30  	return errBrokenConnection{underneath: err}
    31  }
    32  
    33  // IsBrokenConnection returns true if the error is a broken connection error.
    34  func IsBrokenConnection(err error) bool {
    35  	var broken *errBrokenConnection
    36  	return errors.As(err, &broken)
    37  }
    38  
    39  // TODO API
    40  type EventLoopHandler interface {
    41  	EventHandler
    42  	HeartbeatCtx(context.Context) error
    43  }
    44  
    45  // PacemakerLoop provides an event loop with a pacemaker. A zero-value instance
    46  // is a valid instance only when RunAsync is called first.
    47  type PacemakerLoop struct {
    48  	heart.Pacemaker
    49  	Extras   ExtraHandlers
    50  	ErrorLog func(error)
    51  
    52  	events  <-chan Event
    53  	handler func(*OP) error
    54  }
    55  
    56  func (p *PacemakerLoop) errorLog(err error) {
    57  	if p.ErrorLog == nil {
    58  		WSDebug("Uncaught error:", err)
    59  		return
    60  	}
    61  
    62  	p.ErrorLog(err)
    63  }
    64  
    65  // Pace calls the pacemaker's Pace function.
    66  func (p *PacemakerLoop) Pace(ctx context.Context) error {
    67  	return p.Pacemaker.PaceCtx(ctx)
    68  }
    69  
    70  func (p *PacemakerLoop) RunAsync(
    71  	heartrate time.Duration, evs <-chan Event, evl EventLoopHandler, exit func(error)) {
    72  
    73  	WSDebug("Starting the pacemaker loop.")
    74  
    75  	p.Pacemaker = heart.NewPacemaker(heartrate, evl.HeartbeatCtx)
    76  	p.handler = evl.HandleOP
    77  	p.events = evs
    78  
    79  	go func() { exit(p.startLoop()) }()
    80  }
    81  
    82  func (p *PacemakerLoop) startLoop() error {
    83  	defer WSDebug("Pacemaker loop has exited.")
    84  	defer p.Pacemaker.StopTicker()
    85  
    86  	for {
    87  		select {
    88  		case <-p.Pacemaker.Ticks:
    89  			if err := p.Pacemaker.Pace(); err != nil {
    90  				return errors.Wrap(err, "pace failed, reconnecting")
    91  			}
    92  
    93  		case ev, ok := <-p.events:
    94  			if !ok {
    95  				WSDebug("Events channel closed, stopping pacemaker.")
    96  				return nil
    97  			}
    98  
    99  			if ev.Error != nil {
   100  				return errors.Wrap(ev.Error, "event returned error")
   101  			}
   102  
   103  			o, err := DecodeOP(ev)
   104  			if err != nil {
   105  				return errors.Wrap(err, "failed to decode OP")
   106  			}
   107  
   108  			// Check the events before handling.
   109  			p.Extras.Check(o)
   110  
   111  			// Handle the event
   112  			if err := p.handler(o); err != nil {
   113  				if IsBrokenConnection(err) {
   114  					return errors.Wrap(err, "handler failed")
   115  				}
   116  
   117  				p.errorLog(err)
   118  			}
   119  		}
   120  	}
   121  }