DynamoDB - Respect lineage

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)
 }

sources/dynamodb/shard.go

@@ -23,13 +23,22 @@ func (s *StreamStore) ListenToChannel(ctx context.Context, writer writers.Writer
 }
 
 func (s *StreamStore) processShard(ctx context.Context, shard *dynamodbstreams.Shard, writer writers.Writer) {
-	var attempts int
-
 	// Is there another go-routine processing this shard?
 	if s.storage.GetShardProcessing(*shard.ShardId) {
 		return
 	}
 
+	if parentID := shard.ParentShardId; parentID != nil {
+		// 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)
+			return
+		}
+	}
+
 	// If no one is processing it, let's mark it as being processed.
 	s.storage.SetShardProcessing(*shard.ShardId)
 	if s.storage.GetShardProcessed(*shard.ShardId) {
@@ -103,6 +112,7 @@ func (s *StreamStore) processShard(ctx context.Context, shard *dynamodbstreams.S
 			logger.Panic("Failed to publish messages, exiting...", slog.Any("err", err))
 		}
 
+		var attempts int
 		if len(getRecordsOutput.Records) > 0 {
 			attempts = 0
 			lastRecord := getRecordsOutput.Records[len(getRecordsOutput.Records)-1]

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
 		}