Fix issue #167 - Concurrent write/write of in-flight shard map (#168)

Co-authored-by: Sanket Deshpande <sanket@clearblade.com>

consumer.go

@@ -118,18 +118,18 @@ func (c *Consumer) Scan(ctx context.Context, fn ScanFunc) error {
 
 	wg := new(sync.WaitGroup)
 	// process each of the shards
-	shardsInProcess := make(map[string]struct{})
+	s := newShardsInProcess()
 	for shard := range shardC {
 		shardId := aws.ToString(shard.ShardId)
-		if _, ok := shardsInProcess[shardId]; ok {
+		if s.doesShardExist(shardId) {
 			// safetynet: if shard already in process by another goroutine, just skipping the request
 			continue
 		}
 		wg.Add(1)
 		go func(shardID string) {
-			shardsInProcess[shardID] = struct{}{}
+			s.addShard(shardID)
 			defer func() {
-				delete(shardsInProcess, shardID)
+				s.deleteShard(shardID)
 			}()
 			defer wg.Done()
 			var err error
@@ -311,3 +311,24 @@ func isRetriableError(err error) bool {
 func isShardClosed(nextShardIterator, currentShardIterator *string) bool {
 	return nextShardIterator == nil || currentShardIterator == nextShardIterator
 }
+
+type shards struct {
+	shardsInProcess sync.Map
+}
+
+func newShardsInProcess() *shards {
+	return &shards{}
+}
+
+func (s *shards) addShard(shardId string) {
+	s.shardsInProcess.Store(shardId, struct{}{})
+}
+
+func (s *shards) doesShardExist(shardId string) bool {
+	_, ok := s.shardsInProcess.Load(shardId)
+	return ok
+}
+
+func (s *shards) deleteShard(shardId string) {
+	s.shardsInProcess.Delete(shardId)
+}