Issue 3070 - Wait for state store commits to finish based on logs

java/cdk/src/main/java/sleeper/cdk/stack/StateStoreCommitterStack.java

@@ -21,6 +21,7 @@
 import software.amazon.awscdk.services.iam.IGrantable;
 import software.amazon.awscdk.services.lambda.IFunction;
 import software.amazon.awscdk.services.lambda.eventsources.SqsEventSource;
+import software.amazon.awscdk.services.logs.LogGroup;
 import software.amazon.awscdk.services.s3.Bucket;
 import software.amazon.awscdk.services.s3.IBucket;
 import software.amazon.awscdk.services.sns.Topic;
@@ -41,6 +42,7 @@
 import static sleeper.cdk.Utils.createLambdaLogGroup;
 import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.STATESTORE_COMMITTER_DLQ_ARN;
 import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.STATESTORE_COMMITTER_DLQ_URL;
+import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.STATESTORE_COMMITTER_LOG_GROUP;
 import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.STATESTORE_COMMITTER_QUEUE_ARN;
 import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.STATESTORE_COMMITTER_QUEUE_URL;
 import static sleeper.configuration.properties.instance.CommonProperty.STATESTORE_COMMITTER_BATCH_SIZE;
@@ -116,6 +118,8 @@ private void lambdaToCommitStateStoreUpdates(
 
         String functionName = String.join("-", "sleeper",
                 Utils.cleanInstanceId(instanceProperties), "statestore-committer");
+        LogGroup logGroup = createLambdaLogGroup(this, "StateStoreCommitterLogGroup", functionName, instanceProperties);
+        instanceProperties.set(STATESTORE_COMMITTER_LOG_GROUP, logGroup.getLogGroupName());
 
         IFunction handlerFunction = committerJar.buildFunction(this, "StateStoreCommitter", builder -> builder
                 .functionName(functionName)
@@ -125,7 +129,7 @@ private void lambdaToCommitStateStoreUpdates(
                 .timeout(Duration.seconds(instanceProperties.getInt(STATESTORE_COMMITTER_LAMBDA_TIMEOUT_IN_SECONDS)))
                 .handler("sleeper.statestore.committer.lambda.StateStoreCommitterLambda::handleRequest")
                 .environment(environmentVariables)
-                .logGroup(createLambdaLogGroup(this, "StateStoreCommitterLogGroup", functionName, instanceProperties)));
+                .logGroup(logGroup));
 
         handlerFunction.addEventSource(SqsEventSource.Builder.create(commitQueue)
                 .batchSize(instanceProperties.getInt(STATESTORE_COMMITTER_BATCH_SIZE))

java/configuration/src/main/java/sleeper/configuration/properties/instance/CdkDefinedInstanceProperty.java

@@ -179,6 +179,10 @@ public interface CdkDefinedInstanceProperty extends InstanceProperty {
             .description("The ARN of the dead letter queue for statestore commit requests.")
             .propertyGroup(InstancePropertyGroup.COMMON)
             .build();
+    CdkDefinedInstanceProperty STATESTORE_COMMITTER_LOG_GROUP = Index.propertyBuilder("sleeper.statestore.committer.log.group")
+            .description("The nane of the log group for the state store committer.")
+            .propertyGroup(InstancePropertyGroup.COMMON)
+            .build();
 
     // Table metrics
     CdkDefinedInstanceProperty TABLE_METRICS_LAMBDA_FUNCTION = Index.propertyBuilder("sleeper.table.metrics.lambda.function")

java/statestore-commit/src/main/java/sleeper/commit/StateStoreCommitRequest.java

@@ -17,8 +17,8 @@
 
 import sleeper.compaction.job.commit.CompactionJobCommitRequest;
 import sleeper.compaction.job.commit.CompactionJobIdAssignmentCommitRequest;
+import sleeper.core.statestore.StateStoreException;
 import sleeper.core.statestore.commit.SplitPartitionCommitRequest;
-import sleeper.core.statestore.commit.StateStoreCommitRequestInS3;
 import sleeper.ingest.job.commit.IngestAddFilesCommitRequest;
 
 import java.util.Objects;
@@ -29,6 +29,8 @@
 public class StateStoreCommitRequest {
 
     private final Object request;
+    private final String tableId;
+    private final ApplyRequest applyRequest;
 
     /**
      * Creates a request to commit the results of a compaction job.
@@ -37,7 +39,7 @@ public class StateStoreCommitRequest {
      * @return         a state store commit request
      */
     public static StateStoreCommitRequest forCompactionJob(CompactionJobCommitRequest request) {
-        return new StateStoreCommitRequest(request);
+        return new StateStoreCommitRequest(request, request.getJob().getTableId(), committer -> committer.commitCompaction(request));
     }
 
     /**
@@ -47,7 +49,7 @@ public static StateStoreCommitRequest forCompactionJob(CompactionJobCommitReques
      * @return         a state store commit request
      */
     public static StateStoreCommitRequest forCompactionJobIdAssignment(CompactionJobIdAssignmentCommitRequest request) {
-        return new StateStoreCommitRequest(request);
+        return new StateStoreCommitRequest(request, request.getTableId(), committer -> committer.assignCompactionInputFiles(request));
     }
 
     /**
@@ -57,7 +59,7 @@ public static StateStoreCommitRequest forCompactionJobIdAssignment(CompactionJob
      * @return         a state store commit request
      */
     public static StateStoreCommitRequest forIngestAddFiles(IngestAddFilesCommitRequest request) {
-        return new StateStoreCommitRequest(request);
+        return new StateStoreCommitRequest(request, request.getTableId(), committer -> committer.addFiles(request));
     }
 
     /**
@@ -67,27 +69,27 @@ public static StateStoreCommitRequest forIngestAddFiles(IngestAddFilesCommitRequ
      * @return         a state store commit request
      */
     public static StateStoreCommitRequest forSplitPartition(SplitPartitionCommitRequest request) {
-        return new StateStoreCommitRequest(request);
+        return new StateStoreCommitRequest(request, request.getTableId(), committer -> committer.splitPartition(request));
     }
 
-    /**
-     * Creates a request which is stored in S3.
-     *
-     * @param  request the commit request
-     * @return         a state store commit request
-     */
-    public static StateStoreCommitRequest storedInS3(StateStoreCommitRequestInS3 request) {
-        return new StateStoreCommitRequest(request);
-    }
-
-    private StateStoreCommitRequest(Object request) {
+    private StateStoreCommitRequest(Object request, String tableId, ApplyRequest applyRequest) {
         this.request = request;
+        this.tableId = tableId;
+        this.applyRequest = applyRequest;
     }
 
     public Object getRequest() {
         return request;
     }
 
+    public String getTableId() {
+        return tableId;
+    }
+
+    void apply(StateStoreCommitter committer) throws StateStoreException {
+        applyRequest.apply(committer);
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(request);
@@ -110,4 +112,11 @@ public String toString() {
         return "StateStoreCommitRequest{request=" + request + "}";
     }
 
+    /**
+     * Applies the current request with a given committer.
+     */
+    @FunctionalInterface
+    private interface ApplyRequest {
+        void apply(StateStoreCommitter committer) throws StateStoreException;
+    }
 }

java/statestore-commit/src/main/java/sleeper/commit/StateStoreCommitRequestDeserialiser.java

@@ -24,6 +24,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import sleeper.commit.StateStoreCommitter.LoadS3ObjectFromDataBucket;
 import sleeper.compaction.job.CompactionJob;
 import sleeper.compaction.job.CompactionJobJsonSerDe;
 import sleeper.compaction.job.commit.CompactionJobCommitRequest;
@@ -46,11 +47,20 @@ public class StateStoreCommitRequestDeserialiser {
     public static final Logger LOGGER = LoggerFactory.getLogger(StateStoreCommitRequestDeserialiser.class);
 
     private final Gson gson;
+    private final Gson gsonFromDataBucket;
+    private final LoadS3ObjectFromDataBucket loadFromDataBucket;
+
+    public StateStoreCommitRequestDeserialiser(
+            TablePropertiesProvider tablePropertiesProvider, LoadS3ObjectFromDataBucket loadFromDataBucket) {
+        gson = gson(tablePropertiesProvider, this::fromDataBucket);
+        gsonFromDataBucket = gson(tablePropertiesProvider, DeserialiseFromDataBucket.refuseWhileReadingFromBucket());
+        this.loadFromDataBucket = loadFromDataBucket;
+    }
 
-    public StateStoreCommitRequestDeserialiser(TablePropertiesProvider tablePropertiesProvider) {
-        gson = GsonConfig.standardBuilder()
+    private static Gson gson(TablePropertiesProvider tablePropertiesProvider, DeserialiseFromDataBucket readFromDataBucket) {
+        return GsonConfig.standardBuilder()
                 .registerTypeAdapter(CompactionJob.class, new CompactionJobJsonSerDe())
-                .registerTypeAdapter(StateStoreCommitRequest.class, new WrapperDeserialiser())
+                .registerTypeAdapter(StateStoreCommitRequest.class, new WrapperDeserialiser(readFromDataBucket))
                 .registerTypeAdapter(SplitPartitionCommitRequest.class, new SplitPartitionDeserialiser(tablePropertiesProvider))
                 .serializeNulls()
                 .create();
@@ -66,11 +76,22 @@ public StateStoreCommitRequest fromJson(String jsonString) {
         return gson.fromJson(jsonString, StateStoreCommitRequest.class);
     }
 
+    private StateStoreCommitRequest fromDataBucket(StateStoreCommitRequestInS3 request) {
+        String json = loadFromDataBucket.loadFromDataBucket(request.getKeyInS3());
+        return gsonFromDataBucket.fromJson(json, StateStoreCommitRequest.class);
+    }
+
     /**
      * Deserialises the commit request by reading the type.
      */
     private static class WrapperDeserialiser implements JsonDeserializer<StateStoreCommitRequest> {
 
+        private final DeserialiseFromDataBucket fromDataBucket;
+
+        private WrapperDeserialiser(DeserialiseFromDataBucket fromDataBucket) {
+            this.fromDataBucket = fromDataBucket;
+        }
+
         @Override
         public StateStoreCommitRequest deserialize(JsonElement json, Type wrapperType, JsonDeserializationContext context) throws JsonParseException {
 
@@ -88,15 +109,15 @@ public StateStoreCommitRequest deserialize(JsonElement json, Type wrapperType, J
                 case INGEST_ADD_FILES:
                     return StateStoreCommitRequest.forIngestAddFiles(
                             context.deserialize(requestObj, IngestAddFilesCommitRequest.class));
-                case STORED_IN_S3:
-                    return StateStoreCommitRequest.storedInS3(
-                            context.deserialize(requestObj, StateStoreCommitRequestInS3.class));
                 case COMPACTION_JOB_ID_ASSIGNMENT:
                     return StateStoreCommitRequest.forCompactionJobIdAssignment(
                             context.deserialize(requestObj, CompactionJobIdAssignmentCommitRequest.class));
                 case SPLIT_PARTITION:
                     return StateStoreCommitRequest.forSplitPartition(
                             context.deserialize(requestObj, SplitPartitionCommitRequest.class));
+                case STORED_IN_S3:
+                    return fromDataBucket.read(
+                            context.deserialize(requestObj, StateStoreCommitRequestInS3.class));
                 default:
                     throw new CommitRequestValidationException("Unrecognised request type");
             }
@@ -138,4 +159,19 @@ public SplitPartitionCommitRequest deserialize(JsonElement jsonElement, Type typ
             return new SplitPartitionCommitRequest(tableId, parentPartition, leftChildPartition, rightChildPartition);
         }
     }
+
+    /**
+     * Reads and deserialises a commit request from the data bucket. An alternative implementation will refuse reading
+     * from the bucket because the pointer was already stored in S3.
+     */
+    @FunctionalInterface
+    private interface DeserialiseFromDataBucket {
+        StateStoreCommitRequest read(StateStoreCommitRequestInS3 request);
+
+        static DeserialiseFromDataBucket refuseWhileReadingFromBucket() {
+            return request -> {
+                throw new IllegalArgumentException("Found a request stored in S3 pointing to another S3 object: " + request.getKeyInS3());
+            };
+        }
+    }
 }

java/statestore-commit/src/main/java/sleeper/commit/StateStoreCommitter.java

@@ -30,7 +30,6 @@
 import sleeper.core.statestore.StateStore;
 import sleeper.core.statestore.StateStoreException;
 import sleeper.core.statestore.commit.SplitPartitionCommitRequest;
-import sleeper.core.statestore.commit.StateStoreCommitRequestInS3;
 import sleeper.core.statestore.exception.FileAlreadyExistsException;
 import sleeper.core.statestore.exception.FileNotFoundException;
 import sleeper.core.statestore.exception.FileReferenceNotAssignedToJobException;
@@ -59,7 +58,6 @@ public class StateStoreCommitter {
     private final CompactionJobStatusStore compactionJobStatusStore;
     private final IngestJobStatusStore ingestJobStatusStore;
     private final GetStateStoreByTableId stateStoreProvider;
-    private final LoadS3ObjectFromDataBucket loadFromDataBucket;
     private final Supplier<Instant> timeSupplier;
 
     public StateStoreCommitter(
@@ -69,21 +67,23 @@ public StateStoreCommitter(
             GetStateStoreByTableId stateStoreProvider,
             LoadS3ObjectFromDataBucket loadFromDataBucket,
             Supplier<Instant> timeSupplier) {
-        this.deserialiser = new StateStoreCommitRequestDeserialiser(tablePropertiesProvider);
+        this.deserialiser = new StateStoreCommitRequestDeserialiser(tablePropertiesProvider, loadFromDataBucket);
         this.compactionJobStatusStore = compactionJobStatusStore;
         this.ingestJobStatusStore = ingestJobStatusStore;
         this.stateStoreProvider = stateStoreProvider;
-        this.loadFromDataBucket = loadFromDataBucket;
         this.timeSupplier = timeSupplier;
     }
 
     /**
      * Applies a state store commit request.
      *
-     * @param json the commit request JSON string
+     * @param  json the commit request JSON string
+     * @return      the commit request
      */
-    public void applyFromJson(String json) throws StateStoreException {
-        apply(deserialiser.fromJson(json));
+    public StateStoreCommitRequest applyFromJson(String json) throws StateStoreException {
+        StateStoreCommitRequest request = deserialiser.fromJson(json);
+        apply(request);
+        return request;
     }
 
     /**
@@ -92,30 +92,19 @@ public void applyFromJson(String json) throws StateStoreException {
      * @param request the commit request
      */
     public void apply(StateStoreCommitRequest request) throws StateStoreException {
-        Object requestObj = request.getRequest();
-        if (requestObj instanceof CompactionJobCommitRequest) {
-            apply((CompactionJobCommitRequest) requestObj);
-        } else if (requestObj instanceof IngestAddFilesCommitRequest) {
-            apply((IngestAddFilesCommitRequest) requestObj);
-        } else if (requestObj instanceof StateStoreCommitRequestInS3) {
-            apply((StateStoreCommitRequestInS3) requestObj);
-        } else if (requestObj instanceof CompactionJobIdAssignmentCommitRequest) {
-            apply((CompactionJobIdAssignmentCommitRequest) requestObj);
-        } else if (requestObj instanceof SplitPartitionCommitRequest) {
-            apply((SplitPartitionCommitRequest) requestObj);
-        } else {
-            throw new IllegalArgumentException("Unsupported commit request type: " + requestObj.getClass().getName());
-        }
+        request.apply(this);
+        LOGGER.info("Applied request to table ID {} with type {} at time {}",
+                request.getTableId(), request.getRequest().getClass().getSimpleName(), Instant.now());
     }
 
-    private void apply(CompactionJobCommitRequest request) throws StateStoreException {
+    void commitCompaction(CompactionJobCommitRequest request) throws StateStoreException {
         CompactionJob job = request.getJob();
         try {
             CompactionJobCommitter.updateStateStoreSuccess(job, request.getRecordsWritten(),
                     stateStoreProvider.getByTableId(job.getTableId()));
             compactionJobStatusStore.jobCommitted(compactionJobCommitted(job, timeSupplier.get())
                     .taskId(request.getTaskId()).jobRunId(request.getJobRunId()).build());
-            LOGGER.info("Successfully committed compaction job {} to table with ID {}", job.getId(), job.getTableId());
+            LOGGER.debug("Successfully committed compaction job {}", job.getId());
         } catch (ReplaceRequestsFailedException e) {
             Exception failure = e.getFailures().get(0);
             if (failure instanceof FileNotFoundException
@@ -133,32 +122,28 @@ private void apply(CompactionJobCommitRequest request) throws StateStoreExceptio
         }
     }
 
-    private void apply(IngestAddFilesCommitRequest request) throws StateStoreException {
+    void addFiles(IngestAddFilesCommitRequest request) throws StateStoreException {
         StateStore stateStore = stateStoreProvider.getByTableId(request.getTableId());
         List<AllReferencesToAFile> files = AllReferencesToAFile.newFilesWithReferences(request.getFileReferences());
         stateStore.addFilesWithReferences(files);
         IngestJob job = request.getJob();
         if (job != null) {
             ingestJobStatusStore.jobAddedFiles(IngestJobAddedFilesEvent.ingestJobAddedFiles(job, files, request.getWrittenTime())
                     .taskId(request.getTaskId()).jobRunId(request.getJobRunId()).build());
-            LOGGER.info("Successfully committed new files for ingest job {} to table with ID {}", job.getId(), request.getTableId());
+            LOGGER.debug("Successfully committed new files for ingest job {}", job.getId());
         } else {
-            LOGGER.info("Successfully committed new files for ingest to table with ID {}", request.getTableId());
+            LOGGER.debug("Successfully committed new files for ingest with no job");
         }
     }
 
-    private void apply(SplitPartitionCommitRequest request) throws StateStoreException {
+    void splitPartition(SplitPartitionCommitRequest request) throws StateStoreException {
         StateStore stateStore = stateStoreProvider.getByTableId(request.getTableId());
         stateStore.atomicallyUpdatePartitionAndCreateNewOnes(request.getParentPartition(), request.getLeftChild(), request.getRightChild());
     }
 
-    private void apply(StateStoreCommitRequestInS3 request) throws StateStoreException {
-        String json = loadFromDataBucket.loadFromDataBucket(request.getKeyInS3());
-        StateStoreCommitRequest requestFromS3 = deserialiser.fromJson(json);
-        if (requestFromS3.getRequest() instanceof StateStoreCommitRequestInS3) {
-            throw new IllegalArgumentException("Found a request stored in S3 pointing to another S3 object: " + request.getKeyInS3());
-        }
-        apply(requestFromS3);
+    void assignCompactionInputFiles(CompactionJobIdAssignmentCommitRequest request) throws StateStoreException {
+        StateStore stateStore = stateStoreProvider.getByTableId(request.getTableId());
+        stateStore.assignJobIds(request.getAssignJobIdRequests());
     }
 
     /**
@@ -173,9 +158,4 @@ public interface LoadS3ObjectFromDataBucket {
          */
         String loadFromDataBucket(String key);
     }
-
-    private void apply(CompactionJobIdAssignmentCommitRequest request) throws StateStoreException {
-        StateStore stateStore = stateStoreProvider.getByTableId(request.getTableId());
-        stateStore.assignJobIds(request.getAssignJobIdRequests());
-    }
 }