Merge pull request redis#737 from gbbr/context

cluster: add support for context.Context to ClusterClient and Ring
monjovi · Mar 8, 2018 · 9133634 · 9133634
2 parents b533525 + 731dd72
commit 9133634
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/cluster.go b/cluster.go
@@ -1,6 +1,7 @@
 package redis
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math"
@@ -557,6 +558,8 @@ func (c *clusterStateHolder) Get() (*clusterState, error) {
 type ClusterClient struct {
 	cmdable
 
+	ctx context.Context
+
 	opt           *ClusterOptions
 	nodes         *clusterNodes
 	state         *clusterStateHolder
@@ -593,6 +596,22 @@ func NewClusterClient(opt *ClusterOptions) *ClusterClient {
 	return c
 }
 
+func (c *ClusterClient) Context() context.Context {
+	if c.ctx != nil {
+		return c.ctx
+	}
+	return context.Background()
+}
+
+func (c *ClusterClient) WithContext(ctx context.Context) *ClusterClient {
+	if ctx == nil {
+		panic("nil context")
+	}
+	c2 := c.copy()
+	c2.ctx = ctx
+	return c2
+}
+
 func (c *ClusterClient) copy() *ClusterClient {
 	cp := *c
 	return &cp

diff --git a/ring.go b/ring.go
@@ -1,6 +1,7 @@
 package redis
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -289,6 +290,9 @@ func (c *ringShards) Close() error {
 // Otherwise you should use Redis Cluster.
 type Ring struct {
 	cmdable
+
+	ctx context.Context
+
 	opt           *RingOptions
 	shards        *ringShards
 	cmdsInfoCache *cmdsInfoCache
@@ -318,6 +322,22 @@ func NewRing(opt *RingOptions) *Ring {
 	return ring
 }
 
+func (c *Ring) Context() context.Context {
+	if c.ctx != nil {
+		return c.ctx
+	}
+	return context.Background()
+}
+
+func (c *Ring) WithContext(ctx context.Context) *Ring {
+	if ctx == nil {
+		panic("nil context")
+	}
+	c2 := c.copy()
+	c2.ctx = ctx
+	return c2
+}
+
 func (c *Ring) copy() *Ring {
 	cp := *c
 	return &cp