add wgWaitGroup

.gitignore

@@ -4,6 +4,7 @@
 *.dll
 *.so
 *.dylib
+bin
 
 # Test binary, built with `go test -c`
 *.test
@@ -13,3 +14,11 @@
 
 # Dependency directories (remove the comment below to include it)
 # vendor/
+
+# Go workspace file
+go.work
+
+# General
+.DS_Store
+.idea
+.vscode

go.mod

@@ -2,6 +2,10 @@ module go.unistack.org/micro-broker-kgo/v3
 
 go 1.19
 
+replace (
+	go.unistack.org/micro/v3 => ../micro
+)
+
 require (
 	github.com/google/uuid v1.6.0
 	github.com/twmb/franz-go v1.16.1

subscriber.go

@@ -106,6 +106,12 @@ func (s *subscriber) killConsumers(ctx context.Context, lost map[string][]int32)
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
+	if s.kopts.Wait != nil {
+		ctx, cancel := context.WithTimeout(ctx, s.kopts.GracefulTimeout)
+		defer cancel()
+		s.kopts.Wait.WaitContext(ctx)
+	}
+
 	for topic, partitions := range lost {
 		for _, partition := range partitions {
 			tp := tp{topic, partition}
@@ -148,7 +154,14 @@ func (s *subscriber) assigned(_ context.Context, c *kgo.Client, assigned map[str
 				opts:    s.opts,
 			}
 			s.consumers[tp{topic, partition}] = pc
-			go pc.consume()
+
+			go func() {
+				if s.kopts.Wait != nil {
+					s.kopts.Wait.Add(1)
+					defer s.kopts.Wait.Done()
+				}
+				pc.consume()
+			}()
 		}
 	}
 }