Repository: gchq/sleeper
Issue: Ingest batcher state store tests fail intermittently

A test failed to query pending records and assert the sort order. No records were returned at all, when 3 records were added to the store.

It's not clear why this would happen. We can update this issue when it happens and try to figure it out.

Here are several times this happened in DynamoDBIngestBatcherStoreIT$OrderFilesReturnedFromStore.shouldListPendingFilesInOrderRequestsReceivedOldestFirst:

https://github.com/gchq/sleeper/actions/runs/5344363020/jobs/9688654804
https://github.com/gchq/sleeper/runs/14467112289

https://github.com/gchq/sleeper/actions/runs/5389911405/jobs/9784528532
https://github.com/gchq/sleeper/runs/14589403331

Here it happened in DynamoDBIngestBatcherStoreIT$AssignFilesToJobs.shouldFailToAssignFileWhenAssignmentAlreadyExists:

https://github.com/gchq/sleeper/actions/runs/5654678408/job/15318234077
https://github.com/gchq/sleeper/runs/15318336515

Here it happened in DynamoDBIngestBatcherStoreIT$AssignFilesToJobs.shouldFailToAssignFilesWhenOneIsAlreadyAssigned, with just one of the files missing:

https://github.com/mh674563/sleeper/actions/runs/5586589682/jobs/10210939461
https://github.com/mh674563/sleeper/actions/runs/5586589706/attempts/1

Here it happened in DynamoDBIngestBatcherStoreIT$OrderFilesReturnedFromStore.shouldReportAllFilesInOrderRequestsReceivedMostRecentFirst:

https://github.com/gchq/sleeper/actions/runs/5760737290/attempts/1

Here are some transaction cancellation failures:
https://github.com/gchq/sleeper/actions/runs/5521500088/jobs/10069693700
https://github.com/gchq/sleeper/runs/14949826850
https://github.com/gchq/sleeper/actions/runs/5654361346/job/15317281795?pr=1043


Changed files:
java/ingest/ingest-batcher-job-creator/src/test/java/sleeper/ingest/batcher/job/creator/IngestBatcherJobCreatorLambdaIT.java
java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreTestBase.java
java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreIT.java
java/ingest/ingest-batcher-store/src/main/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStore.java

diff --git a/java/ingest/ingest-batcher-job-creator/src/test/java/sleeper/ingest/batcher/job/creator/IngestBatcherJobCreatorLambdaIT.java b/java/ingest/ingest-batcher-job-creator/src/test/java/sleeper/ingest/batcher/job/creator/IngestBatcherJobCreatorLambdaIT.java
index 2ea932ba2..82f66f3f8 100644
--- a/java/ingest/ingest-batcher-job-creator/src/test/java/sleeper/ingest/batcher/job/creator/IngestBatcherJobCreatorLambdaIT.java
+++ b/java/ingest/ingest-batcher-job-creator/src/test/java/sleeper/ingest/batcher/job/creator/IngestBatcherJobCreatorLambdaIT.java
@@ -73,8 +73,8 @@ public class IngestBatcherJobCreatorLambdaIT {
         properties.set(DEFAULT_INGEST_BATCHER_MIN_JOB_SIZE, "0");
     });
     private final TableProperties tableProperties = createTestTableProperties(instanceProperties, schemaWithKey("key"), s3);
