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
-	s := newShardsInProcess()
+	shardsInProcess := make(map[string]struct{})
 	for shard := range shardC {
 		shardId := aws.ToString(shard.ShardId)
-		if s.doesShardExist(shardId) {
+		if _, ok := shardsInProcess[shardId]; ok {
 			// safetynet: if shard already in process by another goroutine, just skipping the request
 			continue
 		}
 		wg.Add(1)
 		go func(shardID string) {
-			s.addShard(shardID)
+			shardsInProcess[shardID] = struct{}{}
 			defer func() {
-				s.deleteShard(shardID)
+				delete(shardsInProcess, shardID)
 			}()
 			defer wg.Done()
 			var err error
@@ -311,24 +311,3 @@ 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)
-}