mediocre-go-lib/mrun/hook.go

package mrun

import (
	"context"

	"github.com/mediocregopher/mediocre-go-lib/mcmp"
)

// Hook describes a function which can be registered to trigger on an event via
// the WithHook function.
type Hook func(context.Context) error

type hookKey struct {
	key interface{}
}

// AddHook registers a Hook under a typed key. The Hook will be called when
// TriggerHooks is called with that same key. Multiple Hooks can be registered
// for the same key, and will be called sequentially when triggered.
//
// Hooks will be called with whatever Context is passed into TriggerHooks.
func AddHook(cmp *mcmp.Component, key interface{}, hook Hook) {
	mcmp.AddSeriesValue(cmp, hookKey{key}, hook)
}

func triggerHooks(
	ctx context.Context,
	cmp *mcmp.Component,
	key interface{},
	next func([]mcmp.SeriesElement) (mcmp.SeriesElement, []mcmp.SeriesElement),
) error {
	els := mcmp.SeriesElements(cmp, hookKey{key})
	var el mcmp.SeriesElement
	for {
		if len(els) == 0 {
			break
		}
		el, els = next(els)
		if el.Child != nil {
			if err := triggerHooks(ctx, el.Child, key, next); err != nil {
				return err
			}
			continue
		}
		hook := el.Value.(Hook)
		if err := hook(ctx); err != nil {
			return err
		}
	}
	return nil
}

// TriggerHooks causes all Hooks registered with AddHook on the Component under
// the given key to be called in the order they were registered. The given
// Context is passed into all Hooks being called.
//
// If any Hook returns an error no further Hooks will be called and that error
// will be returned.
//
// If the Component has children (see the mcmp package), and those children have
// Hooks registered under this key, then their Hooks will be called in the
// expected order. See package docs for an example.
func TriggerHooks(
	ctx context.Context,
	cmp *mcmp.Component,
	key interface{},
) error {
	next := func(els []mcmp.SeriesElement) (
		mcmp.SeriesElement, []mcmp.SeriesElement,
	) {
		return els[0], els[1:]
	}
	return triggerHooks(ctx, cmp, key, next)
}

// TriggerHooksReverse is the same as TriggerHooks except that registered Hooks
// are called in the reverse order in which they were registered.
func TriggerHooksReverse(ctx context.Context, cmp *mcmp.Component, key interface{}) error {
	next := func(els []mcmp.SeriesElement) (
		mcmp.SeriesElement, []mcmp.SeriesElement,
	) {
		last := len(els) - 1
		return els[last], els[:last]
	}
	return triggerHooks(ctx, cmp, key, next)
}

type builtinEvent int

const (
	initEvent builtinEvent = iota
	shutdownEvent
)

// InitHook registers the given Hook to run when Init is called. This is a
// special case of AddHook.
//
// As a convention Hooks running on the init event should block only as long as
// it takes to ensure that whatever is running can do so successfully. For
// short-lived tasks this isn't a problem, but long-lived tasks (e.g. a web
// server) will want to use the Hook only to initialize, and spawn off a
// go-routine to do their actual work. Long-lived tasks should set themselves up
// to shutdown on the shutdown event (see ShutdownHook).
func InitHook(cmp *mcmp.Component, hook Hook) {
	AddHook(cmp, initEvent, hook)
}

// Init runs all Hooks registered using InitHook. This is a special case of
// TriggerHooks.
func Init(ctx context.Context, cmp *mcmp.Component) error {
	return TriggerHooks(ctx, cmp, initEvent)
}

// ShutdownHook registers the given Hook to run when Shutdown is called. This is
// a special case of AddHook.
//
// See InitHook for more on the relationship between Init(Hook) and
// Shutdown(Hook).
func ShutdownHook(cmp *mcmp.Component, hook Hook) {
	AddHook(cmp, shutdownEvent, hook)
}

// Shutdown runs all Hooks registered using ShutdownHook in the reverse order in
// which they were registered. This is a special case of TriggerHooks.
func Shutdown(ctx context.Context, cmp *mcmp.Component) error {
	return TriggerHooksReverse(ctx, cmp, shutdownEvent)
}