From ee3d7c11c8197e0ca8f758409e0f732b315600d2 Mon Sep 17 00:00:00 2001
From: Mateusz Czubak <mateusz@firebolt.io>
Date: Mon, 15 May 2023 11:08:44 +0200
Subject: [PATCH] kgo: fix race condition in close

---
 pkg/kgo/client.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pkg/kgo/client.go b/pkg/kgo/client.go
index 8a36433c..b421932b 100644
--- a/pkg/kgo/client.go
+++ b/pkg/kgo/client.go
@@ -948,7 +948,7 @@ func (cl *Client) Close() {
 
 	sessCloseCtx, sessCloseCancel := context.WithTimeout(cl.ctx, time.Second)
 	var wg sync.WaitGroup
-	for _, sns := range cl.sinksAndSources {
+	cl.allSinksAndSources(func(sns sinkAndSource) {
 		if sns.source.session.id != 0 {
 			sns := sns
 			wg.Add(1)
@@ -957,7 +957,7 @@ func (cl *Client) Close() {
 				sns.source.killSessionOnClose(sessCloseCtx)
 			}()
 		}
-	}
+	})
 	wg.Wait()
 	sessCloseCancel()