Avoid panic when using dedicated client after being recycled, return an error instead (#593)

* avoid panic on use after close, return error instead

* introduce ErrDedicatedClientRecycled

* better comment for ErrDedicatedClientRecycled

Author: FZambia

client.go

@@ -177,7 +177,9 @@ func (c *dedicatedSingleClient) B() Builder {
 
 func (c *dedicatedSingleClient) Do(ctx context.Context, cmd Completed) (resp RedisResult) {
 retry:
-	c.check()
+	if err := c.check(); err != nil {
+		return newErrResult(err)
+	}
 	resp = c.wire.Do(ctx, cmd)
 	if c.retry && cmd.IsReadOnly() && isRetryable(resp.NonRedisError(), c.wire, ctx) {
 		goto retry
@@ -197,7 +199,9 @@ func (c *dedicatedSingleClient) DoMulti(ctx context.Context, multi ...Completed)
 		retryable = allReadOnly(multi)
 	}
 retry:
-	c.check()
+	if err := c.check(); err != nil {
+		return fillErrs(len(multi), err)
+	}
 	resp = c.wire.DoMulti(ctx, multi...).s
 	if retryable && anyRetryable(resp, c.wire, ctx) {
 		goto retry
@@ -212,7 +216,9 @@ retry:
 
 func (c *dedicatedSingleClient) Receive(ctx context.Context, subscribe Completed, fn func(msg PubSubMessage)) (err error) {
 retry:
-	c.check()
+	if err := c.check(); err != nil {
+		return err
+	}
 	err = c.wire.Receive(ctx, subscribe, fn)
 	if c.retry {
 		if _, ok := err.(*RedisError); !ok && isRetryable(err, c.wire, ctx) {
@@ -226,7 +232,11 @@ retry:
 }
 
 func (c *dedicatedSingleClient) SetPubSubHooks(hooks PubSubHooks) <-chan error {
-	c.check()
+	if err := c.check(); err != nil {
+		ch := make(chan error, 1)
+		ch <- err
+		return ch
+	}
 	return c.wire.SetPubSubHooks(hooks)
 }
 
@@ -235,10 +245,11 @@ func (c *dedicatedSingleClient) Close() {
 	c.release()
 }
 
-func (c *dedicatedSingleClient) check() {
+func (c *dedicatedSingleClient) check() error {
 	if atomic.LoadUint32(&c.mark) != 0 {
-		panic(dedicatedClientUsedAfterReleased)
+		return ErrDedicatedClientRecycled
 	}
+	return nil
 }
 
 func (c *dedicatedSingleClient) release() {

client_test.go

@@ -548,10 +548,10 @@ func TestSingleClient(t *testing.T) {
 		}
 	})
 
-	t.Run("Dedicate panic after released", func(t *testing.T) {
+	t.Run("Dedicate ErrDedicatedClientRecycled after released", func(t *testing.T) {
 		m.AcquireFn = func() wire { return &mockWire{} }
-		check := func() {
-			if err := recover(); err != dedicatedClientUsedAfterReleased {
+		check := func(err error) {
+			if !errors.Is(err, ErrDedicatedClientRecycled) {
 				t.Fatalf("unexpected err %v", err)
 			}
 		}
@@ -567,20 +567,22 @@ func TestSingleClient(t *testing.T) {
 			closeFn(c, cancel)
 			for _, fn := range []func(){
 				func() {
-					defer check()
-					c.Do(context.Background(), c.B().Get().Key("k").Build())
+					resp := c.Do(context.Background(), c.B().Get().Key("k").Build())
+					check(resp.Error())
 				},
 				func() {
-					defer check()
-					c.DoMulti(context.Background(), c.B().Get().Key("k").Build())
+					resp := c.DoMulti(context.Background(), c.B().Get().Key("k").Build())
+					for _, r := range resp {
+						check(r.Error())
+					}
 				},
 				func() {
-					defer check()
-					c.Receive(context.Background(), c.B().Subscribe().Channel("k").Build(), func(msg PubSubMessage) {})
+					err := c.Receive(context.Background(), c.B().Subscribe().Channel("k").Build(), func(msg PubSubMessage) {})
+					check(err)
 				},
 				func() {
-					defer check()
-					c.SetPubSubHooks(PubSubHooks{})
+					ch := c.SetPubSubHooks(PubSubHooks{})
+					check(<-ch)
 				},
 			} {
 				fn()

cluster.go

@@ -1111,7 +1111,7 @@ func (c *dedicatedClusterClient) acquire(ctx context.Context, slot uint16) (wire
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.mark {
-		panic(dedicatedClientUsedAfterReleased)
+		return nil, ErrDedicatedClientRecycled
 	}
 	if c.slot == cmds.NoSlot {
 		c.slot = slot
@@ -1241,7 +1241,9 @@ func (c *dedicatedClusterClient) SetPubSubHooks(hooks PubSubHooks) <-chan error
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.mark {
-		panic(dedicatedClientUsedAfterReleased)
+		ch := make(chan error, 1)
+		ch <- ErrDedicatedClientRecycled
+		return ch
 	}
 	if p := c.pshks; p != nil {
 		c.pshks = nil

cluster_test.go

@@ -1520,9 +1520,9 @@ func TestClusterClient(t *testing.T) {
 		}
 	})
 
-	t.Run("Dedicate panic after released", func(t *testing.T) {
-		check := func() {
-			if err := recover(); err != dedicatedClientUsedAfterReleased {
+	t.Run("Dedicate ErrDedicatedClientRecycled after released", func(t *testing.T) {
+		check := func(err error) {
+			if !errors.Is(err, ErrDedicatedClientRecycled) {
 				t.Fatalf("unexpected err %v", err)
 			}
 		}
@@ -1538,20 +1538,22 @@ func TestClusterClient(t *testing.T) {
 			closeFn(c, cancel)
 			for _, fn := range []func(){
 				func() {
-					defer check()
-					c.Do(context.Background(), c.B().Get().Key("k").Build())
+					resp := c.Do(context.Background(), c.B().Get().Key("k").Build())
+					check(resp.Error())
 				},
 				func() {
-					defer check()
-					c.DoMulti(context.Background(), c.B().Get().Key("k").Build())
+					resp := c.DoMulti(context.Background(), c.B().Get().Key("k").Build())
+					for _, r := range resp {
+						check(r.Error())
+					}
 				},
 				func() {
-					defer check()
-					c.Receive(context.Background(), c.B().Subscribe().Channel("k").Build(), func(msg PubSubMessage) {})
+					err := c.Receive(context.Background(), c.B().Subscribe().Channel("k").Build(), func(msg PubSubMessage) {})
+					check(err)
 				},
 				func() {
-					defer check()
-					c.SetPubSubHooks(PubSubHooks{})
+					ch := c.SetPubSubHooks(PubSubHooks{})
+					check(<-ch)
 				},
 			} {
 				fn()

rueidis.go

@@ -51,6 +51,8 @@ var (
 	ErrReplicaOnlyNotSupported = errors.New("ReplicaOnly is not supported for single client")
 	// ErrWrongPipelineMultiplex means wrong value for ClientOption.PipelineMultiplex
 	ErrWrongPipelineMultiplex = errors.New("ClientOption.PipelineMultiplex must not be bigger than MaxPipelineMultiplex")
+	// ErrDedicatedClientRecycled means the caller attempted to use the dedicated client which has been already recycled (after canceled/closed).
+	ErrDedicatedClientRecycled = errors.New("dedicated client should not be used after recycled")
 )
 
 // ClientOption should be passed to NewClient to construct a Client
@@ -392,4 +394,3 @@ func dial(dst string, opt *ClientOption) (conn net.Conn, err error) {
 }
 
 const redisErrMsgCommandNotAllow = "command is not allowed"
-const dedicatedClientUsedAfterReleased = "DedicatedClient should not be used after recycled"