Fix a rare panic during consumer group rebalance (#791)

* Fix a rare panic during consumer group rebalance

Currently, there is a sync.WaitGroup that is used for process accounting
with go routines started by Generation.Start.  It is used when a
generation is ending to ensure that all of the Start-ed functions exit
before the next generation can start.  There is an edge case that can
cause panic as written up in #786 due to unsafe use of the WaitGroup.
Basically, it is possible that close() is already in process and is
waiting on WaitGroup.Wait() when Start() comes in and calls
WaitGroup.Add(1).

The contract of Start() is that the provided function is to be run.
Code in the wild may depend on this behavior, so it's not an option to
return from Start() without runing the provided function.

Accordingly, this PR updates the code to coordinate between close() and
Start() such that the panic case is no longer possible while preserving
the existing contract.  It uses channels and a mutex in order to create
two cases: the normal case where the generation is alive and the edge
case where the generation has already ended.

* Update consumergroup.go

Co-authored-by: Achille <[email protected]>

Co-authored-by: Achille <[email protected]>

Author: Steve van Loben Sels

consumergroup.go

@@ -322,9 +322,17 @@ type Generation struct {
 
 	conn coordinator
 
-	once sync.Once
-	done chan struct{}
-	wg   sync.WaitGroup
+	// the following fields are used for process accounting to synchronize
+	// between Start and close.  lock protects all of them.  done is closed
+	// when the generation is ending in order to signal that the generation
+	// should start self-desructing.  closed protects against double-closing
+	// the done chan.  routines is a count of runing go routines that have been
+	// launched by Start.  joined will be closed by the last go routine to exit.
+	lock     sync.Mutex
+	done     chan struct{}
+	closed   bool
+	routines int
+	joined   chan struct{}
 
 	retentionMillis int64
 	log             func(func(Logger))
@@ -334,10 +342,21 @@ type Generation struct {
 // close stops the generation and waits for all functions launched via Start to
 // terminate.
 func (g *Generation) close() {
-	g.once.Do(func() {
+	g.lock.Lock()
+	if !g.closed {
 		close(g.done)
-	})
-	g.wg.Wait()
+		g.closed = true
+	}
+	// determine whether any go routines are running that we need to wait for.
+	// waiting needs to happen outside of the critical section.
+	r := g.routines
+	g.lock.Unlock()
+
+	// NOTE: r will be zero if no go routines were ever launched.  no need to
+	// wait in that case.
+	if r > 0 {
+		<-g.joined
+	}
 }
 
 // Start launches the provided function in a go routine and adds accounting such
@@ -354,15 +373,42 @@ func (g *Generation) close() {
 // progress for this consumer and potentially cause consumer group membership
 // churn.
 func (g *Generation) Start(fn func(ctx context.Context)) {
-	g.wg.Add(1)
+	g.lock.Lock()
+	defer g.lock.Unlock()
+
+	// this is an edge case: if the generation has already closed, then it's
+	// possible that the close func has already waited on outstanding go
+	// routines and exited.
+	//
+	// nonetheless, it's important to honor that the fn is invoked in case the
+	// calling function is waiting e.g. on a channel send or a WaitGroup.  in
+	// such a case, fn should immediately exit because ctx.Err() will return
+	// ErrGenerationEnded.
+	if g.closed {
+		go fn(genCtx{g})
+		return
+	}
+
+	// register that there is one more go routine that's part of this gen.
+	g.routines++
+
 	go func() {
 		fn(genCtx{g})
+		g.lock.Lock()
 		// shut down the generation as soon as one function exits.  this is
-		// different from close() in that it doesn't wait on the wg.
-		g.once.Do(func() {
+		// different from close() in that it doesn't wait for all go routines in
+		// the generation to exit.
+		if !g.closed {
 			close(g.done)
-		})
-		g.wg.Done()
+			g.closed = true
+		}
+		g.routines--
+		// if this was the last go routine in the generation, close the joined
+		// chan so that close() can exit if it's waiting.
+		if g.routines == 0 {
+			close(g.joined)
+		}
+		g.lock.Unlock()
 	}()
 }
 
@@ -781,6 +827,7 @@ func (cg *ConsumerGroup) nextGeneration(memberID string) (string, error) {
 		Assignments:     cg.makeAssignments(assignments, offsets),
 		conn:            conn,
 		done:            make(chan struct{}),
+		joined:          make(chan struct{}),
 		retentionMillis: int64(cg.config.RetentionTime / time.Millisecond),
 		log:             cg.withLogger,
 		logError:        cg.withErrorLogger,

consumergroup_test.go

@@ -630,6 +630,7 @@ func TestGenerationExitsOnPartitionChange(t *testing.T) {
 	gen := Generation{
 		conn:     conn,
 		done:     make(chan struct{}),
+		joined:   make(chan struct{}),
 		log:      func(func(Logger)) {},
 		logError: func(func(Logger)) {},
 	}
@@ -649,3 +650,30 @@ func TestGenerationExitsOnPartitionChange(t *testing.T) {
 		}
 	}
 }
+
+func TestGenerationStartsFunctionAfterClosed(t *testing.T) {
+	gen := Generation{
+		conn:     &mockCoordinator{},
+		done:     make(chan struct{}),
+		joined:   make(chan struct{}),
+		log:      func(func(Logger)) {},
+		logError: func(func(Logger)) {},
+	}
+
+	gen.close()
+
+	ch := make(chan error)
+	gen.Start(func(ctx context.Context) {
+		<-ctx.Done()
+		ch <- ctx.Err()
+	})
+
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for func to run")
+	case err := <-ch:
+		if err != ErrGenerationEnded {
+			t.Fatalf("expected %v but got %v", ErrGenerationEnded, err)
+		}
+	}
+}