diff options
Diffstat (limited to 'stmutil')
-rw-r--r-- | stmutil/context.go | 8 |
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() {} |