From e8c572c45b4d028cfc0c608699434bc8e8808c67 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 5 Nov 2019 18:57:25 +1100
Subject: Cache ContextDoneVars

Note that nothing currently flushes them. Could probably flush them when they're done, and the early return will take care of the rest.
---
 stmutil/context.go      | 23 ++++++++++++++++++-----
 stmutil/context_test.go | 20 ++++++++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)
 create mode 100644 stmutil/context_test.go

(limited to 'stmutil')

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() {}
 }
diff --git a/stmutil/context_test.go b/stmutil/context_test.go
new file mode 100644
index 0000000..0a6b7c0
--- /dev/null
+++ b/stmutil/context_test.go
@@ -0,0 +1,20 @@
+package stmutil
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestContextEquality(t *testing.T) {
+	ctx := context.Background()
+	assert.True(t, ctx == context.Background())
+	childCtx, cancel := context.WithCancel(ctx)
+	assert.True(t, childCtx != ctx)
+	assert.True(t, childCtx != ctx)
+	assert.Equal(t, context.Background(), ctx)
+	cancel()
+	assert.Equal(t, context.Background(), ctx)
+	assert.NotEqual(t, ctx, childCtx)
+}
-- 
cgit v1.2.3