diff --git a/context.go b/context.go index dab603e8..dcfb09f1 100644 --- a/context.go +++ b/context.go @@ -24,6 +24,7 @@ type Context struct { context.Context moduleInstances map[string][]interface{} cfg *Config + cleanupFuncs []func() } // NewContext provides a new context derived from the given @@ -40,6 +41,10 @@ func NewContext(ctx Context) (Context, context.CancelFunc) { wrappedCancel := func() { cancel() + for _, f := range ctx.cleanupFuncs { + f() + } + for modName, modInstances := range newCtx.moduleInstances { for _, inst := range modInstances { if cu, ok := inst.(CleanerUpper); ok { @@ -55,6 +60,11 @@ func NewContext(ctx Context) (Context, context.CancelFunc) { 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 // returns the value. If mod.New() does not return a pointer // value, it is converted to one so that it is unmarshaled