-    private final IngestBatcherStore store = new DynamoDBIngestBatcherStore(dynamoDB, instanceProperties,
-            new TablePropertiesProvider(s3, instanceProperties));
+    private final IngestBatcherStore store = DynamoDBIngestBatcherStore.withConsistentReads(
+            dynamoDB, instanceProperties, new TablePropertiesProvider(s3, instanceProperties));
 
     private IngestBatcherJobCreatorLambda lambdaWithTimesAndJobIds(List<Instant> times, List<String> jobIds) {
         return new IngestBatcherJobCreatorLambda(
@@ -103,6 +103,7 @@ public class IngestBatcherJobCreatorLambdaIT {
                 .fileSizeBytes(1024)
                 .receivedTime(Instant.parse("2023-05-25T14:43:00Z"))
                 .build());
+        store.getAllFilesNewestFirst(); // Use consistent reads to ensure PutItem is completed
 
         // When
         lambdaWithTimesAndJobIds(
diff --git a/java/ingest/ingest-batcher-store/src/main/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStore.java b/java/ingest/ingest-batcher-store/src/main/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStore.java
index c85079569..ce96a7459 100644
--- a/java/ingest/ingest-batcher-store/src/main/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStore.java
+++ b/java/ingest/ingest-batcher-store/src/main/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStore.java
@@ -54,13 +54,29 @@ public class DynamoDBIngestBatcherStore implements IngestBatcherStore {
     private final AmazonDynamoDB dynamoDB;
     private final String requestsTableName;
     private final TablePropertiesProvider tablePropertiesProvider;
+    private final boolean consistentReads;
 
     public DynamoDBIngestBatcherStore(AmazonDynamoDB dynamoDB,
                                       InstanceProperties instanceProperties,
                                       TablePropertiesProvider tablePropertiesProvider) {
+        this(dynamoDB, instanceProperties, tablePropertiesProvider, false);
+    }
+
+    private DynamoDBIngestBatcherStore(AmazonDynamoDB dynamoDB,
+                                       InstanceProperties instanceProperties,
+                                       TablePropertiesProvider tablePropertiesProvider,
+                                       boolean consistentReads) {
         this.dynamoDB = dynamoDB;
         this.requestsTableName = ingestRequestsTableName(instanceProperties.get(ID));
         this.tablePropertiesProvider = tablePropertiesProvider;
+        this.consistentReads = consistentReads;
+    }
+
+    public static DynamoDBIngestBatcherStore withConsistentReads(
+            AmazonDynamoDB dynamoDB,
+            InstanceProperties instanceProperties,
+            TablePropertiesProvider tablePropertiesProvider) {
+        return new DynamoDBIngestBatcherStore(dynamoDB, instanceProperties, tablePropertiesProvider, true);
     }
 
     public static String ingestRequestsTableName(String instanceId) {
@@ -97,7 +113,8 @@ public class DynamoDBIngestBatcherStore implements IngestBatcherStore {
     @Override
     public List<FileIngestRequest> getAllFilesNewestFirst() {
         return streamPagedItems(dynamoDB, new ScanRequest()
-                .withTableName(requestsTableName))
+                .withTableName(requestsTableName)
+                .withConsistentRead(consistentReads))
                 .map(DynamoDBIngestRequestFormat::readRecord)
                 .sorted(comparing(FileIngestRequest::getReceivedTime).reversed())
                 .collect(Collectors.toList());
@@ -107,6 +124,7 @@ public class DynamoDBIngestBatcherStore implements IngestBatcherStore {
     public List<FileIngestRequest> getPendingFilesOldestFirst() {
         return streamPagedItems(dynamoDB, new QueryRequest()
                 .withTableName(requestsTableName)
+                .withConsistentRead(consistentReads)
                 .withKeyConditionExpression("#JobId = :not_assigned")
                 .withExpressionAttributeNames(Map.of("#JobId", JOB_ID))
                 .withExpressionAttributeValues(new DynamoDBRecordBuilder()
diff --git a/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreIT.java b/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreIT.java
index 0eab8ac6c..c619fc269 100644
--- a/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreIT.java
+++ b/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreIT.java
@@ -118,6 +118,13 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
             // When
             store.addFile(fileIngestRequest1);
             store.addFile(fileIngestRequest2);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job", List.of(fileIngestRequest1, fileIngestRequest2));
 
             // Then
@@ -137,6 +144,13 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
 
             // When
             store.addFile(fileIngestRequest1);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job", List.of(fileIngestRequest1));
             store.addFile(fileIngestRequest2);
 
@@ -159,6 +173,13 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
 
             // When
             store.addFile(fileIngestRequest);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job", List.of(fileIngestRequest));
 
             // Then
@@ -178,6 +199,13 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
             FileIngestRequest fileIngestRequest = fileRequest()
                     .file("test-bucket/test.parquet").build();
             store.addFile(fileIngestRequest);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job-1", List.of(fileIngestRequest));
 
             // When / Then
@@ -197,6 +225,13 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
                     .file("test-bucket/test-2.parquet").build();
             store.addFile(fileIngestRequest1);
             store.addFile(fileIngestRequest2);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job-1", List.of(fileIngestRequest1));
 
             // When / Then
@@ -227,9 +262,22 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
             FileIngestRequest fileIngestRequest = fileRequest()
                     .file("test-bucket/sendTwice.parquet").build();
             store.addFile(fileIngestRequest);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("duplicate-job", List.of(fileIngestRequest));
             store.addFile(fileIngestRequest);
 
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             // When / Then
             List<FileIngestRequest> duplicateJob = List.of(fileIngestRequest);
             assertThatThrownBy(() -> store.assignJob("duplicate-job", duplicateJob))
@@ -300,7 +348,8 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
             store.addFile(fileIngestRequest);
 
             // Then
-            assertThat(streamPagedItems(dynamoDBClient, new ScanRequest().withTableName(requestsTableName)))
+            assertThat(streamPagedItems(dynamoDBClient, new ScanRequest()
+                    .withTableName(requestsTableName).withConsistentRead(true)))
                     .extracting(item -> getLongAttribute(item, EXPIRY_TIME, 0L))
                     .containsExactly(expectedExpiryTime.getEpochSecond());
         }
@@ -316,10 +365,18 @@ public class DynamoDBIngestBatcherStoreIT extends DynamoDBIngestBatcherStoreTest
 
             // When
             store.addFile(fileIngestRequest);
+
+            // Scan with consistent reads waits for PutItems to complete.
+            // TransactWriteItems will fail if a PutItem has not yet been committed.
+            // In production, a scan will happen before the update, so only files that exist in the store will be
+            // assigned to jobs.
+            store.getAllFilesNewestFirst();
+
             store.assignJob("test-job", List.of(fileIngestRequest));
 
             // Then
-            assertThat(streamPagedItems(dynamoDBClient, new ScanRequest().withTableName(requestsTableName)))
+            assertThat(streamPagedItems(dynamoDBClient, new ScanRequest()
+                    .withTableName(requestsTableName).withConsistentRead(true)))
                     .extracting(item -> getLongAttribute(item, EXPIRY_TIME, 0L))
                     .containsExactly(expectedExpiryTime.getEpochSecond());
         }
diff --git a/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreTestBase.java b/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreTestBase.java
index 9386f9987..a61193cbd 100644
--- a/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreTestBase.java
+++ b/java/ingest/ingest-batcher-store/src/test/java/sleeper/ingest/batcher/store/DynamoDBIngestBatcherStoreTestBase.java
@@ -43,7 +43,7 @@ public class DynamoDBIngestBatcherStoreTestBase extends DynamoDBTestBase {
     private final TablePropertiesProvider tablePropertiesProvider = new FixedTablePropertiesProvider(
             List.of(table1, table2));
     protected final String requestsTableName = DynamoDBIngestBatcherStore.ingestRequestsTableName(instanceProperties.get(ID));
-    protected final IngestBatcherStore store = new DynamoDBIngestBatcherStore(
+    protected final IngestBatcherStore store = DynamoDBIngestBatcherStore.withConsistentReads(
             dynamoDBClient, instanceProperties, tablePropertiesProvider);
 
     @BeforeEach