consumer.go

@@ -313,22 +313,32 @@ func isShardClosed(nextShardIterator, currentShardIterator *string) bool {
 }
 
 type shards struct {
-	shardsInProcess sync.Map
+	*sync.RWMutex
+	shardsInProcess map[string]struct{}
 }
 
 func newShardsInProcess() *shards {
-	return &shards{}
+	return &shards{
+		RWMutex:         &sync.RWMutex{},
+		shardsInProcess: make(map[string]struct{}),
+	}
 }
 
 func (s *shards) addShard(shardId string) {
-	s.shardsInProcess.Store(shardId, struct{}{})
+	s.Lock()
+	defer s.Unlock()
+	s.shardsInProcess[shardId] = struct{}{}
 }
 
 func (s *shards) doesShardExist(shardId string) bool {
-	_, ok := s.shardsInProcess.Load(shardId)
+	s.RLock()
+	defer s.RUnlock()
+	_, ok := s.shardsInProcess[shardId]
 	return ok
 }
 
 func (s *shards) deleteShard(shardId string) {
-	s.shardsInProcess.Delete(shardId)
+	s.Lock()
+	defer s.Unlock()
+	delete(s.shardsInProcess, shardId)
 }