Protect partitionsHeld in consumer by lock

Signed-off-by: Ivan Babrou <[email protected]>
jaegertracing · Apr 23, 2019 · 2582040 · 2582040
1 parent fa84803
commit 2582040
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/cmd/ingester/app/consumer/consumer.go b/cmd/ingester/app/consumer/consumer.go
@@ -110,17 +110,19 @@ func (c *Consumer) Close() error {
 
 func (c *Consumer) handleMessages(pc sc.PartitionConsumer) {
 	c.logger.Info("Starting message handler", zap.Int32("partition", pc.Partition()))
+	c.partitionMapLock.Lock()
 	c.partitionsHeld++
 	c.partitionsHeldGauge.Update(c.partitionsHeld)
+	wg := &c.partitionIDToState[pc.Partition()].wg
+	c.partitionMapLock.Unlock()
 	defer func() {
+		c.closePartition(pc)
+		wg.Done()
+		c.partitionMapLock.Lock()
 		c.partitionsHeld--
 		c.partitionsHeldGauge.Update(c.partitionsHeld)
+		c.partitionMapLock.Unlock()
 	}()
-	c.partitionMapLock.Lock()
-	wg := &c.partitionIDToState[pc.Partition()].wg
-	c.partitionMapLock.Unlock()
-	defer wg.Done()
-	defer c.closePartition(pc)
 
 	msgMetrics := c.newMsgMetrics(pc.Partition())