Provide GetAndSet(k string, setFn func) to allow synchronized atomic upsert behavior when key doesn't exist in cache

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

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