diff --git a/sources/dynamodb/offsets/offsets.go b/sources/dynamodb/offsets/offsets.go index f0ac6880..9b8a9e53 100644 --- a/sources/dynamodb/offsets/offsets.go +++ b/sources/dynamodb/offsets/offsets.go @@ -14,6 +14,23 @@ type OffsetStorage struct { ttlMap *ttlmap.TTLMap } +func shardSeenKey(shardId string) string { + return fmt.Sprintf("seen#shardId#%s", shardId) +} + +func (o *OffsetStorage) SetShardSeen(shardID string) { + o.ttlMap.Set(ttlmap.SetArgs{ + Key: shardSeenKey(shardID), + Value: true, + DoNotFlushToDisk: true, + }, ShardExpirationAndBuffer) +} + +func (o *OffsetStorage) GetShardSeen(shardID string) bool { + _, isOk := o.ttlMap.Get(shardSeenKey(shardID)) + return isOk +} + func shardProcessingKey(shardId string) string { return fmt.Sprintf("processing#shardId#%s", shardId) } diff --git a/sources/dynamodb/shard.go b/sources/dynamodb/shard.go index 3e39781e..18062d7d 100644 --- a/sources/dynamodb/shard.go +++ b/sources/dynamodb/shard.go @@ -29,9 +29,9 @@ func (s *StreamStore) processShard(ctx context.Context, shard *dynamodbstreams.S } if parentID := shard.ParentShardId; parentID != nil { - // If the parent shard exists, is it still being processed? If so, let's wait a bit and then retry. - // We must process the parent shard first before processing the child shard. - if s.storage.GetShardProcessing(*parentID) && !s.storage.GetShardProcessed(*parentID) { + // Have we seen the parent? If so, let's wait for processing to finish + // If we haven't seen the parent, then we can assume this is the parent and we don't need to wait. + if s.storage.GetShardSeen(*parentID) && !s.storage.GetShardProcessed(*parentID) { slog.Info("Parent shard is being processed, let's sleep 3s and retry", slog.String("shardId", *shard.ShardId), slog.String("parentShardId", *parentID)) time.Sleep(3 * time.Second) s.processShard(ctx, shard, writer) diff --git a/sources/dynamodb/stream.go b/sources/dynamodb/stream.go index 4dd205ab..e74707da 100644 --- a/sources/dynamodb/stream.go +++ b/sources/dynamodb/stream.go @@ -65,6 +65,11 @@ func (s *StreamStore) scanForNewShards() error { return fmt.Errorf("failed to describe stream: %w", err) } + // We need two loops because we need to mark all the shards as "SEEN" before we process. + for _, shard := range result.StreamDescription.Shards { + s.storage.SetShardSeen(*shard.ShardId) + } + for _, shard := range result.StreamDescription.Shards { s.shardChan <- shard }