From d02f2e6b259072386a967e02492e2587a0cd82eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20R=C3=BChl?= <sebastian@mapped.com>
Date: Fri, 5 Jan 2024 11:23:01 +0100
Subject: [PATCH] fix: add waitGroup to start go func

---
 client.go | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/client.go b/client.go
index dd46cb5..f7b49d1 100644
--- a/client.go
+++ b/client.go
@@ -163,12 +163,15 @@ type client struct {
 	lastID            int64
 	backoffFactory    func() backoff.BackOff
 	cancelFunc        context.CancelFunc
+	wg                sync.WaitGroup
 }
 
 func (c *client) Start() {
 	c.setState(ClientConnecting)
 	boff := c.backoffFactory()
+	c.wg.Add(1)
 	go func() {
+		defer c.wg.Done()
 		for {
 			c.setErr(nil)
 			// Listen for state change to ClientConnected and signal backoff Reset then.
@@ -221,9 +224,8 @@ func (c *client) Start() {
 func (c *client) Stop() {
 	if c.cancelFunc != nil {
 		c.cancelFunc()
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) // in practice, it is faster than 5 seconds so this is just to avoid infinite block
-		defer cancel()
-		c.WaitForState(ctx, ClientClosed)
+		c.wg.Wait()
+		c.setState(ClientClosed)
 	}
 }