From 3edcbbd04f8d38eb63c225c21514e7eb30118d78 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 | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/rpcclient/infrastructure.go b/rpcclient/infrastructure.go index 86a18149c0..d4c09a8051 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() @@ -1398,6 +1399,10 @@ func (c *Client) Connect(ctx context.Context, retry bool) error { c.wsConn = wsConn close(c.connEstablished) c.start() + go func() { + <-ctx.Done() + c.Shutdown() + }() if !c.config.DisableAutoReconnect { c.wg.Add(2) go c.wsReconnectHandler(ctx)