Add cleanup callbacks to context

This commit is contained in:
Matthew Holt 2019-05-29 23:10:12 -06:00
parent 2265db9028
commit 1b6b422c63

View File

@ -24,6 +24,7 @@ type Context struct {
context.Context context.Context
moduleInstances map[string][]interface{} moduleInstances map[string][]interface{}
cfg *Config cfg *Config
cleanupFuncs []func()
} }
// NewContext provides a new context derived from the given // NewContext provides a new context derived from the given
@ -40,6 +41,10 @@ func NewContext(ctx Context) (Context, context.CancelFunc) {
wrappedCancel := func() { wrappedCancel := func() {
cancel() cancel()
for _, f := range ctx.cleanupFuncs {
f()
}
for modName, modInstances := range newCtx.moduleInstances { for modName, modInstances := range newCtx.moduleInstances {
for _, inst := range modInstances { for _, inst := range modInstances {
if cu, ok := inst.(CleanerUpper); ok { if cu, ok := inst.(CleanerUpper); ok {
@ -55,6 +60,11 @@ func NewContext(ctx Context) (Context, context.CancelFunc) {
return newCtx, wrappedCancel return newCtx, wrappedCancel
} }
// OnCancel executes f when ctx is cancelled.
func (ctx *Context) OnCancel(f func()) {
ctx.cleanupFuncs = append(ctx.cleanupFuncs, f)
}
// LoadModule decodes rawMsg into a new instance of mod and // LoadModule decodes rawMsg into a new instance of mod and
// returns the value. If mod.New() does not return a pointer // returns the value. If mod.New() does not return a pointer
// value, it is converted to one so that it is unmarshaled // value, it is converted to one so that it is unmarshaled