Issue 3799 - System test creating many compaction jobs at once

.vscode/settings.json

@@ -24,12 +24,14 @@
     "org.mockito.Mockito.*",
     "org.mockito.ArgumentMatchers.*",
     "org.mockito.Answers.*",
+    "sleeper.core.partition.PartitionTreeTestHelper.*",
     "sleeper.core.schema.SchemaTestHelper.*",
     "sleeper.core.statestore.inmemory.StateStoreTestHelper.*",
     "sleeper.configuration.properties.InstancePropertiesTestHelper.*",
     "sleeper.configuration.properties.table.TablePropertiesTestHelper.*",
     "sleeper.core.util.ExponentialBackoffWithJitterTestHelper.*",
-    "sleeper.core.testutils.SupplierTestHelper.*"
+    "sleeper.core.testutils.SupplierTestHelper.*",
+    "sleeper.systemtest.dsl.testutil.SystemTestPartitionsTestHelper.*"
   ],
   "comment.java.configuration.runtimes": [
     {

java/compaction/compaction-rust/src/main/java/sleeper/compaction/rust/RustBridge.java

@@ -74,7 +74,7 @@ public static synchronized Compaction getRustCompactor() throws IOException {
             return nativeLib;
 
         } catch (UnsatisfiedLinkError err) {
-            throw (IOException) new IOException().initCause(err);
+            throw (IOException) new IOException("Could not initialise Rust bridge", err);
         }
     }
 

java/core/src/main/java/sleeper/core/partition/PartitionsFromSplitPoints.java

@@ -95,7 +95,7 @@ public List<Partition> construct() {
         List<Partition> partitionsInLayer = List.of(tree.getRootPartition());
         while (!partitionsInLayer.isEmpty()) {
             LOGGER.debug("Layer {}", layer++);
-            partitionsInLayer.forEach(partition -> LOGGER.debug(partition.toString()));
+            partitionsInLayer.forEach(partition -> LOGGER.debug("{}", partition));
             partitionsInLayer = partitionsInLayer.stream()
                     .map(partition -> partition.getChildPartitionIds())
                     .flatMap(List::stream)
@@ -136,7 +136,6 @@ private List<Partition.Builder> addLayer(List<Partition.Builder> partitionsInLay
                 leftPartition.parentPartitionId(id);
                 rightPartition.parentPartitionId(id);
                 parents.add(parent);
-                LOGGER.debug("Created parent partition {} joining partitions {}", id, childPartitionIds);
             }
         }
         allPartitions.addAll(parents);
@@ -165,7 +164,6 @@ private List<Partition.Builder> createLeafPartitions() {
                     .childPartitionIds(new ArrayList<>())
                     .dimension(-1);
             leafPartitions.add(partition);
-            LOGGER.debug("Created leaf partition {} for region {}", id, region);
         }
         LOGGER.info("Created {} leaf partitions from {} split points", leafPartitions.size(), splitPoints.size());
         return leafPartitions;

java/core/src/test/java/sleeper/core/partition/PartitionTreeTestHelper.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.core.partition;
+
+import sleeper.core.record.Record;
+import sleeper.core.schema.Field;
+import sleeper.core.schema.Schema;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A test helper for creating partitions based on split points.
+ */
+public class PartitionTreeTestHelper {
+
+    private PartitionTreeTestHelper() {
+    }
+
+    /**
+     * Creates a partition tree by finding records at the boundaries. This can be used when we are generating predefined
+     * test data, and we can generate the nth record in sort order at will. This will create a leaf partition for each
+     * group of records where we specify the number of records in a group. The boundaries of each partition will be the
+     * value of the split field for the record at the boundary. This will take the first row key as the field to split
+     * on. This may be used to avoid needing to hold all the test data in memory at once.
+     *
+     * @param  recordsPerPartition the number of records to fit in each partition
+     * @param  totalRecords        the total number of records that will be in the table
+     * @param  generator           a method to find the nth record in sort order
+     * @param  schema              the schema
+     * @return                     the partition tree
+     */
+    public static PartitionTree createPartitionTreeWithRecordsPerPartitionAndTotal(int recordsPerPartition, long totalRecords, NthRecordGenerator generator, Schema schema) {
+        List<Object> splitPoints = new ArrayList<>();
+        Field splitField = schema.getRowKeyFields().get(0);
+        for (long i = recordsPerPartition; i < totalRecords; i += recordsPerPartition) {
+            splitPoints.add(generator.getNthRecord(i).get(splitField.getName()));
+        }
+        return PartitionsFromSplitPoints.treeFrom(schema, splitPoints);
+    }
+
+    /**
+     * A generator for the nth record in sort order to generate test data on the fly. Used when the test data is defined
+     * deterministically to avoid needing to hold it all in memory at once.
+     */
+    @FunctionalInterface
+    public interface NthRecordGenerator {
+
+        /**
+         * Gets the nth record in the test data in sort order.
+         *
+         * @param  n the index of the record, starting from 0
+         * @return   the record
+         */
+        Record getNthRecord(long n);
+    }
+}

java/core/src/test/java/sleeper/core/partition/PartitionTreeTestHelperTest.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.core.partition;
+
+import org.junit.jupiter.api.Test;
+
+import sleeper.core.record.Record;
+import sleeper.core.schema.Schema;
+import sleeper.core.schema.type.IntType;
+import sleeper.core.testutils.printers.PartitionsPrinter;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static sleeper.core.partition.PartitionTreeTestHelper.createPartitionTreeWithRecordsPerPartitionAndTotal;
+import static sleeper.core.schema.SchemaTestHelper.schemaWithKey;
+
+public class PartitionTreeTestHelperTest {
+
+    @Test
+    void shouldCreateSplitPointsFromRecordRangeAndRecordsPerPartition() {
+        // Given
+        Schema schema = schemaWithKey("key", new IntType());
+        List<Record> records = List.of(
+                new Record(Map.of("key", 10)),
+                new Record(Map.of("key", 20)),
+                new Record(Map.of("key", 30)));
+
+        // When
+        PartitionTree tree = createPartitionTreeWithRecordsPerPartition(1, records, schema);
+
+        // Then
+        assertThat(PartitionsPrinter.printPartitions(schema, tree))
+                .isEqualTo("""
+                        Leaf partition at LL:
+                        {"key":{"min":-2147483648,"minInclusive":true,"max":20,"maxInclusive":false},"stringsBase64Encoded":true}
+                        Leaf partition at LR:
+                        {"key":{"min":20,"minInclusive":true,"max":30,"maxInclusive":false},"stringsBase64Encoded":true}
+                        Leaf partition at R:
+                        {"key":{"min":30,"minInclusive":true,"max":null,"maxInclusive":false},"stringsBase64Encoded":true}
+                        Partition at L:
+                        {"key":{"min":-2147483648,"minInclusive":true,"max":30,"maxInclusive":false},"stringsBase64Encoded":true}
+                        Partition at root:
+                        {"key":{"min":-2147483648,"minInclusive":true,"max":null,"maxInclusive":false},"stringsBase64Encoded":true}
+                        """);
+    }
+
+    private PartitionTree createPartitionTreeWithRecordsPerPartition(int recordsPerPartition, List<Record> records, Schema schema) {
+        return createPartitionTreeWithRecordsPerPartitionAndTotal(recordsPerPartition, records.size(),
+                index -> records.get((int) index),
+                schema);
+    }
+}

java/system-test/system-test-drivers/pom.xml

@@ -90,11 +90,6 @@
             <artifactId>testcontainers</artifactId>
             <scope>test</scope>
         </dependency>
-        <dependency>
-            <groupId>org.testcontainers</groupId>
-            <artifactId>junit-jupiter</artifactId>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>org.testcontainers</groupId>
             <artifactId>localstack</artifactId>

java/system-test/system-test-drivers/src/main/java/sleeper/systemtest/drivers/compaction/AwsCompactionDriver.java

@@ -25,8 +25,15 @@
 import software.amazon.awssdk.services.ec2.Ec2Client;
 import software.amazon.awssdk.services.ecs.EcsClient;
 import software.amazon.awssdk.services.lambda.LambdaClient;
+import software.amazon.awssdk.services.sqs.SqsClient;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
+import software.amazon.awssdk.services.sqs.model.Message;
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
 
 import sleeper.clients.deploy.InvokeLambda;
+import sleeper.compaction.core.job.CompactionJob;
+import sleeper.compaction.core.job.CompactionJobSerDe;
 import sleeper.compaction.core.job.CompactionJobStatusStore;
 import sleeper.compaction.core.job.creation.CreateCompactionJobs;
 import sleeper.compaction.core.task.CompactionTaskStatus;
@@ -45,8 +52,12 @@
 import sleeper.task.common.RunCompactionTasks;
 
 import java.io.IOException;
+import java.util.List;
+import java.util.stream.Stream;
 
+import static java.util.function.Predicate.not;
 import static sleeper.core.properties.instance.CdkDefinedInstanceProperty.COMPACTION_JOB_CREATION_TRIGGER_LAMBDA_FUNCTION;
+import static sleeper.core.properties.instance.CdkDefinedInstanceProperty.COMPACTION_JOB_QUEUE_URL;
 import static sleeper.core.properties.instance.CdkDefinedInstanceProperty.COMPACTION_TASK_CREATION_LAMBDA_FUNCTION;
 
 public class AwsCompactionDriver implements CompactionDriver {
@@ -57,16 +68,19 @@ public class AwsCompactionDriver implements CompactionDriver {
     private final AmazonDynamoDB dynamoDBClient;
     private final AmazonS3 s3Client;
     private final AmazonSQS sqsClient;
+    private final SqsClient sqsClientV2;
     private final EcsClient ecsClient;
     private final AutoScalingClient asClient;
     private final Ec2Client ec2Client;
+    private final CompactionJobSerDe serDe = new CompactionJobSerDe();
 
     public AwsCompactionDriver(SystemTestInstanceContext instance, SystemTestClients clients) {
         this.instance = instance;
         this.lambdaClient = clients.getLambda();
         this.dynamoDBClient = clients.getDynamoDB();
         this.s3Client = clients.getS3();
         this.sqsClient = clients.getSqs();
+        this.sqsClientV2 = clients.getSqsV2();
         this.ecsClient = clients.getEcs();
         this.asClient = clients.getAutoScaling();
         this.ec2Client = clients.getEc2();
@@ -138,4 +152,41 @@ public void forceStartTasks(int numberOfTasks, PollWithRetries poll) {
     public void scaleToZero() {
         EC2Scaler.create(instance.getInstanceProperties(), asClient, ec2Client).scaleTo(0);
     }
+
+    @Override
+    public List<CompactionJob> drainJobsQueueForWholeInstance() {
+        String queueUrl = instance.getInstanceProperties().get(COMPACTION_JOB_QUEUE_URL);
+        LOGGER.info("Draining compaction jobs queue: {}", queueUrl);
+        List<CompactionJob> jobs = Stream.iterate(
+                receiveJobs(queueUrl), not(List::isEmpty), lastJobs -> receiveJobs(queueUrl))
+                .flatMap(List::stream).toList();
+        LOGGER.info("Found {} compaction jobs", jobs.size());
+        return jobs;
+    }
+
+    private List<CompactionJob> receiveJobs(String queueUrl) {
+        ReceiveMessageResponse receiveResult = sqsClientV2.receiveMessage(request -> request
+                .queueUrl(queueUrl)
+                .maxNumberOfMessages(10)
+                .waitTimeSeconds(5));
+        List<Message> messages = receiveResult.messages();
+        if (messages.isEmpty()) {
+            return List.of();
+        }
+        DeleteMessageBatchResponse deleteResult = sqsClientV2.deleteMessageBatch(request -> request
+                .queueUrl(queueUrl)
+                .entries(messages.stream()
+                        .map(message -> DeleteMessageBatchRequestEntry.builder()
+                                .id(message.messageId())
+                                .receiptHandle(message.receiptHandle())
+                                .build())
+                        .toList()));
+        if (!deleteResult.failed().isEmpty()) {
+            throw new RuntimeException("Failed deleting compaction job messages: " + deleteResult.failed());
+        }
+        return messages.stream()
+                .map(Message::body)
+                .map(serDe::fromJson)
+                .toList();
+    }
 }

java/system-test/system-test-drivers/src/test/java/sleeper/systemtest/drivers/compaction/AwsCompactionDriverDrainCompactionsIT.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.systemtest.drivers.compaction;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.services.sqs.SqsClient;
+import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
+
+import sleeper.compaction.core.job.CompactionJob;
+import sleeper.compaction.core.job.CompactionJobFactory;
+import sleeper.compaction.core.job.CompactionJobSerDe;
+import sleeper.core.statestore.FileReferenceFactory;
+import sleeper.core.util.SplitIntoBatches;
+import sleeper.systemtest.drivers.testutil.LocalStackDslTest;
+import sleeper.systemtest.drivers.testutil.LocalStackSystemTestDrivers;
+import sleeper.systemtest.dsl.SleeperSystemTest;
+import sleeper.systemtest.dsl.SystemTestContext;
+import sleeper.systemtest.dsl.compaction.CompactionDriver;
+import sleeper.systemtest.dsl.instance.SystemTestInstanceContext;
+
+import java.util.List;
+import java.util.stream.IntStream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static sleeper.core.properties.instance.CdkDefinedInstanceProperty.COMPACTION_JOB_QUEUE_URL;
+import static sleeper.systemtest.drivers.testutil.LocalStackTestInstance.DRAIN_COMPACTIONS;
+
+@LocalStackDslTest
+public class AwsCompactionDriverDrainCompactionsIT {
+    public static final Logger LOGGER = LoggerFactory.getLogger(AwsCompactionDriverDrainCompactionsIT.class);
+
+    SqsClient sqs;
+    CompactionDriver driver;
+    SystemTestInstanceContext instance;
+
+    @BeforeEach
+    void setUp(SleeperSystemTest sleeper, SystemTestContext context, LocalStackSystemTestDrivers drivers) {
+        sleeper.connectToInstance(DRAIN_COMPACTIONS);
+        sqs = drivers.clients().getSqsV2();
+        driver = drivers.compaction(context);
+        instance = context.instance();
+    }
+
+    @Test
+    void shouldDrainCompactionJobsFromQueue() {
+        // Given
+        FileReferenceFactory fileFactory = FileReferenceFactory.from(instance.getStateStore());
+        CompactionJobFactory jobFactory = new CompactionJobFactory(instance.getInstanceProperties(), instance.getTableProperties());
+        List<CompactionJob> jobs = IntStream.rangeClosed(1, 20)
+                .mapToObj(i -> jobFactory.createCompactionJob(
+                        List.of(fileFactory.rootFile("file" + i + ".parquet", i)), "root"))
+                .toList();
+        send(jobs);
+
+        // When / Then
+        assertThat(driver.drainJobsQueueForWholeInstance())
+                .containsExactlyElementsOf(jobs);
+    }
+
+    private void send(List<CompactionJob> jobs) {
+        String queueUrl = instance.getInstanceProperties().get(COMPACTION_JOB_QUEUE_URL);
+        LOGGER.info("Sending to compaction jobs queue: {}", queueUrl);
+        CompactionJobSerDe serDe = new CompactionJobSerDe();
+        for (List<CompactionJob> batch : SplitIntoBatches.splitListIntoBatchesOf(10, jobs)) {
+            sqs.sendMessageBatch(request -> request
+                    .queueUrl(queueUrl)
+                    .entries(batch.stream()
+                            .map(job -> SendMessageBatchRequestEntry.builder()
+                                    .messageBody(serDe.toJson(job))
+                                    .id(job.getId())
+                                    .build())
+                            .toList()));
+        }
+    }
+
+}