Browse Source

Add Increment & Decrement

Brad Fitzpatrick 14 years ago
parent
commit
8e208d84a7
2 changed files with 100 additions and 21 deletions
  1. 53 3
      memcache.go
  2. 47 18
      memcache_test.go

+ 53 - 3
memcache.go

@@ -25,6 +25,7 @@ import (
 	"io/ioutil"
 	"net"
 	"os"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -102,6 +103,8 @@ var (
 	resultNotFound  = []byte("NOT_FOUND\r\n")
 	resultDeleted   = []byte("DELETED\r\n")
 	resultEnd       = []byte("END\r\n")
+
+	resultClientErrorPrefix = []byte("CLIENT_ERROR ")
 )
 
 // New returns a memcache client using the provided server(s)
@@ -488,15 +491,20 @@ func (c *Client) populateOne(rw *bufio.ReadWriter, verb string, item *Item) os.E
 	return fmt.Errorf("memcache: unexpected response line from %q: %q", verb, string(line))
 }
 
-func writeExpectf(rw *bufio.ReadWriter, expect []byte, format string, args ...interface{}) os.Error {
+func writeReadLine(rw *bufio.ReadWriter, format string, args ...interface{}) ([]byte, os.Error) {
 	_, err := fmt.Fprintf(rw, format, args...)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	if err := rw.Flush(); err != nil {
-		return err
+		return nil, err
 	}
 	line, err := rw.ReadSlice('\n')
+	return line, err
+}
+
+func writeExpectf(rw *bufio.ReadWriter, expect []byte, format string, args ...interface{}) os.Error {
+	line, err := writeReadLine(rw, format, args...)
 	if err != nil {
 		return err
 	}
@@ -520,3 +528,45 @@ func (c *Client) Delete(key string) os.Error {
 		return writeExpectf(rw, resultDeleted, "delete %s\r\n", key)
 	})
 }
+
+// Increment atomically increments key by delta. The return value is
+// the new value after being incremented or an error. If the value
+// didn't exist in memcached the error is ErrCacheMiss. The value in
+// memcached must be an decimal number, or an error will be returned.
+// On 64-bit overflow, the new value wraps around.
+func (c *Client) Increment(key string, delta uint64) (newValue uint64, err os.Error) {
+	return c.incrDecr("incr", key, delta)
+}
+
+// Decrement atomically decrements key by delta. The return value is
+// the new value after being decremented or an error. If the value
+// didn't exist in memcached the error is ErrCacheMiss. The value in
+// memcached must be an decimal number, or an error will be returned.
+// On underflow, the new value is capped at zero and does not wrap
+// around.
+func (c *Client) Decrement(key string, delta uint64) (newValue uint64, err os.Error) {
+	return c.incrDecr("decr", key, delta)
+}
+
+func (c *Client) incrDecr(verb, key string, delta uint64) (uint64, os.Error) {
+	var val uint64
+	err := c.withKeyRw(key, func(rw *bufio.ReadWriter) os.Error {
+		line, err := writeReadLine(rw, "%s %s %d\r\n", verb, key, delta)
+		if err != nil {
+			return err
+		}
+		switch {
+		case bytes.Equal(line, resultNotFound):
+			return ErrCacheMiss
+		case bytes.HasPrefix(line, resultClientErrorPrefix):
+			errMsg := line[len(resultClientErrorPrefix) : len(line)-2]
+			return os.NewError("memcache: client error: " + string(errMsg))
+		}
+		val, err = strconv.Atoui64(string(line[:len(line)-2]))
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	return val, err
+}

+ 47 - 18
memcache_test.go

@@ -19,6 +19,8 @@ package memcache
 
 import (
 	"net"
+	"os"
+	"strings"
 	"testing"
 )
 
@@ -39,22 +41,30 @@ func TestMemcache(t *testing.T) {
 	if !setup(t) {
 		return
 	}
+	checkErr := func(err os.Error, format string, args ...interface{}) {
+		if err != nil {
+			t.Fatalf(format, args...)
+		}
+	}
+
 	c := New(testServer)
 
+	mustSet := func(it *Item) {
+		if err := c.Set(it); err != nil {
+			t.Fatalf("failed to Set %#v: %v", *it, err)
+		}
+	}
+
 	// Set
 	foo := &Item{Key: "foo", Value: []byte("fooval"), Flags: 123}
-	if err := c.Set(foo); err != nil {
-		t.Fatalf("first set(foo): %v", err)
-	}
-	if err := c.Set(foo); err != nil {
-		t.Fatalf("second set(foo): %v", err)
-	}
+	err := c.Set(foo)
+	checkErr(err, "first set(foo): %v", err)
+	err = c.Set(foo)
+	checkErr(err, "second set(foo): %v", err)
 
 	// Get
 	it, err := c.Get("foo")
-	if err != nil {
-		t.Fatalf("get(foo): %v", err)
-	}
+	checkErr(err, "get(foo): %v", err)
 	if it.Key != "foo" {
 		t.Errorf("get(foo) Key = %q, want foo", it.Key)
 	}
@@ -67,18 +77,15 @@ func TestMemcache(t *testing.T) {
 
 	// Add
 	bar := &Item{Key: "bar", Value: []byte("barval")}
-	if err := c.Add(bar); err != nil {
-		t.Fatalf("first add(foo): %v", err)
-	}
+	err = c.Add(bar)
+	checkErr(err, "first add(foo): %v", err)
 	if err := c.Add(bar); err != ErrNotStored {
 		t.Fatalf("second add(foo) want ErrNotStored, got %v", err)
 	}
 
 	// GetMulti
 	m, err := c.GetMulti([]string{"foo", "bar"})
-	if err != nil {
-		t.Fatalf("GetMulti: %v", err)
-	}
+	checkErr(err, "GetMulti: %v", err)
 	if g, e := len(m), 2; g != e {
 		t.Errorf("GetMulti: got len(map) = %d, want = %d", g, e)
 	}
@@ -97,12 +104,34 @@ func TestMemcache(t *testing.T) {
 
 	// Delete
 	err = c.Delete("foo")
-	if err != nil {
-		t.Errorf("Delete: %v", err)
-	}
+	checkErr(err, "Delete: %v", err)
 	it, err = c.Get("foo")
 	if err != ErrCacheMiss {
 		t.Errorf("post-Delete want ErrCacheMiss, got %v", err)
 	}
 
+	// Incr/Decr
+	mustSet(&Item{Key: "num", Value: []byte("42")})
+	n, err := c.Increment("num", 8)
+	checkErr(err, "Increment num + 8: %v", err)
+	if n != 50 {
+		t.Fatalf("Increment num + 8: want=50, got=%d", n)
+	}
+	n, err = c.Decrement("num", 49)
+	checkErr(err, "Decrement: %v", err)
+	if n != 1 {
+		t.Fatalf("Decrement 49: want=1, got=%d", n)
+	}
+	err = c.Delete("num")
+	checkErr(err, "delete num: %v", err)
+	n, err = c.Increment("num", 1)
+	if err != ErrCacheMiss {
+		t.Fatalf("increment post-delete: want ErrCacheMiss, got %v", err)
+	}
+	mustSet(&Item{Key: "num", Value: []byte("not-numeric")})
+	n, err = c.Increment("num", 1)
+	if err == nil || !strings.Contains(err.String(), "client error") {
+		t.Fatalf("increment non-number: want client error, got %v", err)
+	}
+
 }