[fix] consumer map race

consumer/worker.go

@@ -5,7 +5,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/aliyun/aliyun-log-go-sdk"
+	sls "github.com/aliyun/aliyun-log-go-sdk"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
 	"go.uber.org/atomic"
@@ -16,7 +16,7 @@ type ConsumerWorker struct {
 	consumerHeatBeat   *ConsumerHeatBeat
 	client             *ConsumerClient
 	workerShutDownFlag *atomic.Bool
-	shardConsumer      map[int]*ShardConsumerWorker
+	shardConsumer      sync.Map // map[int]*ShardConsumerWorker
 	do                 func(shard int, logGroup *sls.LogGroupList) string
 	waitGroup          sync.WaitGroup
 	Logger             log.Logger
@@ -30,9 +30,9 @@ func InitConsumerWorker(option LogHubConfig, do func(int, *sls.LogGroupList) str
 		consumerHeatBeat:   consumerHeatBeat,
 		client:             consumerClient,
 		workerShutDownFlag: atomic.NewBool(false),
-		shardConsumer:      make(map[int]*ShardConsumerWorker),
-		do:                 do,
-		Logger:             logger,
+		//shardConsumer:      make(map[int]*ShardConsumerWorker),
+		do:     do,
+		Logger: logger,
 	}
 	consumerClient.createConsumerGroup()
 	return consumerWorker
@@ -82,50 +82,62 @@ func (consumerWorker *ConsumerWorker) run() {
 func (consumerWorker *ConsumerWorker) shutDownAndWait() {
 	for {
 		time.Sleep(500 * time.Millisecond)
-		for shard, consumer := range consumerWorker.shardConsumer {
-			if !consumer.isShutDownComplete() {
-				consumer.consumerShutDown()
-			} else if consumer.isShutDownComplete() {
-				delete(consumerWorker.shardConsumer, shard)
-			}
-		}
-		if len(consumerWorker.shardConsumer) == 0 {
+		count := 0
+		consumerWorker.shardConsumer.Range(
+			func(key, value interface{}) bool {
+				count++
+				consumer := value.(*ShardConsumerWorker)
+				if !consumer.isShutDownComplete() {
+					consumer.consumerShutDown()
+				} else if consumer.isShutDownComplete() {
+					consumerWorker.shardConsumer.Delete(key)
+				}
+				return true
+			},
+		)
+		if count == 0 {
 			break
 		}
 	}
 
 }
 
 func (consumerWorker *ConsumerWorker) getShardConsumer(shardId int) *ShardConsumerWorker {
-	consumer := consumerWorker.shardConsumer[shardId]
-	if consumer != nil {
-		return consumer
+	consumer, ok := consumerWorker.shardConsumer.Load(shardId)
+	if ok {
+		return consumer.(*ShardConsumerWorker)
 	}
-	consumer = initShardConsumerWorker(shardId, consumerWorker.client, consumerWorker.do, consumerWorker.Logger)
-	consumerWorker.shardConsumer[shardId] = consumer
-	return consumer
+	consumerIns := initShardConsumerWorker(shardId, consumerWorker.client, consumerWorker.do, consumerWorker.Logger)
+	consumerWorker.shardConsumer.Store(shardId, consumerIns)
+	return consumerIns
 
 }
 
 func (consumerWorker *ConsumerWorker) cleanShardConsumer(owned_shards []int) {
-	for shard, consumer := range consumerWorker.shardConsumer {
 
-		if !Contain(shard, owned_shards) {
-			level.Info(consumerWorker.Logger).Log("msg", "try to call shut down for unassigned consumer shard", "shardId", shard)
-			consumer.consumerShutDown()
-			level.Info(consumerWorker.Logger).Log("msg", "Complete call shut down for unassigned consumer shard", "shardId", shard)
-		}
+	consumerWorker.shardConsumer.Range(
+		func(key, value interface{}) bool {
+			shard := key.(int)
+			consumer := value.(*ShardConsumerWorker)
 
-		if consumer.isShutDownComplete() {
-			isDeleteShard := consumerWorker.consumerHeatBeat.removeHeartShard(shard)
-			if isDeleteShard {
-				level.Info(consumerWorker.Logger).Log("msg", "Remove an assigned consumer shard", "shardId", shard)
-				delete(consumerWorker.shardConsumer, shard)
-			} else {
-				level.Info(consumerWorker.Logger).Log("msg", "Remove an assigned consumer shard failed", "shardId", shard)
+			if !Contain(shard, owned_shards) {
+				level.Info(consumerWorker.Logger).Log("msg", "try to call shut down for unassigned consumer shard", "shardId", shard)
+				consumer.consumerShutDown()
+				level.Info(consumerWorker.Logger).Log("msg", "Complete call shut down for unassigned consumer shard", "shardId", shard)
 			}
-		}
-	}
+
+			if consumer.isShutDownComplete() {
+				isDeleteShard := consumerWorker.consumerHeatBeat.removeHeartShard(shard)
+				if isDeleteShard {
+					level.Info(consumerWorker.Logger).Log("msg", "Remove an assigned consumer shard", "shardId", shard)
+					consumerWorker.shardConsumer.Delete(shard)
+				} else {
+					level.Info(consumerWorker.Logger).Log("msg", "Remove an assigned consumer shard failed", "shardId", shard)
+				}
+			}
+			return true
+		},
+	)
 }
 
 // This function is used to initialize the global log configuration

producer/io_thread_pool.go

@@ -36,6 +36,9 @@ func (threadPool *IoThreadPool) addTask(batch *ProducerBatch) {
 func (threadPool *IoThreadPool) popTask() *ProducerBatch {
 	defer threadPool.lock.Unlock()
 	threadPool.lock.Lock()
+	if threadPool.queue.Len() <= 0 {
+		return nil
+	}
 	ele := threadPool.queue.Front()
 	threadPool.queue.Remove(ele)
 	return ele.Value.(*ProducerBatch)
@@ -50,12 +53,12 @@ func (threadPool *IoThreadPool) hasTask() bool {
 func (threadPool *IoThreadPool) start(ioWorkerWaitGroup *sync.WaitGroup, ioThreadPoolwait *sync.WaitGroup) {
 	defer ioThreadPoolwait.Done()
 	for {
-		if threadPool.hasTask() {
-			select {
-			case threadPool.ioworker.maxIoWorker <- 1:
-				ioWorkerWaitGroup.Add(1)
-				go threadPool.ioworker.sendToServer(threadPool.popTask(), ioWorkerWaitGroup)
-			}
+		if task := threadPool.popTask(); task != nil {
+			threadPool.ioworker.startSendTask(ioWorkerWaitGroup)
+			go func(producerBatch *ProducerBatch) {
+				defer threadPool.ioworker.closeSendTask(ioWorkerWaitGroup)
+				threadPool.ioworker.sendToServer(producerBatch)
+			}(task)
 		} else {
 			if !threadPool.threadPoolShutDownFlag.Load() {
 				time.Sleep(100 * time.Millisecond)

producer/io_worker.go

@@ -41,14 +41,9 @@ func initIoWorker(client sls.ClientInterface, retryQueue *RetryQueue, logger log
 	}
 }
 
-func (ioWorker *IoWorker) sendToServer(producerBatch *ProducerBatch, ioWorkerWaitGroup *sync.WaitGroup) {
-	if producerBatch == nil || ioWorkerWaitGroup == nil {
-		return
-	}
+func (ioWorker *IoWorker) sendToServer(producerBatch *ProducerBatch) {
 	level.Debug(ioWorker.logger).Log("msg", "ioworker send data to server")
-	defer ioWorker.closeSendTask(ioWorkerWaitGroup)
 	var err error
-	atomic.AddInt64(&ioWorker.taskCount, 1)
 	if producerBatch.shardHash != nil {
 		err = ioWorker.client.PostLogStoreLogs(producerBatch.getProject(), producerBatch.getLogstore(), producerBatch.logGroup, producerBatch.getShardHash())
 	} else {
@@ -116,9 +111,15 @@ func (ioWorker *IoWorker) addErrorMessageToBatchAttempt(producerBatch *ProducerB
 }
 
 func (ioWorker *IoWorker) closeSendTask(ioWorkerWaitGroup *sync.WaitGroup) {
-	ioWorkerWaitGroup.Done()
-	atomic.AddInt64(&ioWorker.taskCount, -1)
 	<-ioWorker.maxIoWorker
+	atomic.AddInt64(&ioWorker.taskCount, -1)
+	ioWorkerWaitGroup.Done()
+}
+
+func (ioWorker *IoWorker) startSendTask(ioWorkerWaitGroup *sync.WaitGroup) {
+	atomic.AddInt64(&ioWorker.taskCount, 1)
+	ioWorker.maxIoWorker <- 1
+	ioWorkerWaitGroup.Add(1)
 }
 
 func (ioWorker *IoWorker) excuteFailedCallback(producerBatch *ProducerBatch) {