From 03da66c18a66a00ff0f796370438a58f6d2b3b13 Mon Sep 17 00:00:00 2001
From: Artem Chernyshev <artem.chernyshev@emc.com>
Date: Mon, 3 Oct 2016 12:38:13 +0300
Subject: [PATCH] Add new parameter to the RingOptions to override CommandInfo
 first key

There is problem with `eval` and `evalsha` commands.
`COMMAND INFO eval` returns first key position equals `0`.
After that, redis ring chooses `eval` as a value for sharding.
They, if you try to delete created value, ring may choose another shard
and delete won't work.

Eval command should be parsed, to be sharded properly, according to
redis specs: http://redis.io/commands/command .

I've introduced a new flag in the `RingOptions`, which will enable new
behavior: `EnableKeyLocationParsing`.

If it is enabled, `cmdFirstKey` will try to get key position using
`cmd.getFirstKeyPos()`. This function is defined for `eval` and
`evalsha` commands.
If it has parameters, it will return `3`, otherwise it will return `0`.
---
 main_test.go |  1 +
 ring.go      | 21 +++++++++++++++++++++
 ring_test.go | 17 +++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/main_test.go b/main_test.go
index a9afa82..d79a143 100644
--- a/main_test.go
+++ b/main_test.go
@@ -138,6 +138,7 @@ func redisRingOptions() *redis.RingOptions {
 		PoolTimeout:        30 * time.Second,
 		IdleTimeout:        500 * time.Millisecond,
 		IdleCheckFrequency: 500 * time.Millisecond,
+		RouteByEvalKeys:    true,
 	}
 }
 
diff --git a/ring.go b/ring.go
index b0bba58..08d52ef 100644
--- a/ring.go
+++ b/ring.go
@@ -40,6 +40,9 @@ type RingOptions struct {
 	PoolTimeout        time.Duration
 	IdleTimeout        time.Duration
 	IdleCheckFrequency time.Duration
+
+	// RouteByEvalKeys flag to enable eval and evalsha key position parsing for sharding
+	RouteByEvalKeys bool
 }
 
 func (opt *RingOptions) init() {
@@ -132,6 +135,8 @@ type Ring struct {
 	cmdsInfoOnce *sync.Once
 
 	closed bool
+
+	routeByEvalKeys bool
 }
 
 var _ Cmdable = (*Ring)(nil)
@@ -154,6 +159,7 @@ func NewRing(opt *RingOptions) *Ring {
 		clopt.Addr = addr
 		ring.addClient(name, NewClient(clopt))
 	}
+	ring.routeByEvalKeys = opt.RouteByEvalKeys
 	go ring.heartbeat()
 	return ring
 }
@@ -221,7 +227,22 @@ func (c *Ring) cmdInfo(name string) *CommandInfo {
 	return c.cmdsInfo[name]
 }
 
+func (c *Ring) getEvalFirstKey(cmd Cmder) string {
+	if c.routeByEvalKeys && cmd.arg(2) != "0" {
+		return cmd.arg(3)
+	} else {
+		return cmd.arg(0)
+	}
+}
+
 func (c *Ring) cmdFirstKey(cmd Cmder) string {
+	switch cmd.arg(0) {
+	case "eval":
+		return c.getEvalFirstKey(cmd)
+	case "evalsha":
+		return c.getEvalFirstKey(cmd)
+	}
+
 	cmdInfo := c.cmdInfo(cmd.arg(0))
 	if cmdInfo == nil {
 		internal.Logf("info for cmd=%s not found", cmd.arg(0))
diff --git a/ring_test.go b/ring_test.go
index 3c06686..0ab258d 100644
--- a/ring_test.go
+++ b/ring_test.go
@@ -89,6 +89,23 @@ var _ = Describe("Redis Ring", func() {
 		Expect(ringShard2.Info().Val()).To(ContainSubstring("keys=100"))
 	})
 
+	It("supports eval key search", func() {
+		script := redis.NewScript(`
+			local r = redis.call('SET', KEYS[1], ARGV[1])
+			return r
+		`)
+
+		var key string
+		for i := 0; i < 100; i++ {
+			key = fmt.Sprintf("key{%d}", i)
+			err := script.Run(ring, []string{key}, "value").Err()
+			Expect(err).NotTo(HaveOccurred())
+		}
+
+		Expect(ringShard1.Info().Val()).To(ContainSubstring("keys=52"))
+		Expect(ringShard2.Info().Val()).To(ContainSubstring("keys=48"))
+	})
+
 	Describe("pipelining", func() {
 		It("returns an error when all shards are down", func() {
 			ring := redis.NewRing(&redis.RingOptions{})