From 1b6b422c638532d49e697242d9fcf1aa0c3fdc53 Mon Sep 17 00:00:00 2001
From: Matthew Holt <mholt@users.noreply.github.com>
Date: Wed, 29 May 2019 23:10:12 -0600
Subject: [PATCH] Add cleanup callbacks to context

---
 context.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

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