Update to fix blocking pops.

Now correctly handles redis multi bulk reply.
qm1004 · Nov 20, 2010 · df9826b · df9826b
1 parent 3004586
commit df9826b
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 12 deletions.
diff --git a/redis.go b/redis.go
@@ -601,22 +601,27 @@ func (client *Client) Rpop(key string) ([]byte, os.Error) {
     return res.([]byte), nil
 }
 
-func (client *Client) Blpop(key string) ([]byte, os.Error) {
-    res, err := client.sendCommand("BLPOP", key)
-    if err != nil {
-        return nil, err
-    }
-
-    return res.([]byte), nil
+func (client *Client) Blpop(keys []string, timeoutSecs uint) (*string, []byte, os.Error) {
+    return client.bpop("BLPOP", keys, timeoutSecs)
+}
+func (client *Client) Brpop(keys []string, timeoutSecs uint) (*string, []byte, os.Error) {
+    return client.bpop("BRPOP", keys, timeoutSecs)
 }
 
-func (client *Client) Brpop(key string) ([]byte, os.Error) {
-    res, err := client.sendCommand("BRPOP", key)
+func (client *Client) bpop(cmd string, keys []string, timeoutSecs uint) (*string, []byte, os.Error) {
+    args := append(keys, strconv.Uitoa(timeoutSecs))
+    res, err := client.sendCommand(cmd, args...)
     if err != nil {
-        return nil, err
+        return nil, nil, err
     }
-
-    return res.([]byte), nil
+    kv := res.([][]byte)
+    // Check for timeout
+    if len(kv) != 2 {
+        return nil, nil, nil
+    }
+    k := string(kv[0])
+    v := kv[1]
+    return &k, v, nil
 }
 
 func (client *Client) Rpoplpush(src string, dst string) ([]byte, os.Error) {

diff --git a/redis_test.go b/redis_test.go
@@ -9,6 +9,7 @@ import (
     "runtime"
     "strconv"
     "strings"
+    "time"
     "testing"
 )
 
@@ -173,6 +174,70 @@ func TestList(t *testing.T) {
 
 }
 
+func TestBrpop(t *testing.T) {
+    go func() {
+        time.Sleep(100 * 1000)
+        if err := client.Lpush("l", []byte("a")); err != nil {
+            t.Fatal("Lpush failed", err.String())
+        }
+    }()
+    key, value, err := client.Brpop([]string{"l"}, 1)
+    if err != nil {
+        t.Fatal("Brpop failed", err.String())
+    }
+    if *key != "l" {
+        t.Fatalf("Expected %s but got %s", "l", *key)
+    }
+    if string(value) != "a" {
+        t.Fatalf("Expected %s but got %s", "a", string(value))
+    }
+}
+
+func TestBlpop(t *testing.T) {
+    go func() {
+        time.Sleep(100 * 1000)
+        if err := client.Lpush("l", []byte("a")); err != nil {
+            t.Fatal("Lpush failed", err.String())
+        }
+    }()
+    key, value, err := client.Blpop([]string{"l"}, 1)
+    if err != nil {
+        t.Fatal("Blpop failed", err.String())
+    }
+    if *key != "l" {
+        t.Fatalf("Expected %s but got %s", "l", *key)
+    }
+    if string(value) != "a" {
+        t.Fatalf("Expected %s but got %s", "a", string(value))
+    }
+}
+
+func TestBrpopTimeout(t *testing.T) {
+    key, value, err := client.Brpop([]string{"l"}, 1)
+    if err != nil {
+        t.Fatal("BrpopTimeout failed", err.String())
+    }
+    if key != nil {
+        t.Fatalf("Expected %s but got %s", "", key)
+    }
+    if value != nil {
+        t.Fatalf("Expected %s but got %s", nil, value)
+    }
+}
+
+func TestBlpopTimeout(t *testing.T) {
+    key, value, err := client.Blpop([]string{"l"}, 1)
+    if err != nil {
+        t.Fatal("BlpopTimeout failed", err.String())
+    }
+    if key != nil {
+        t.Fatalf("Expected %s but got %s", "", key)
+    }
+    if value != nil {
+        t.Fatalf("Expected %s but got %s", nil, value)
+    }
+}
+
 func verifyHash(t *testing.T, key string, expected map[string][]byte) {
     //test Hget
     m1 := make(map[string][]byte)