[FLINK-36523][Connectors/DynamoDB][BugFix] Fixing LATEST behaviour fo…

…r DDB connector

flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/enumerator/DynamoDbStreamsSourceEnumerator.java

@@ -42,6 +42,7 @@
 import javax.annotation.Nullable;
 
 import java.io.IOException;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -75,6 +76,7 @@ public class DynamoDbStreamsSourceEnumerator
     private final DynamoDbStreamsShardAssigner shardAssigner;
     private final ShardAssignerContext shardAssignerContext;
     private final SplitTracker splitTracker;
+    private final Instant startTimestamp;
 
     private final Map<Integer, Set<DynamoDbStreamsShardSplit>> splitAssignment = new HashMap<>();
 
@@ -93,10 +95,16 @@ public DynamoDbStreamsSourceEnumerator(
         this.shardAssignerContext = new ShardAssignerContext(splitAssignment, context);
         InitialPosition initialPosition = sourceConfig.get(STREAM_INITIAL_POSITION);
         if (state == null) {
-            this.splitTracker = new SplitTracker(streamArn, initialPosition);
+            this.startTimestamp = Instant.now();
+            this.splitTracker = new SplitTracker(streamArn, initialPosition, this.startTimestamp);
         } else {
             this.splitTracker =
-                    new SplitTracker(state.getKnownSplits(), streamArn, initialPosition);
+                    new SplitTracker(
+                            state.getKnownSplits(),
+                            streamArn,
+                            initialPosition,
+                            state.getStartTimestamp());
+            this.startTimestamp = state.getStartTimestamp();
         }
     }
 
