diff options
Diffstat (limited to 'stmutil/context.go')
-rw-r--r-- | stmutil/context.go | 23 |
1 files changed, 18 insertions, 5 deletions
diff --git a/stmutil/context.go b/stmutil/context.go index e82e522..98eb8aa 100644 --- a/stmutil/context.go +++ b/stmutil/context.go @@ -2,19 +2,32 @@ package stmutil import ( "context" + "sync" "github.com/anacrolix/stm" ) +var ( + mu sync.Mutex + ctxVars = map[context.Context]*stm.Var{} +) + func ContextDoneVar(ctx context.Context) (*stm.Var, func()) { + mu.Lock() + defer mu.Unlock() + if v, ok := ctxVars[ctx]; ok { + return v, func() {} + } if ctx.Err() != nil { - return stm.NewVar(true), func() {} + v := stm.NewVar(true) + ctxVars[ctx] = v + return v, func() {} } - ctx, cancel := context.WithCancel(ctx) - _var := stm.NewVar(false) + v := stm.NewVar(false) go func() { <-ctx.Done() - stm.AtomicSet(_var, true) + stm.AtomicSet(v, true) }() - return _var, cancel + ctxVars[ctx] = v + return v, func() {} } |