From 97e6ed817821d91977b51fdec185ca2dbb917fa6 Mon Sep 17 00:00:00 2001
From: Andrea Spacca <andrea.spacca@gmail.com>
Date: Wed, 24 Apr 2019 08:33:36 +0200
Subject: [PATCH] Ring instrumentantions (#1017)

* Ring instrumentantions
---
 ring.go      |  9 ++++++++-
 ring_test.go | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/ring.go b/ring.go
index a582aefc..5956b71a 100644
--- a/ring.go
+++ b/ring.go
@@ -361,7 +361,8 @@ func NewRing(opt *RingOptions) *Ring {
 
 	ring.process = ring.defaultProcess
 	ring.processPipeline = ring.defaultProcessPipeline
-	ring.cmdable.setProcessor(ring.Process)
+
+	ring.init()
 
 	for name, addr := range opt.Addrs {
 		clopt := opt.clientOptions()
@@ -374,6 +375,10 @@ func NewRing(opt *RingOptions) *Ring {
 	return ring
 }
 
+func (c *Ring) init() {
+	c.cmdable.setProcessor(c.Process)
+}
+
 func (c *Ring) Context() context.Context {
 	if c.ctx != nil {
 		return c.ctx
@@ -392,6 +397,8 @@ func (c *Ring) WithContext(ctx context.Context) *Ring {
 
 func (c *Ring) clone() *Ring {
 	cp := *c
+	cp.init()
+
 	return &cp
 }
 
diff --git a/ring_test.go b/ring_test.go
index 4ff08986..d498e034 100644
--- a/ring_test.go
+++ b/ring_test.go
@@ -1,6 +1,7 @@
 package redis_test
 
 import (
+	"context"
 	"crypto/rand"
 	"fmt"
 	"net"
@@ -104,6 +105,27 @@ var _ = Describe("Redis Ring", func() {
 		Expect(ringShard2.Info("keyspace").Val()).To(ContainSubstring("keys=100"))
 	})
 
+	It("propagates process for WithContext", func() {
+		var fromWrap []string
+		wrapper := func(oldProcess func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+			return func(cmd redis.Cmder) error {
+				fromWrap = append(fromWrap, cmd.Name())
+
+				return oldProcess(cmd)
+			}
+		}
+
+		ctx := context.Background()
+		ring = ring.WithContext(ctx)
+		ring.WrapProcess(wrapper)
+
+		ring.Ping()
+		Expect(fromWrap).To(Equal([]string{"ping"}))
+
+		ring.Ping()
+		Expect(fromWrap).To(Equal([]string{"ping", "ping"}))
+	})
+
 	Describe("pipeline", func() {
 		It("distributes keys", func() {
 			pipe := ring.Pipeline()