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{})