From cb65787f37ed8c14c10db93fa207128186bbd3b7 Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Sun, 28 Oct 2018 22:17:33 -0400
Subject: [PATCH] mctx: GetSetMutableValue. Not sure if I'm really going to use
 this or not

---
 mctx/ctx.go      | 71 ++++++++++++++++++++++++++++++++++--------------
 mctx/ctx_test.go | 39 +++++++++++++++++---------
 2 files changed, 76 insertions(+), 34 deletions(-)

diff --git a/mctx/ctx.go b/mctx/ctx.go
index 886b0e6..098954b 100644
--- a/mctx/ctx.go
+++ b/mctx/ctx.go
@@ -131,14 +131,6 @@ func ChildOf(ctx Context, name string) Context {
 	//	return mlog.DefaultWriteFn(w, msg)
 	//})
 
-	// copy mutable values
-	if len(s.mutVals) > 0 {
-		childS.mutVals = make(map[interface{}]interface{}, len(s.mutVals))
-		for key, val := range s.mutVals {
-			childS.mutVals[key] = val
-		}
-	}
-
 	// create child's ctx and store it in parent
 	childCtx := withCtxState(ctx, childS)
 	if s.children == nil {
@@ -151,19 +143,6 @@ func ChildOf(ctx Context, name string) Context {
 // TODO these might not be worth the effort, they're very subject to
 // race-conditions
 
-// SetMutableValue is like WithMutable, except rather than leaving the original
-// ctx unaffected it modifies the value in that context. Children of this
-// context will inherit an independent copy of its immutable values.
-func SetMutableValue(ctx Context, key, value interface{}) {
-	s := getCtxState(ctx)
-	s.l.Lock()
-	defer s.l.Unlock()
-	if s.mutVals == nil {
-		s.mutVals = map[interface{}]interface{}{}
-	}
-	s.mutVals[key] = value
-}
-
 // MutableValue acts like the Value method, except that it only deals with
 // keys/values set by SetMutableValue.
 func MutableValue(ctx Context, key interface{}) interface{} {
@@ -175,3 +154,53 @@ func MutableValue(ctx Context, key interface{}) interface{} {
 	}
 	return s.mutVals[key]
 }
+
+// GetSetMutableValue is used to interact with a mutable value on the context in
+// a thread-safe way. The key's value is retrieved and passed to the callback.
+// The value returned from the callback is stored to back into the context as
+// well as being returned from this function.
+//
+// If noCallbackIfSet is set to true, then this function will only return the
+// value for the key if it's already set and not use the callback at all.
+//
+// The callback returning nil is equivalent to unsetting the value.
+//
+// Children of this context will _not_ inherit any of its mutable values.
+func GetSetMutableValue(
+	ctx Context, noCallbackIfSet bool,
+	key interface{}, fn func(interface{}) interface{},
+) interface{} {
+	s := getCtxState(ctx)
+
+	if noCallbackIfSet {
+		s.l.RLock()
+		if s.mutVals != nil && s.mutVals[key] != nil {
+			defer s.l.RUnlock()
+			return s.mutVals[key]
+		}
+		s.l.RUnlock()
+	}
+
+	s.l.Lock()
+	defer s.l.Unlock()
+
+	if s.mutVals == nil {
+		s.mutVals = map[interface{}]interface{}{}
+	}
+	val := s.mutVals[key]
+
+	// It's possible something happened between the first check inside the
+	// read-lock and now, so double check this case. It's still good to have the
+	// read-lock check there, it'll handle 99% of the cases.
+	if noCallbackIfSet && val != nil {
+		return val
+	}
+
+	val = fn(val)
+	if val == nil {
+		delete(s.mutVals, key)
+	} else {
+		s.mutVals[key] = val
+	}
+	return val
+}
diff --git a/mctx/ctx_test.go b/mctx/ctx_test.go
index d15f7de..6cf6761 100644
--- a/mctx/ctx_test.go
+++ b/mctx/ctx_test.go
@@ -6,15 +6,11 @@ import (
 	"github.com/mediocregopher/mediocre-go-lib/mtest/massert"
 )
 
-func TestContext(t *T) {
+func TestInheritance(t *T) {
 	ctx := New()
-	SetMutableValue(ctx, "one", 1)
-
 	ctx1 := ChildOf(ctx, "1")
 	ctx1a := ChildOf(ctx1, "a")
-	SetMutableValue(ctx1, "one", 2)
 	ctx1b := ChildOf(ctx1, "b")
-	SetMutableValue(ctx1b, "one", 3)
 	ctx2 := ChildOf(ctx, "2")
 
 	massert.Fatal(t, massert.All(
@@ -47,12 +43,29 @@ func TestContext(t *T) {
 		massert.Equal(Parent(ctx1b), ctx1),
 		massert.Equal(Parent(ctx2), ctx),
 	))
-
-	massert.Fatal(t, massert.All(
-		massert.Equal(MutableValue(ctx, "one"), 1),
-		massert.Equal(MutableValue(ctx1, "one"), 2),
-		massert.Equal(MutableValue(ctx1a, "one"), 1),
-		massert.Equal(MutableValue(ctx1b, "one"), 3),
-		massert.Equal(MutableValue(ctx2, "one"), 1),
-	))
+}
+
+func TestMutableValues(t *T) {
+	fn := func(v interface{}) interface{} {
+		if v == nil {
+			return 0
+		}
+		return v.(int) + 1
+	}
+
+	var aa []massert.Assertion
+
+	ctx := New()
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx, false, 0, fn), 0))
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx, false, 0, fn), 1))
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx, true, 0, fn), 1))
+
+	aa = append(aa, massert.Equal(MutableValue(ctx, 0), 1))
+
+	ctx1 := ChildOf(ctx, "one")
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx1, true, 0, fn), 0))
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx1, false, 0, fn), 1))
+	aa = append(aa, massert.Equal(GetSetMutableValue(ctx1, true, 0, fn), 1))
+
+	massert.Fatal(t, massert.All(aa...))
 }