Refactor Set, SetNX and SetXX with expiration.

cluster_test.go

@@ -167,7 +167,7 @@ var _ = Describe("Cluster", func() {
 			Expect(err).To(Equal(redis.Nil))
 			Expect(val).To(Equal(""))
 
-			val, err = client.Set("A", "VALUE").Result()
+			val, err = client.Set("A", "VALUE", 0).Result()
 			Expect(err).NotTo(HaveOccurred())
 			Expect(val).To(Equal("OK"))
 
@@ -181,7 +181,7 @@ var _ = Describe("Cluster", func() {
 		})
 
 		It("should follow redirects", func() {
-			Expect(client.Set("A", "VALUE").Err()).NotTo(HaveOccurred())
+			Expect(client.Set("A", "VALUE", 0).Err()).NotTo(HaveOccurred())
 			Expect(redis.HashSlot("A")).To(Equal(6373))
 
 			// Slot 6373 is stored on the second node

command.go

@@ -338,12 +338,25 @@ func (cmd *BoolCmd) String() string {
 
 func (cmd *BoolCmd) parseReply(rd *bufio.Reader) error {
 	v, err := parseReply(rd, nil)
+	// `SET key value NX` returns nil when key already exists.
+	if err == Nil {
+		cmd.val = false
+		return nil
+	}
 	if err != nil {
 		cmd.err = err
 		return err
 	}
-	cmd.val = v.(int64) == 1
-	return nil
+	switch vv := v.(type) {
+	case int64:
+		cmd.val = vv == 1
+		return nil
+	case string:
+		cmd.val = vv == "OK"
+		return nil
+	default:
+		return fmt.Errorf("got %T, wanted int64 or string")
+	}
 }
 
 //------------------------------------------------------------------------------

command_test.go

@@ -26,15 +26,15 @@ var _ = Describe("Command", func() {
 	})
 
 	It("should have a plain string result", func() {
-		set := client.Set("foo", "bar")
+		set := client.Set("foo", "bar", 0)
 		Expect(set.String()).To(Equal("SET foo bar: OK"))
 
 		get := client.Get("foo")
 		Expect(get.String()).To(Equal("GET foo: bar"))
 	})
 
 	It("should have correct val/err states", func() {
-		set := client.Set("key", "hello")
+		set := client.Set("key", "hello", 0)
 		Expect(set.Err()).NotTo(HaveOccurred())
 		Expect(set.Val()).To(Equal("OK"))
 
@@ -47,7 +47,7 @@ var _ = Describe("Command", func() {
 	})
 
 	It("should escape special chars", func() {
-		set := client.Set("key", "hello1\r\nhello2\r\n")
+		set := client.Set("key", "hello1\r\nhello2\r\n", 0)
 		Expect(set.Err()).NotTo(HaveOccurred())
 		Expect(set.Val()).To(Equal("OK"))
 
@@ -58,7 +58,7 @@ var _ = Describe("Command", func() {
 
 	It("should handle big vals", func() {
 		val := string(bytes.Repeat([]byte{'*'}, 1<<16))
-		set := client.Set("key", val)
+		set := client.Set("key", val, 0)
 		Expect(set.Err()).NotTo(HaveOccurred())
 		Expect(set.Val()).To(Equal("OK"))
 
@@ -70,7 +70,7 @@ var _ = Describe("Command", func() {
 	It("should handle many keys #1", func() {
 		const n = 100000
 		for i := 0; i < n; i++ {
-			client.Set("keys.key"+strconv.Itoa(i), "hello"+strconv.Itoa(i))
+			client.Set("keys.key"+strconv.Itoa(i), "hello"+strconv.Itoa(i), 0)
 		}
 		keys := client.Keys("keys.*")
 		Expect(keys.Err()).NotTo(HaveOccurred())
@@ -83,7 +83,7 @@ var _ = Describe("Command", func() {
 		keys := []string{"non-existent-key"}
 		for i := 0; i < n; i++ {
 			key := "keys.key" + strconv.Itoa(i)
-			client.Set(key, "hello"+strconv.Itoa(i))
+			client.Set(key, "hello"+strconv.Itoa(i), 0)
 			keys = append(keys, key)
 		}
 		keys = append(keys, "non-existent-key")
@@ -100,7 +100,7 @@ var _ = Describe("Command", func() {
 	})
 
 	It("should convert strings via helpers", func() {
-		set := client.Set("key", "10")
+		set := client.Set("key", "10", 0)
 		Expect(set.Err()).NotTo(HaveOccurred())
 
 		n, err := client.Get("key").Int64()

commands.go

@@ -25,6 +25,18 @@ func (c *commandable) Process(cmd Cmder) {
 	c.process(cmd)
 }
 
+func usePrecise(dur time.Duration) bool {
+	return dur < time.Second || dur%time.Second != 0
+}
+
+func formatMs(dur time.Duration) string {
+	return strconv.FormatInt(int64(dur/time.Millisecond), 10)
+}
+
+func formatSec(dur time.Duration) string {
+	return strconv.FormatInt(int64(dur/time.Second), 10)
+}
+
 //------------------------------------------------------------------------------
 
 func (c *commandable) Auth(password string) *StatusCmd {
@@ -77,8 +89,8 @@ func (c *commandable) Exists(key string) *BoolCmd {
 	return cmd
 }
 
-func (c *commandable) Expire(key string, dur time.Duration) *BoolCmd {
-	cmd := NewBoolCmd("EXPIRE", key, strconv.FormatInt(int64(dur/time.Second), 10))
+func (c *commandable) Expire(key string, expiration time.Duration) *BoolCmd {
+	cmd := NewBoolCmd("EXPIRE", key, formatSec(expiration))
 	c.Process(cmd)
 	return cmd
 }
@@ -146,8 +158,8 @@ func (c *commandable) Persist(key string) *BoolCmd {
 	return cmd
 }
 
-func (c *commandable) PExpire(key string, dur time.Duration) *BoolCmd {
-	cmd := NewBoolCmd("PEXPIRE", key, strconv.FormatInt(int64(dur/time.Millisecond), 10))
+func (c *commandable) PExpire(key string, expiration time.Duration) *BoolCmd {
+	cmd := NewBoolCmd("PEXPIRE", key, formatMs(expiration))
 	c.Process(cmd)
 	return cmd
 }
@@ -425,19 +437,22 @@ func (c *commandable) MSetNX(pairs ...string) *BoolCmd {
 	return cmd
 }
 
-func (c *commandable) PSetEx(key string, dur time.Duration, value string) *StatusCmd {
-	cmd := NewStatusCmd(
-		"PSETEX",
-		key,
-		strconv.FormatInt(int64(dur/time.Millisecond), 10),
-		value,
-	)
+func (c *commandable) PSetEx(key string, expiration time.Duration, value string) *StatusCmd {
+	cmd := NewStatusCmd("PSETEX", key, formatMs(expiration), value)
 	c.Process(cmd)
 	return cmd
 }
 
-func (c *commandable) Set(key, value string) *StatusCmd {
-	cmd := NewStatusCmd("SET", key, value)
+func (c *commandable) Set(key, value string, expiration time.Duration) *StatusCmd {
+	args := []string{"SET", key, value}
+	if expiration > 0 {
+		if usePrecise(expiration) {
+			args = append(args, "PX", formatMs(expiration))
+		} else {
+			args = append(args, "EX", formatSec(expiration))
+		}
+	}
+	cmd := NewStatusCmd(args...)
 	c.Process(cmd)
 	return cmd
 }
@@ -453,14 +468,35 @@ func (c *commandable) SetBit(key string, offset int64, value int) *IntCmd {
 	return cmd
 }
 
-func (c *commandable) SetEx(key string, dur time.Duration, value string) *StatusCmd {
-	cmd := NewStatusCmd("SETEX", key, strconv.FormatInt(int64(dur/time.Second), 10), value)
+func (c *commandable) SetEx(key string, expiration time.Duration, value string) *StatusCmd {
+	cmd := NewStatusCmd("SETEX", key, formatSec(expiration), value)
 	c.Process(cmd)
 	return cmd
 }
 
-func (c *commandable) SetNX(key, value string) *BoolCmd {
-	cmd := NewBoolCmd("SETNX", key, value)
+func (c *commandable) SetNX(key, value string, expiration time.Duration) *BoolCmd {
+	var cmd *BoolCmd
+	if expiration == 0 {
+		// Use old `SETNX` to support old Redis versions.
+		cmd = NewBoolCmd("SETNX", key, value)
+	} else {
+		if usePrecise(expiration) {
+			cmd = NewBoolCmd("SET", key, value, "PX", formatMs(expiration), "NX")
+		} else {
+			cmd = NewBoolCmd("SET", key, value, "EX", formatSec(expiration), "NX")
+		}
+	}
+	c.Process(cmd)
+	return cmd
+}
+
+func (c *Client) SetXX(key, value string, expiration time.Duration) *BoolCmd {
+	var cmd *BoolCmd
+	if usePrecise(expiration) {
+		cmd = NewBoolCmd("SET", key, value, "PX", formatMs(expiration), "XX")
+	} else {
+		cmd = NewBoolCmd("SET", key, value, "EX", formatSec(expiration), "XX")
+	}
 	c.Process(cmd)
 	return cmd
 }