fix(panic): partition offset manager

Signed-off-by: napallday <[email protected]>

Author: napallday

balance_strategy.go

@@ -19,6 +19,7 @@ const (
 	// StickyBalanceStrategyName identifies strategies that use the sticky-partition assignment strategy
 	StickyBalanceStrategyName = "sticky"
 
+	// CooperativeStickyBalanceStrategyName identifies strategies that use the cooperative sticky-partition assignment strategy
 	CooperativeStickyBalanceStrategyName = "cooperative-sticky"
 
 	defaultGeneration = -1

offset_manager.go

@@ -140,6 +140,13 @@ func (om *offsetManager) ManagePartition(topic string, partition int32) (Partiti
 }
 
 func (om *offsetManager) RemovePartitions(topicPartitions map[string][]int32) error {
+	// flush one last time
+	if om.conf.Consumer.Offsets.AutoCommit.Enable {
+		for attempt := 0; attempt <= om.conf.Consumer.Offsets.Retry.Max; attempt++ {
+			om.flushToBroker()
+		}
+	}
+
 	var errs ConsumerErrors
 	var errsLock sync.Mutex
 
@@ -150,10 +157,15 @@ func (om *offsetManager) RemovePartitions(topicPartitions map[string][]int32) er
 			go func(topic string, partition int32) {
 				defer wg.Done()
 
-				om.pomsLock.RLock()
+				om.pomsLock.Lock()
 				pom := om.poms[topic][partition]
-				om.pomsLock.RUnlock()
 				err := pom.Close()
+				delete(om.poms[topic], partition)
+				if len(om.poms[topic]) == 0 {
+					delete(om.poms, topic)
+				}
+				om.pomsLock.Unlock()
+
 				if err != nil {
 					errsLock.Lock()
 					var consumerErrs ConsumerErrors
@@ -167,24 +179,6 @@ func (om *offsetManager) RemovePartitions(topicPartitions map[string][]int32) er
 	}
 	wg.Wait()
 
-	// flush one last time
-	if om.conf.Consumer.Offsets.AutoCommit.Enable {
-		for attempt := 0; attempt <= om.conf.Consumer.Offsets.Retry.Max; attempt++ {
-			om.flushToBroker()
-		}
-	}
-
-	om.pomsLock.Lock()
-	for topic, partitions := range topicPartitions {
-		for _, partition := range partitions {
-			delete(om.poms[topic], partition)
-			if len(om.poms[topic]) == 0 {
-				delete(om.poms, topic)
-			}
-		}
-	}
-
-	om.pomsLock.Unlock()
 	if len(errs) > 0 {
 		return errs
 	}