@@ -231,7 +239,7 @@ public void addReader(int subtaskId) {
     public DynamoDbStreamsSourceEnumeratorState snapshotState(long checkpointId) throws Exception {
         List<DynamoDBStreamsShardSplitWithAssignmentStatus> splitStates =
                 splitTracker.snapshotState(checkpointId);
-        return new DynamoDbStreamsSourceEnumeratorState(splitStates);
+        return new DynamoDbStreamsSourceEnumeratorState(splitStates, startTimestamp);
     }
 
     @Override

flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/enumerator/DynamoDbStreamsSourceEnumeratorState.java

@@ -20,6 +20,7 @@
 
 import org.apache.flink.annotation.Internal;
 
+import java.time.Instant;
 import java.util.List;
 
 /**
@@ -29,13 +30,20 @@
 @Internal
 public class DynamoDbStreamsSourceEnumeratorState {
     private final List<DynamoDBStreamsShardSplitWithAssignmentStatus> knownSplits;
+    private final Instant startTimestamp;
 
     public DynamoDbStreamsSourceEnumeratorState(
-            List<DynamoDBStreamsShardSplitWithAssignmentStatus> knownSplits) {
+            final List<DynamoDBStreamsShardSplitWithAssignmentStatus> knownSplits,
+            final Instant startTimestamp) {
         this.knownSplits = knownSplits;
+        this.startTimestamp = startTimestamp;
     }
 
     public List<DynamoDBStreamsShardSplitWithAssignmentStatus> getKnownSplits() {
         return knownSplits;
     }
+
+    public Instant getStartTimestamp() {
+        return startTimestamp;
+    }
 }

flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/enumerator/DynamoDbStreamsSourceEnumeratorStateSerializer.java

@@ -29,6 +29,7 @@
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -66,6 +67,7 @@ public byte[] serialize(
                 out.write(serializedSplit);
                 out.writeInt(split.assignmentStatus().getStatusCode());
             }
+            out.writeLong(dynamoDbStreamsSourceEnumeratorState.getStartTimestamp().toEpochMilli());
 
             out.flush();
 
@@ -109,11 +111,12 @@ public DynamoDbStreamsSourceEnumeratorState deserialize(
                 }
             }
 
+            Instant startTimestamp = Instant.ofEpochMilli(in.readLong());
             if (in.available() > 0) {
                 throw new IOException("Unexpected trailing bytes when deserializing.");
             }
 
-            return new DynamoDbStreamsSourceEnumeratorState(knownSplits);
+            return new DynamoDbStreamsSourceEnumeratorState(knownSplits, startTimestamp);
         }
     }
 }

flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/enumerator/tracker/SplitTracker.java

@@ -32,15 +32,18 @@
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.dynamodb.model.Shard;
 
+import java.time.Instant;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
+import static org.apache.flink.connector.dynamodb.source.config.DynamodbStreamsSourceConfigConstants.InitialPosition.LATEST;
 import static org.apache.flink.connector.dynamodb.source.config.DynamodbStreamsSourceConfigConstants.InitialPosition.TRIM_HORIZON;
 import static org.apache.flink.connector.dynamodb.source.enumerator.SplitAssignmentStatus.ASSIGNED;
 import static org.apache.flink.connector.dynamodb.source.enumerator.SplitAssignmentStatus.FINISHED;
@@ -57,18 +60,21 @@ public class SplitTracker {
     private final Set<String> finishedSplits = new HashSet<>();
     private final String streamArn;
     private final InitialPosition initialPosition;
+    private final Instant startTimestamp;
     private static final Logger LOG = LoggerFactory.getLogger(SplitTracker.class);
 
-    public SplitTracker(String streamArn, InitialPosition initialPosition) {
-        this(Collections.emptyList(), streamArn, initialPosition);
+    public SplitTracker(String streamArn, InitialPosition initialPosition, Instant startTimestamp) {
+        this(Collections.emptyList(), streamArn, initialPosition, startTimestamp);
     }
 
     public SplitTracker(
             List<DynamoDBStreamsShardSplitWithAssignmentStatus> initialState,
             String streamArn,
-            DynamodbStreamsSourceConfigConstants.InitialPosition initialPosition) {
+            DynamodbStreamsSourceConfigConstants.InitialPosition initialPosition,
+            Instant startTimestamp) {
         this.streamArn = streamArn;
         this.initialPosition = initialPosition;
+        this.startTimestamp = startTimestamp;
         initialState.forEach(
                 splitWithStatus -> {
                     DynamoDbStreamsShardSplit currentSplit = splitWithStatus.split();
@@ -89,13 +95,71 @@ public SplitTracker(
      * @param shardsToAdd collection of splits to add to tracking
      */
     public void addSplits(Collection<Shard> shardsToAdd) {
-        Set<String> discoveredShardIds =
-                shardsToAdd.stream().map(Shard::shardId).collect(Collectors.toSet());
+        if (TRIM_HORIZON.equals(initialPosition)) {
+            addSplitsForTrimHorizon(shardsToAdd);
+            return;
+        }
+        addSplitsForLatest(shardsToAdd);
+    }
+
+    private void addSplitsForLatest(Collection<Shard> shardsToAdd) {
+        List<Shard> openShards =
+                shardsToAdd.stream()
+                        .filter(shard -> shard.sequenceNumberRange().endingSequenceNumber() == null)
+                        .collect(Collectors.toList());
+        Map<String, Shard> shardIdToShardMap =
+                shardsToAdd.stream().collect(Collectors.toMap(Shard::shardId, shard -> shard));
+        for (Shard shard : openShards) {
+            String shardId = shard.shardId();
+            if (knownSplits.containsKey(shardId)) {
+                continue;
+            }
+            String firstShardIdBeforeTimestamp =
+                    findFirstShardIdBeforeTimestamp(shardId, shardIdToShardMap);
+            putAllAncestorShardsTillFirstShardId(
+                    shardId, firstShardIdBeforeTimestamp, shardIdToShardMap);
+        }
+    }
+
+    private void putAllAncestorShardsTillFirstShardId(
+            String shardId,
+            String firstShardIdBeforeTimestamp,
+            Map<String, Shard> shardIdToShardMap) {
+        String currentShardId = shardId;
+        while (currentShardId != null
+                && shardIdToShardMap.containsKey(currentShardId)
+                && !knownSplits.containsKey(currentShardId)) {
+            Shard currentShard = shardIdToShardMap.get(currentShardId);
+            if (Objects.equals(currentShardId, firstShardIdBeforeTimestamp)) {
+                DynamoDbStreamsShardSplit newSplit = mapToSplit(currentShard, LATEST);
+                knownSplits.putIfAbsent(currentShardId, newSplit);
+                break;
+            } else {
+                DynamoDbStreamsShardSplit newSplit = mapToSplit(currentShard, TRIM_HORIZON);
+                knownSplits.putIfAbsent(currentShardId, newSplit);
+            }
+            currentShardId = currentShard.parentShardId();
+        }
+    }
+
+    private String findFirstShardIdBeforeTimestamp(
+            String shardIdToStartWith, Map<String, Shard> shardIdToShardMap) {
+        String shardId = shardIdToStartWith;
+        while (shardId != null && shardIdToShardMap.containsKey(shardId)) {
+            Shard shard = shardIdToShardMap.get(shardId);
+            if (ShardUtils.isShardCreatedBeforeTimestamp(shardId, startTimestamp)) {
+                return shardId;
+            }
+            shardId = shard.parentShardId();
+        }
+        return null;
+    }
+
+    private void addSplitsForTrimHorizon(Collection<Shard> shardsToAdd) {
         for (Shard shard : shardsToAdd) {
             String shardId = shard.shardId();
             if (!knownSplits.containsKey(shardId)) {
-                DynamoDbStreamsShardSplit newSplit =
-                        mapToSplit(shard, getStartingPosition(shard, discoveredShardIds));
+                DynamoDbStreamsShardSplit newSplit = mapToSplit(shard, TRIM_HORIZON);
                 knownSplits.put(shardId, newSplit);
                 addSplitToMapping(newSplit);
             }
@@ -111,27 +175,6 @@ private void addSplitToMapping(DynamoDbStreamsShardSplit split) {
                 .add(split.splitId());
     }
 
-    /**
-     * If there is no parent in knownSplits of the current shard, that means that the parent has
-     * been read already, so we start the current shard with the specified initial position. We do
-     * this because we can't really differentiate the case of read from an expired snapshot where
-     * there is a node with parent which was trimmed but a grandparent vs a happy case where there
-     * is a node with no parent. This means that for cases when we are reading from an expired
-     * snapshot, we'll start the first node in the lineage with LATEST shard iterator type.
-     *
-     * @param shard Shard whose starting position has to be determined.
-     */
-    private InitialPosition getStartingPosition(Shard shard, Set<String> discoveredShardIds) {
-        if (shard.parentShardId() == null) {
-            return initialPosition;
-        }
-        if (!knownSplits.containsKey(shard.parentShardId())
-                && !discoveredShardIds.contains(shard.parentShardId())) {
-            return initialPosition;
-        }
-        return TRIM_HORIZON;
-    }
-
     private DynamoDbStreamsShardSplit mapToSplit(
             Shard shard, DynamodbStreamsSourceConfigConstants.InitialPosition initialPosition) {
         StartingPosition startingPosition;
@@ -251,6 +294,11 @@ public Set<String> getKnownSplitIds() {
         return knownSplits.keySet();
     }
 
+    @VisibleForTesting
+    public Map<String, DynamoDbStreamsShardSplit> getKnownSplits() {
+        return knownSplits;
+    }
+
     /**
      * finishedSplits needs to be cleaned up. The logic applied to cleaning up finished splits is
      * that if any split has been finished reading, its parent has been finished reading and it is

flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/util/ShardUtils.java

@@ -78,4 +78,14 @@ public static boolean isShardOlderThanInconsistencyDetectionRetentionPeriod(Stri
                                 .plus(
                                         DDB_STREAMS_MAX_RETENTION_PERIOD_FOR_RESOLVING_INCONSISTENCIES));
     }
+
+    /**
+     * Returns true if the shard was created before the given timestamp.
+     *
+     * @param shardId
+     * @param startTimestamp
+     */
+    public static boolean isShardCreatedBeforeTimestamp(String shardId, Instant startTimestamp) {
+        return getShardCreationTime(shardId).isBefore(startTimestamp);
+    }
 }

flink-connector-aws/flink-connector-dynamodb/src/test/java/org/apache/flink/connector/dynamodb/source/enumerator/DynamoDbStreamsSourceEnumeratorStateSerializerTest.java

@@ -25,14 +25,14 @@
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.time.Instant;
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import static java.util.Arrays.copyOf;
-import static org.apache.flink.connector.dynamodb.source.util.TestUtil.generateShardId;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatExceptionOfType;
 
@@ -45,19 +45,29 @@ void testSerializeAndDeserializeEverythingSpecified() throws Exception {
                         IntStream.rangeClosed(0, 3),
                         IntStream.rangeClosed(4, 10),
                         IntStream.rangeClosed(11, 15));
+        Instant startTimestamp = Instant.now();
+        DynamoDbStreamsSourceEnumeratorState deserializedState =
+                getDynamoDbStreamsSourceEnumeratorState(splitsToStore, startTimestamp);
+
+        assertThat(deserializedState.getKnownSplits()).isEqualTo(splitsToStore);
+        assertThat(deserializedState.getStartTimestamp().toEpochMilli())
+                .isEqualTo(startTimestamp.toEpochMilli());
+    }
+
+    private static DynamoDbStreamsSourceEnumeratorState getDynamoDbStreamsSourceEnumeratorState(
+            List<DynamoDBStreamsShardSplitWithAssignmentStatus> splitsToStore,
+            Instant startTimestamp)
+            throws IOException {
         DynamoDbStreamsSourceEnumeratorState initialState =
-                new DynamoDbStreamsSourceEnumeratorState(splitsToStore);
+                new DynamoDbStreamsSourceEnumeratorState(splitsToStore, startTimestamp);
 
         DynamoDbStreamsShardSplitSerializer splitSerializer =
                 new DynamoDbStreamsShardSplitSerializer();
         DynamoDbStreamsSourceEnumeratorStateSerializer serializer =
                 new DynamoDbStreamsSourceEnumeratorStateSerializer(splitSerializer);
 
         byte[] serialized = serializer.serialize(initialState);
-        DynamoDbStreamsSourceEnumeratorState deserializedState =
-                serializer.deserialize(serializer.getVersion(), serialized);
-
-        assertThat(deserializedState).usingRecursiveComparison().isEqualTo(initialState);
+        return serializer.deserialize(serializer.getVersion(), serialized);
     }
 
     @Test
@@ -67,10 +77,10 @@ void testDeserializeWithWrongVersionStateSerializer() throws Exception {
                         IntStream.rangeClosed(0, 3),
                         IntStream.rangeClosed(4, 10),
                         IntStream.rangeClosed(11, 15));
-        String lastSeenShardId = generateShardId(15);
+        Instant startTimestamp = Instant.now();
 
         DynamoDbStreamsSourceEnumeratorState initialState =
-                new DynamoDbStreamsSourceEnumeratorState(splitsToStore);
+                new DynamoDbStreamsSourceEnumeratorState(splitsToStore, startTimestamp);
         DynamoDbStreamsShardSplitSerializer splitSerializer =
                 new DynamoDbStreamsShardSplitSerializer();
         DynamoDbStreamsSourceEnumeratorStateSerializer serializer =
@@ -98,8 +108,9 @@ void testDeserializeWithWrongVersionSplitSerializer() throws Exception {
                         IntStream.rangeClosed(0, 3),
                         IntStream.rangeClosed(4, 10),
                         IntStream.rangeClosed(11, 15));
+        Instant startTimestamp = Instant.now();
         DynamoDbStreamsSourceEnumeratorState initialState =
-                new DynamoDbStreamsSourceEnumeratorState(splitsToStore);
+                new DynamoDbStreamsSourceEnumeratorState(splitsToStore, startTimestamp);
 
         DynamoDbStreamsShardSplitSerializer splitSerializer =
                 new DynamoDbStreamsShardSplitSerializer();
@@ -127,8 +138,9 @@ void testSerializeWithTrailingBytes() throws Exception {
                         IntStream.rangeClosed(0, 3),
                         IntStream.rangeClosed(4, 10),
                         IntStream.rangeClosed(11, 15));
+        Instant startTimestamp = Instant.now();
         DynamoDbStreamsSourceEnumeratorState initialState =
-                new DynamoDbStreamsSourceEnumeratorState(splitsToStore);
+                new DynamoDbStreamsSourceEnumeratorState(splitsToStore, startTimestamp);
 
         DynamoDbStreamsShardSplitSerializer splitSerializer =
                 new DynamoDbStreamsShardSplitSerializer();