From a78bca69e4e8b80515f04fcaf80f4ec04368702f Mon Sep 17 00:00:00 2001
From: Patrick Mylund Nielsen <patrick@patrickmylund.com>
Date: Mon, 2 Jan 2012 14:04:47 +0100
Subject: [PATCH] Added Add and Replace commands

---
 cache.go      | 28 ++++++++++++++++++++++++++--
 cache_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/cache.go b/cache.go
index df24930..5db4146 100644
--- a/cache.go
+++ b/cache.go
@@ -108,8 +108,8 @@ type janitor struct {
 	stop     chan bool
 }
 
-// Adds an item to the cache. If the duration is 0, the cache's default expiration time
-// is used. If it is -1, the item never expires.
+// Adds an item to the cache, replacing any existing item. If the duration is 0, the
+// cache's default expiration time is used. If it is -1, the item never expires.
 func (c *cache) Set(k string, x interface{}, d time.Duration) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -132,6 +132,30 @@ func (c *cache) Set(k string, x interface{}, d time.Duration) {
 	}
 }
 
+// TODO: Add and Replace aren't completely atomic
+
+// Adds an item to the cache only if an item doesn't already exist for the given key,
+// or if the existing item has expired. Returns an error if not.
+func (c *cache) Add(k string, x interface{}, d time.Duration) error {
+	_, found := c.Get(k)
+	if found {
+		return fmt.Errorf("Item %s already exists", k)
+	}
+	c.Set(k, x, d)
+	return nil
+}
+
+// Sets a new value for the cache item only if it already exists. Returns an error if
+// it does not.
+func (c *cache) Replace(k string, x interface{}, d time.Duration) error {
+	_, found := c.Get(k)
+	if !found {
+		return fmt.Errorf("Item %s doesn't exist", k)
+	}
+	c.Set(k, x, d)
+	return nil
+}
+
 // Gets an item from the cache.
 func (c *cache) Get(k string) (interface{}, bool) {
 	c.mu.Lock()
diff --git a/cache_test.go b/cache_test.go
index cb8f977..471da2f 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -349,3 +349,28 @@ func TestDecrementInt64(t *testing.T) {
 		t.Error("int64 is not 3:", x)
 	}
 }
+
+func TestAdd(t *testing.T) {
+	tc := New(0, 0)
+	err := tc.Add("foo", "bar", 0)
+	if err != nil {
+		t.Error("Couldn't add foo even though it shouldn't exist")
+	}
+	err = tc.Add("foo", "baz", 0)
+	if err == nil {
+		t.Error("Successfully added another foo when it should have returned an error")
+	}
+}
+
+func TestReplace(t *testing.T) {
+	tc := New(0, 0)
+	err := tc.Replace("foo", "bar", 0)
+	if err == nil {
+		t.Error("Replaced foo when it shouldn't exist")
+	}
+	tc.Set("foo", "bar", 0)
+	err = tc.Replace("foo", "bar", 0)
+	if err != nil {
+		t.Error("Couldn't replace existing key foo")
+	}
+}