DBZ-8193 Batch write to AWS Kinesis

debezium-server-kinesis/src/main/java/io/debezium/server/kinesis/KinesisChangeConsumer.java

@@ -7,8 +7,11 @@
 
 import java.net.URI;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
@@ -38,7 +41,10 @@
 import software.amazon.awssdk.services.kinesis.KinesisClient;
 import software.amazon.awssdk.services.kinesis.KinesisClientBuilder;
 import software.amazon.awssdk.services.kinesis.model.KinesisException;
-import software.amazon.awssdk.services.kinesis.model.PutRecordRequest;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsRequest;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsRequestEntry;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsResponse;
+import software.amazon.awssdk.services.kinesis.model.PutRecordsResultEntry;
 
 /**
  * Implementation of the consumer that delivers the messages into Amazon Kinesis destination.
@@ -56,12 +62,15 @@ public class KinesisChangeConsumer extends BaseChangeConsumer implements Debeziu
     private static final String PROP_REGION_NAME = PROP_PREFIX + "region";
     private static final String PROP_ENDPOINT_NAME = PROP_PREFIX + "endpoint";
     private static final String PROP_CREDENTIALS_PROFILE = PROP_PREFIX + "credentials.profile";
+    private static final String PROP_BATCH_SIZE = PROP_PREFIX + "batch.size";
+    private static final String PROP_RETRIES = PROP_PREFIX + "default.retries";
 
     private String region;
     private Optional<String> endpointOverride;
     private Optional<String> credentialsProfile;
-    private static final int DEFAULT_RETRIES = 5;
     private static final Duration RETRY_INTERVAL = Duration.ofSeconds(1);
+    private Integer batchSize;
+    private Integer RETRIES;
 
     @ConfigProperty(name = PROP_PREFIX + "null.key", defaultValue = "default")
     String nullKey;
@@ -74,13 +83,23 @@ public class KinesisChangeConsumer extends BaseChangeConsumer implements Debeziu
 
     @PostConstruct
     void connect() {
+        final Config config = ConfigProvider.getConfig();
+        batchSize = config.getOptionalValue(PROP_BATCH_SIZE, Integer.class).orElse(500);
+        RETRIES = config.getOptionalValue(PROP_RETRIES, Integer.class).orElse(5);
+
+        if (batchSize <= 0) {
+            throw new DebeziumException("Batch size must be greater than 0");
+        }
+        else if (batchSize > 500) {
+            throw new DebeziumException("Retries must be less than or equal to 500");
+        }
+
         if (customClient.isResolvable()) {
             client = customClient.get();
             LOGGER.info("Obtained custom configured KinesisClient '{}'", client);
             return;
         }
 
-        final Config config = ConfigProvider.getConfig();
         region = config.getValue(PROP_REGION_NAME, String.class);
         endpointOverride = config.getOptionalValue(PROP_ENDPOINT_NAME, String.class);
         credentialsProfile = config.getOptionalValue(PROP_CREDENTIALS_PROFILE, String.class);
@@ -106,41 +125,102 @@ void close() {
     @Override
     public void handleBatch(List<ChangeEvent<Object, Object>> records, RecordCommitter<ChangeEvent<Object, Object>> committer)
             throws InterruptedException {
-        for (ChangeEvent<Object, Object> record : records) {
-            LOGGER.trace("Received event '{}'", record);
-
-            int attempts = 0;
-            while (!recordSent(record)) {
-                attempts++;
-                if (attempts >= DEFAULT_RETRIES) {
-                    throw new DebeziumException("Exceeded maximum number of attempts to publish event " + record);
+
+        // Guard if records are empty
+        if (records.isEmpty()) {
+            committer.markBatchFinished();
+            return;
+        }
+
+        String streamName;
+        List<ChangeEvent<Object, Object>> batch = new ArrayList<>();
+
+        // Group the records by destination
+        Map<String, List<ChangeEvent<Object, Object>>> segmentedBatches = records.stream().collect(Collectors.groupingBy(record -> record.destination()));
+
+        // Iterate over the segmentedBatches
+        for (List<ChangeEvent<Object, Object>> segmentedBatch : segmentedBatches.values()) {
+            // Iterate over the batch
+
+            for (int i = 0; i < segmentedBatch.size(); i += batchSize) {
+
+                // Create a sublist of the batch given the batchSize
+                batch = segmentedBatch.subList(i, Math.min(i + batchSize, segmentedBatch.size()));
+                List<PutRecordsRequestEntry> putRecordsRequestEntryList = new ArrayList<>();
+                streamName = batch.get(0).destination();
+
+                for (ChangeEvent<Object, Object> record : batch) {
+
+                    Object rv = record.value();
+                    if (rv == null) {
+                        rv = "";
+                    }
+                    PutRecordsRequestEntry putRecordsRequestEntry = PutRecordsRequestEntry.builder()
+                            .partitionKey((record.key() != null) ? getString(record.key()) : nullKey)
+                            .data(SdkBytes.fromByteArray(getBytes(rv))).build();
+                    putRecordsRequestEntryList.add(putRecordsRequestEntry);
+                }
+
+                // Handle Error
+                boolean notSuccesful = true;
+                int attempts = 0;
+                List<PutRecordsRequestEntry> batchRequest = putRecordsRequestEntryList;
+
+                while (notSuccesful) {
+
+                    if (attempts >= RETRIES) {
+                        throw new DebeziumException("Exceeded maximum number of attempts to publish event");
+                    }
+
+                    try {
+                        PutRecordsResponse response = recordsSent(batchRequest, streamName);
+                        attempts++;
+                        if (response.failedRecordCount() > 0) {
+                            LOGGER.warn("Failed to send {} number of records, retrying", response.failedRecordCount());
+                            Metronome.sleeper(RETRY_INTERVAL, Clock.SYSTEM).pause();
+
+                            final List<PutRecordsResultEntry> putRecordsResults = response.records();
+                            List<PutRecordsRequestEntry> failedRecordsList = new ArrayList<>();
+
+                            for (int index = 0; index < putRecordsResults.size(); index++) {
+                                PutRecordsResultEntry entryResult = putRecordsResults.get(index);
+                                if (entryResult.errorCode() != null) {
+                                    failedRecordsList.add(putRecordsRequestEntryList.get(index));
+                                }
+                            }
+                            batchRequest = failedRecordsList;
+
+                        }
+                        else {
+                            notSuccesful = false;
+                        }
+
+                    }
+                    catch (KinesisException exception) {
+                        LOGGER.warn("Failed to send record to {}", streamName, exception);
+                        attempts++;
+                        Metronome.sleeper(RETRY_INTERVAL, Clock.SYSTEM).pause();
+                    }
+                }
+
+                for (ChangeEvent<Object, Object> record : batch) {
+                    committer.markProcessed(record);
                 }
-                Metronome.sleeper(RETRY_INTERVAL, Clock.SYSTEM).pause();
             }
-            committer.markProcessed(record);
         }
+
+        // Mark Batch Finished
         committer.markBatchFinished();
     }
 
-    private boolean recordSent(ChangeEvent<Object, Object> record) {
-        Object rv = record.value();
-        if (rv == null) {
-            rv = "";
-        }
+    private PutRecordsResponse recordsSent(List<PutRecordsRequestEntry> putRecordsRequestEntryList, String streamName) {
 
-        final PutRecordRequest putRecord = PutRecordRequest.builder()
-                .partitionKey((record.key() != null) ? getString(record.key()) : nullKey)
-                .streamName(streamNameMapper.map(record.destination()))
-                .data(SdkBytes.fromByteArray(getBytes(rv)))
-                .build();
+        // Create a PutRecordsRequest
+        PutRecordsRequest putRecordsRequest = PutRecordsRequest.builder().streamName(streamNameMapper.map(streamName)).records(putRecordsRequestEntryList).build();
 
-        try {
-            client.putRecord(putRecord);
-            return true;
-        }
-        catch (KinesisException exception) {
-            LOGGER.warn("Failed to send record to {}", record.destination(), exception);
-            return false;
-        }
+        // Send Request
+        PutRecordsResponse putRecordsResponse = client.putRecords(putRecordsRequest);
+        LOGGER.trace("Response Receieved: " + putRecordsResponse);
+        return putRecordsResponse;
     }
 }