From 464ed4abfa8e7496c2cd6cfb6dc27b804f1f5f5e Mon Sep 17 00:00:00 2001 From: JoeGruff Date: Wed, 6 May 2020 17:57:03 +0900 Subject: [PATCH] rpcclient: Stop client on ctx done. The passed context was ignored until a reconnect. This kills the client when the passed context is done. This is much more useful for consumers who need the client to shut down when their context if finished. --- rpcclient/infrastructure.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/rpcclient/infrastructure.go b/rpcclient/infrastructure.go index 86a18149..02502d15 100644 --- a/rpcclient/infrastructure.go +++ b/rpcclient/infrastructure.go @@ -1355,7 +1355,8 @@ func New(config *ConnConfig, ntfnHandlers *NotificationHandlers) (*Client, error // // This method will error if the client is not configured for websockets, if the // connection has already been established, or if none of the connection -// attempts were successful. +// attempts were successful. The client will be shut down when the passed +// context is terminated. func (c *Client) Connect(ctx context.Context, retry bool) error { c.mtx.Lock() defer c.mtx.Unlock() @@ -1367,6 +1368,14 @@ func (c *Client) Connect(ctx context.Context, retry bool) error { return ErrClientAlreadyConnected } + // Shutdown the client on context cancellation. + c.wg.Add(1) + go func() { + <-ctx.Done() + c.Shutdown() + c.wg.Done() + }() + // Begin connection attempts. Increase the backoff after each failed // attempt, up to a maximum of one minute. var backoff time.Duration