Add cleanup callbacks to context

1 files changed, 10 insertions, 0 deletions

diff --git a/context.go b/context.go
index dab603e8e..dcfb09f1f 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