Add WrapProcessPipeline

Author: yifei.huang

cluster.go

@@ -445,6 +445,10 @@ type ClusterClient struct {
 	cmdsInfoOnce internal.Once
 	cmdsInfo     map[string]*CommandInfo
 
+	process           func(Cmder) error
+	processPipeline   func([]Cmder) error
+	processTxPipeline func([]Cmder) error
+
 	// Reports whether slots reloading is in progress.
 	reloading uint32
 }
@@ -458,7 +462,12 @@ func NewClusterClient(opt *ClusterOptions) *ClusterClient {
 		opt:   opt,
 		nodes: newClusterNodes(opt),
 	}
-	c.setProcessor(c.Process)
+
+	c.process = c.defaultProcess
+	c.processPipeline = c.defaultProcessPipeline
+	c.processTxPipeline = c.defaultProcessTxPipeline
+
+	c.cmdable.setProcessor(c.Process)
 
 	// Add initial nodes.
 	for _, addr := range opt.Addrs {
@@ -628,7 +637,20 @@ func (c *ClusterClient) Close() error {
 	return c.nodes.Close()
 }
 
+func (c *ClusterClient) WrapProcess(
+	fn func(oldProcess func(Cmder) error) func(Cmder) error,
+) {
+	c.process = fn(c.process)
+}
+
 func (c *ClusterClient) Process(cmd Cmder) error {
+	if c.process != nil {
+		return c.process(cmd)
+	}
+	return c.defaultProcess(cmd)
+}
+
+func (c *ClusterClient) defaultProcess(cmd Cmder) error {
 	state, err := c.state()
 	if err != nil {
 		cmd.setErr(err)
@@ -910,17 +932,23 @@ func (c *ClusterClient) reaper(idleCheckFrequency time.Duration) {
 
 func (c *ClusterClient) Pipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.pipelineExec,
+		exec: c.processPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }
 
 func (c *ClusterClient) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 	return c.Pipeline().Pipelined(fn)
 }
 
-func (c *ClusterClient) pipelineExec(cmds []Cmder) error {
+func (c *ClusterClient) WrapProcessPipeline(
+	fn func(oldProcess func([]Cmder) error) func([]Cmder) error,
+) {
+	c.processPipeline = fn(c.processPipeline)
+}
+
+func (c *ClusterClient) defaultProcessPipeline(cmds []Cmder) error {
 	cmdsMap, err := c.mapCmdsByNode(cmds)
 	if err != nil {
 		setCmdsErr(cmds, err)
@@ -1064,17 +1092,17 @@ func (c *ClusterClient) checkMovedErr(
 // TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
 func (c *ClusterClient) TxPipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.txPipelineExec,
+		exec: c.processTxPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }
 
 func (c *ClusterClient) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 	return c.TxPipeline().Pipelined(fn)
 }
 
-func (c *ClusterClient) txPipelineExec(cmds []Cmder) error {
+func (c *ClusterClient) defaultProcessTxPipeline(cmds []Cmder) error {
 	state, err := c.state()
 	if err != nil {
 		return err

example_instrumentation_test.go

@@ -2,58 +2,47 @@ package redis_test
 
 import (
 	"fmt"
-	"sync/atomic"
-	"time"
 
 	"github.com/go-redis/redis"
 )
 
 func Example_instrumentation() {
-	ring := redis.NewRing(&redis.RingOptions{
-		Addrs: map[string]string{
-			"shard1": ":6379",
-		},
+	cl := redis.NewClient(&redis.Options{
+		Addr: ":6379",
 	})
-	ring.ForEachShard(func(client *redis.Client) error {
-		wrapRedisProcess(client)
-		return nil
+	cl.WrapProcess(func(old func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+		return func(cmd redis.Cmder) error {
+			fmt.Printf("starting processing: <%s>\n", cmd)
+			err := old(cmd)
+			fmt.Printf("finished processing: <%s>\n", cmd)
+			return err
+		}
 	})
 
-	for {
-		ring.Ping()
-	}
+	cl.Ping()
+	// Output: starting processing: <ping: >
+	// finished processing: <ping: PONG>
 }
 
-func wrapRedisProcess(client *redis.Client) {
-	const precision = time.Microsecond
-	var count, avgDur uint32
-
-	go func() {
-		for range time.Tick(3 * time.Second) {
-			n := atomic.LoadUint32(&count)
-			dur := time.Duration(atomic.LoadUint32(&avgDur)) * precision
-			fmt.Printf("%s: processed=%d avg_dur=%s\n", client, n, dur)
-		}
-	}()
-
-	client.WrapProcess(func(oldProcess func(redis.Cmder) error) func(redis.Cmder) error {
-		return func(cmd redis.Cmder) error {
-			start := time.Now()
-			err := oldProcess(cmd)
-			dur := time.Since(start)
-
-			const decay = float64(1) / 100
-			ms := float64(dur / precision)
-			for {
-				avg := atomic.LoadUint32(&avgDur)
-				newAvg := uint32((1-decay)*float64(avg) + decay*ms)
-				if atomic.CompareAndSwapUint32(&avgDur, avg, newAvg) {
-					break
-				}
-			}
-			atomic.AddUint32(&count, 1)
+func Example_Pipeline_instrumentation() {
+	client := redis.NewClient(&redis.Options{
+		Addr: ":6379",
+	})
 
+	client.WrapProcessPipeline(func(old func([]redis.Cmder) error) func([]redis.Cmder) error {
+		return func(cmds []redis.Cmder) error {
+			fmt.Printf("pipeline starting processing: %v\n", cmds)
+			err := old(cmds)
+			fmt.Printf("pipeline finished processing: %v\n", cmds)
 			return err
 		}
 	})
+
+	client.Pipelined(func(pipe redis.Pipeliner) error {
+		pipe.Ping()
+		pipe.Ping()
+		return nil
+	})
+	// Output: pipeline starting processing: [ping:  ping: ]
+	// pipeline finished processing: [ping: PONG ping: PONG]
 }

redis.go

@@ -22,6 +22,12 @@ func SetLogger(logger *log.Logger) {
 	internal.Logger = logger
 }
 
+func (c *baseClient) init() {
+	c.process = c.defaultProcess
+	c.processPipeline = c.defaultProcessPipeline
+	c.processTxPipeline = c.defaultProcessTxPipeline
+}
+
 func (c *baseClient) String() string {
 	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
 }
@@ -85,7 +91,8 @@ func (c *baseClient) initConn(cn *pool.Conn) error {
 			connPool: pool.NewSingleConnPool(cn),
 		},
 	}
-	conn.setProcessor(conn.Process)
+	conn.baseClient.init()
+	conn.statefulCmdable.setProcessor(conn.Process)
 
 	_, err := conn.Pipelined(func(pipe Pipeliner) error {
 		if c.opt.Password != "" {
@@ -117,14 +124,11 @@ func (c *baseClient) initConn(cn *pool.Conn) error {
 // an input and returns the new wrapper process func. createWrapper should
 // use call the old process func within the new process func.
 func (c *baseClient) WrapProcess(fn func(oldProcess func(cmd Cmder) error) func(cmd Cmder) error) {
-	c.process = fn(c.defaultProcess)
+	c.process = fn(c.process)
 }
 
 func (c *baseClient) Process(cmd Cmder) error {
-	if c.process != nil {
-		return c.process(cmd)
-	}
-	return c.defaultProcess(cmd)
+	return c.process(cmd)
 }
 
 func (c *baseClient) defaultProcess(cmd Cmder) error {
@@ -198,35 +202,48 @@ func (c *baseClient) getAddr() string {
 	return c.opt.Addr
 }
 
+func (c *baseClient) WrapProcessPipeline(
+	fn func(oldProcess func([]Cmder) error) func([]Cmder) error,
+) {
+	c.processPipeline = fn(c.processPipeline)
+	c.processTxPipeline = fn(c.processTxPipeline)
+}
+
+func (c *baseClient) defaultProcessPipeline(cmds []Cmder) error {
+	return c.generalProcessPipeline(cmds, c.pipelineProcessCmds)
+}
+
+func (c *baseClient) defaultProcessTxPipeline(cmds []Cmder) error {
+	return c.generalProcessPipeline(cmds, c.txPipelineProcessCmds)
+}
+
 type pipelineProcessor func(*pool.Conn, []Cmder) (bool, error)
 
-func (c *baseClient) pipelineExecer(p pipelineProcessor) pipelineExecer {
-	return func(cmds []Cmder) error {
-		for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
-			if attempt > 0 {
-				time.Sleep(c.retryBackoff(attempt))
-			}
+func (c *baseClient) generalProcessPipeline(cmds []Cmder, p pipelineProcessor) error {
+	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
+		if attempt > 0 {
+			time.Sleep(c.retryBackoff(attempt))
+		}
 
-			cn, _, err := c.getConn()
-			if err != nil {
-				setCmdsErr(cmds, err)
-				return err
-			}
+		cn, _, err := c.getConn()
+		if err != nil {
+			setCmdsErr(cmds, err)
+			return err
+		}
 
-			canRetry, err := p(cn, cmds)
+		canRetry, err := p(cn, cmds)
 
-			if err == nil || internal.IsRedisError(err) {
-				_ = c.connPool.Put(cn)
-				break
-			}
-			_ = c.connPool.Remove(cn)
+		if err == nil || internal.IsRedisError(err) {
+			_ = c.connPool.Put(cn)
+			break
+		}
+		_ = c.connPool.Remove(cn)
 
-			if !canRetry || !internal.IsRetryableError(err, true) {
-				break
-			}
+		if !canRetry || !internal.IsRetryableError(err, true) {
+			break
 		}
-		return firstCmdsErr(cmds)
 	}
+	return firstCmdsErr(cmds)
 }
 
 func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
@@ -324,14 +341,15 @@ type Client struct {
 }
 
 func newClient(opt *Options, pool pool.Pooler) *Client {
-	client := Client{
+	c := Client{
 		baseClient: baseClient{
 			opt:      opt,
 			connPool: pool,
 		},
 	}
-	client.setProcessor(client.Process)
-	return &client
+	c.baseClient.init()
+	c.cmdable.setProcessor(c.Process)
+	return &c
 }
 
 // NewClient returns a client to the Redis Server specified by Options.
@@ -343,7 +361,7 @@ func NewClient(opt *Options) *Client {
 func (c *Client) copy() *Client {
 	c2 := new(Client)
 	*c2 = *c
-	c2.setProcessor(c2.Process)
+	c2.cmdable.setProcessor(c2.Process)
 	return c2
 }
 
@@ -366,9 +384,9 @@ func (c *Client) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 
 func (c *Client) Pipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.pipelineExecer(c.pipelineProcessCmds),
+		exec: c.processPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }
 
@@ -379,9 +397,9 @@ func (c *Client) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 // TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
 func (c *Client) TxPipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.pipelineExecer(c.txPipelineProcessCmds),
+		exec: c.processTxPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }
 
@@ -430,9 +448,9 @@ func (c *Conn) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 
 func (c *Conn) Pipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.pipelineExecer(c.pipelineProcessCmds),
+		exec: c.processPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }
 
@@ -443,8 +461,8 @@ func (c *Conn) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
 // TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
 func (c *Conn) TxPipeline() Pipeliner {
 	pipe := Pipeline{
-		exec: c.pipelineExecer(c.txPipelineProcessCmds),
+		exec: c.processTxPipeline,
 	}
-	pipe.setProcessor(pipe.Process)
+	pipe.statefulCmdable.setProcessor(pipe.Process)
 	return &pipe
 }

redis_context.go

@@ -12,7 +12,10 @@ type baseClient struct {
 	connPool pool.Pooler
 	opt      *Options
 
-	process func(Cmder) error
+	process           func(Cmder) error
+	processPipeline   func([]Cmder) error
+	processTxPipeline func([]Cmder) error
+
 	onClose func() error // hook called when client is closed
 
 	ctx context.Context

redis_no_context.go

@@ -10,6 +10,9 @@ type baseClient struct {
 	connPool pool.Pooler
 	opt      *Options
 
-	process func(Cmder) error
+	process           func(Cmder) error
+	processPipeline   func([]Cmder) error
+	processTxPipeline func([]Cmder) error
+
 	onClose func() error // hook called when client is closed
 }