125 lines
3.2 KiB
Go
125 lines
3.2 KiB
Go
// Package routines provides utilities for managing long-running goroutines.
|
|
package routines
|
|
|
|
import "io"
|
|
import "fmt"
|
|
import "log"
|
|
import "time"
|
|
import "sync"
|
|
import "errors"
|
|
import "context"
|
|
|
|
// Func is a routine created from a run function.
|
|
type Func func (context.Context) error
|
|
|
|
// Run runs the routine.
|
|
func (fun Func) Run (ctx context.Context) error {
|
|
return fun(ctx)
|
|
}
|
|
|
|
// Routine is an object that can be run and stopped.
|
|
type Routine interface {
|
|
// Run is a long-running function that does not return until it is
|
|
// finished, or its context is cancelled. If the context is cancelled,
|
|
// the function must perform necessary cleanup/shutdown operations and
|
|
// exit. An error is returned if the routine exited due to an error.
|
|
Run (context.Context) error
|
|
}
|
|
|
|
// Manager is a system capable of managing multiple routines, and restarting
|
|
// them if they fail.
|
|
type Manager struct {
|
|
// Routines specifies a list of routines to manage. These are started
|
|
// when Run() is called.
|
|
Routines []Routine
|
|
|
|
// RestartDeadline specifies the amount of time a routine has to be
|
|
// running before failing to be restarted. This is to prevent routines
|
|
// that immediately fail from just being restarted over and over again.
|
|
RestartDeadline time.Duration
|
|
|
|
// Logger, if non-nil, is where log messages will be written to. If it
|
|
// is nil, messages will be written to the standard logger. To disable
|
|
// logging altogether, this can be set to io.Discard.
|
|
Logger io.Writer
|
|
|
|
ctx context.Context
|
|
}
|
|
|
|
// Run spawns all routines in the Routines slice. If a routine exits with an
|
|
// error and it was running for longer than RestartDeadline, it is restarted.
|
|
// Run returns only when all routines have exited.
|
|
func (this *Manager) Run (ctx context.Context) error {
|
|
ctx, done := context.WithCancel(ctx)
|
|
this.ctx = ctx
|
|
|
|
var waitGroup sync.WaitGroup
|
|
for _, routine := range this.Routines {
|
|
if routine != nil {
|
|
waitGroup.Add(1)
|
|
go this.runRoutine(routine, &waitGroup)
|
|
}
|
|
}
|
|
waitGroup.Wait()
|
|
|
|
done()
|
|
return nil
|
|
}
|
|
|
|
// Append adds one or more routines to the Routines slice. This has no effect if
|
|
// the manager is already running.
|
|
func (this *Manager) Append (routines ...Routine) {
|
|
this.Routines = append(this.Routines, routines...)
|
|
}
|
|
|
|
func (this *Manager) log (message ...any) {
|
|
if this.Logger == nil {
|
|
log.Println(message...)
|
|
} else {
|
|
fmt.Fprintln(this.Logger, message...)
|
|
}
|
|
}
|
|
|
|
func (this *Manager) runRoutine (routine Routine, group *sync.WaitGroup) {
|
|
defer group.Done()
|
|
|
|
for {
|
|
lastStart := time.Now()
|
|
err := panicWrap(routine.Run, this.ctx)
|
|
|
|
if ctxErr := this.ctx.Err(); ctxErr != nil {
|
|
if err == nil {
|
|
this.log("(i) stopped routine")
|
|
} else {
|
|
this.log("!!! stopped routine, with error:", err)
|
|
}
|
|
break
|
|
}
|
|
|
|
if err == nil {
|
|
this.log("(i) routine exited")
|
|
break
|
|
} else {
|
|
this.log("XXX routine failed:", err)
|
|
}
|
|
|
|
if time.Since(lastStart) < this.RestartDeadline {
|
|
this.log("!!! not restarting routine, failed too soon")
|
|
break
|
|
} else {
|
|
this.log("(i) routine is being restarted")
|
|
}
|
|
}
|
|
}
|
|
|
|
func panicWrap (f func (context.Context) error, ctx context.Context) (err error) {
|
|
defer func () {
|
|
if pan := recover(); pan != nil {
|
|
err = errors.New(fmt.Sprint(pan))
|
|
}
|
|
} ()
|
|
|
|
err = f(ctx)
|
|
return
|
|
}
|