From 731dd72b84618db9d6dd300e43e5b096e837c37a Mon Sep 17 00:00:00 2001
From: Gabriel Aszalos <gabriel.aszalos@datadoghq.com>
Date: Wed, 7 Mar 2018 14:39:56 +0100
Subject: [PATCH] {cluster,ring}: add support for context to ClusterClient and
 Ring

---
 cluster.go | 19 +++++++++++++++++++
 ring.go    | 20 ++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/cluster.go b/cluster.go
index 767df5f1..1e787add 100644
--- 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
index c1c16d77..6d287741 100644
--- 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