From 15c04393c8778fb24459f832eb256c876d016ab0 Mon Sep 17 00:00:00 2001
From: Douglas Daniels <daniels.douglas@gmail.com>
Date: Wed, 10 Feb 2016 14:56:17 -0600
Subject: [PATCH] Provide GetAndSet(k string, setFn func) to allow synchronized
 atomic upsert behavior when key doesn't exist in cache

---
 cache.go      | 32 ++++++++++++++++++++++++++++
 cache_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/cache.go b/cache.go
index 70e4dad..3f572bb 100644
--- a/cache.go
+++ b/cache.go
@@ -149,6 +149,38 @@ func (c *cache) get(k string) (interface{}, bool) {
 	return item.Object, true
 }
 
+func (c *cache) getWithExpiration(k string) (interface{}, time.Time, bool) {
+	item, found := c.items[k]
+	if !found {
+		return nil, time.Time{}, false
+	}
+	// "Inlining" of Expired
+	if item.Expiration > 0 {
+		if time.Now().UnixNano() > item.Expiration {
+			return nil, time.Time{}, false
+		}
+	}
+	return item.Object, time.Unix(0, item.Expiration), true
+}
+
+// SetTransaction given a values current state and expiration and whether it was found in the cache
+// atomically update and return the value and the new expiration time
+type SetTransaction func(interface{}, time.Time, bool) (interface{}, time.Duration)
+
+// GetAndSet allows retrieval of a cache key and setting the value based on a setFn in a synchronized atomic manner
+func (c *cache) GetAndSet(k string, setFn SetTransaction) (interface{}, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	v, expiration, found := c.getWithExpiration(k)
+
+	newV, duration := setFn(v, expiration, found)
+
+	c.set(k, newV, duration)
+
+	return newV, found
+}
+
 // Increment an item of type int, int8, int16, int32, int64, uintptr, uint,
 // uint8, uint32, or uint64, float32 or float64 by n. Returns an error if the
 // item's value is not an integer, if it was not found, or if it is not
diff --git a/cache_test.go b/cache_test.go
index 6e81693..e8e3f95 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -68,6 +68,63 @@ func TestCache(t *testing.T) {
 	}
 }
 
+func TestGetAndSet(t *testing.T) {
+	var found bool
+
+	tc := New(50*time.Millisecond, 1*time.Millisecond)
+
+	setHandler := func(v interface{}, expiration time.Time, found bool) (interface{}, time.Duration) {
+		// Cache miss so set to 1 with default expiration
+		if !found {
+			return int64(1), DefaultExpiration
+		}
+
+		return v.(int64) + 1, time.Now().Sub(expiration)
+	}
+
+	x, found := tc.GetAndSet("a", func(v interface{}, expiration time.Time, found bool) (interface{}, time.Duration) {
+		if found {
+			t.Error("a was found while GetAndSet")
+		}
+		if v != nil {
+			t.Error("expected key[a] to be nil")
+		}
+		if !expiration.IsZero() {
+			t.Error("expected expiration to not exist")
+		}
+		return setHandler(v, expiration, found)
+	})
+
+	if x != int64(1) {
+		t.Errorf("Expected x = 1 after GetAndSet but was %d", x)
+	}
+
+	if found {
+		t.Errorf("Expected found[false] on initial GetAndSet but was %t", found)
+	}
+
+	x, found = tc.GetAndSet("a", func(v interface{}, expiration time.Time, found bool) (interface{}, time.Duration) {
+		if !found {
+			t.Error("a was not found in second GetAndSet")
+		}
+		if v == nil {
+			t.Error("expected key[a] to not be nil")
+		}
+		if expiration.IsZero() {
+			t.Error("expected expiration to exist")
+		}
+		return setHandler(v, expiration, found)
+	})
+
+	if x != int64(2) {
+		t.Errorf("Expected x = 2 after GetAndSet but was %d", x)
+	}
+
+	if !found {
+		t.Errorf("Expected found[true] on second GetAndSet was %t", found)
+	}
+}
+
 func TestCacheTimes(t *testing.T) {
 	var found bool
 
@@ -1459,7 +1516,7 @@ func BenchmarkRWMutexMapGet(b *testing.B) {
 
 func BenchmarkRWMutexInterfaceMapGetStruct(b *testing.B) {
 	b.StopTimer()
-	s := struct{name string}{name: "foo"}
+	s := struct{ name string }{name: "foo"}
 	m := map[interface{}]string{
 		s: "bar",
 	}