From c9521b468f957479cd0802b1f5c250e9ecc1b947 Mon Sep 17 00:00:00 2001 From: Dave Collins Date: Tue, 13 Jul 2021 19:17:50 -0500 Subject: [PATCH] rpcclient: Cancel client context on shutdown. This introduces a child context with cancellation to the manual Connect method and cancels it when the client is manually stopped via Shutdown to ensure callers are still able to shutdown the client via that method when using manual connections in addition to when using automatic ones. --- rpcclient/infrastructure.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/rpcclient/infrastructure.go b/rpcclient/infrastructure.go index e7e8e080a6..593177cb68 100644 --- a/rpcclient/infrastructure.go +++ b/rpcclient/infrastructure.go @@ -152,6 +152,10 @@ type Client struct { // It is protected by mtx. wsConn *websocket.Conn + // cancel is a function used to cancel the client context which forces a + // shutdown. + cancel func() + // disconnected indicated whether or not the server is disconnected. disconnected bool @@ -1014,6 +1018,11 @@ func (c *Client) doShutdown() bool { log.Tracef("Shutting down RPC client %s", c.config.Host) close(c.shutdown) + c.mtx.Lock() + if c.cancel != nil { + c.cancel() + } + c.mtx.Unlock() return true } @@ -1380,6 +1389,7 @@ func (c *Client) Connect(ctx context.Context, retry bool) error { } // Shutdown the client on context cancellation. + ctx, c.cancel = context.WithCancel(ctx) c.wg.Add(1) go func() { <-ctx.Done()