Rework ContextDoneVar to not leak done Contexts

1 files changed, 7 insertions, 1 deletions

diff --git a/stmutil/context.go b/stmutil/context.go
index 98eb8aa..315abf4 100644
--- a/stmutil/context.go
+++ b/stmutil/context.go
@@ -12,6 +12,8 @@ var (
 	ctxVars = map[context.Context]*stm.Var{}
 )
 
+// Returns an STM var that contains a bool equal to `ctx.Err != nil`, and a cancel function to be
+// called when the user is no longer interested in the var.
 func ContextDoneVar(ctx context.Context) (*stm.Var, func()) {
 	mu.Lock()
 	defer mu.Unlock()
@@ -19,14 +21,18 @@ func ContextDoneVar(ctx context.Context) (*stm.Var, func()) {
 		return v, func() {}
 	}
 	if ctx.Err() != nil {
+		// TODO: What if we had read-only Vars? Then we could have a global one for this that we
+		// just reuse.
 		v := stm.NewVar(true)
-		ctxVars[ctx] = v
 		return v, func() {}
 	}
 	v := stm.NewVar(false)
 	go func() {
 		<-ctx.Done()
 		stm.AtomicSet(v, true)
+		mu.Lock()
+		delete(ctxVars, ctx)
+		mu.Unlock()
 	}()
 	ctxVars[ctx] = v
 	return v, func() {}