[FLINK-37298] Added Pluggable Components for BatchStrategy & BufferWrapper in AsyncSinkWriter.

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java

@@ -18,6 +18,7 @@
 package org.apache.flink.connector.base.sink.writer;
 
 import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.common.operators.MailboxExecutor;
 import org.apache.flink.api.common.operators.ProcessingTimeService;
 import org.apache.flink.api.connector.sink2.StatefulSinkWriter;
@@ -31,13 +32,12 @@
 import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
 import org.apache.flink.util.Preconditions;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.io.Serializable;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Deque;
 import java.util.List;
 import java.util.ListIterator;
 import java.util.concurrent.ScheduledFuture;
@@ -101,19 +101,25 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
 
     /**
      * Buffer to hold request entries that should be persisted into the destination, along with its
-     * size in bytes.
+     * total size in bytes.
      *
-     * <p>A request entry contain all relevant details to make a call to the destination. Eg, for
-     * Kinesis Data Streams a request entry contains the payload and partition key.
+     * <p>This buffer is managed using {@link BufferWrapper}, allowing sink implementations to
+     * define their own optimized buffering strategies. By default, {@link DequeBufferWrapper} is
+     * used.
      *
-     * <p>It seems more natural to buffer InputT, ie, the events that should be persisted, rather
-     * than RequestEntryT. However, in practice, the response of a failed request call can make it
-     * very hard, if not impossible, to reconstruct the original event. It is much easier, to just
-     * construct a new (retry) request entry from the response and add that back to the queue for
-     * later retry.
+     * <p>The buffer stores {@link RequestEntryWrapper} objects rather than raw {@link
+     * RequestEntryT} instances, as buffering wrapped entries allows for better tracking of size and
+     * metadata. This also makes it easier to handle retries by prioritising failed entries back
+     * into the buffer.
+     */
+    private final BufferWrapper<RequestEntryT> bufferedRequestEntries;
+
+    /**
+     * Batch component responsible for forming a batch of request entries from the buffer when the
+     * sink is ready to flush. This determines the logic of including entries in a batch from the
+     * buffered requests.
      */
-    private final Deque<RequestEntryWrapper<RequestEntryT>> bufferedRequestEntries =
-            new ArrayDeque<>();
+    private final BatchCreator<RequestEntryT> batchCreator;
 
     /**
      * Tracks all pending async calls that have been executed since the last checkpoint. Calls that
@@ -126,12 +132,6 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      */
     private int inFlightRequestsCount;
 
-    /**
-     * Tracks the cumulative size of all elements in {@code bufferedRequestEntries} to facilitate
-     * the criterion for flushing after {@code maxBatchSizeInBytes} is reached.
-     */
-    private double bufferedRequestEntriesTotalSizeInBytes;
-
     private boolean existsActiveTimerCallback = false;
 
     /**
@@ -213,11 +213,26 @@ protected void submitRequestEntries(
      */
     protected abstract long getSizeInBytes(RequestEntryT requestEntry);
 
+    /**
+     * This constructor is deprecated. Users should use {@link #AsyncSinkWriter(ElementConverter,
+     * WriterInitContext, AsyncSinkWriterConfiguration, Collection, BatchCreator, BufferWrapper)}.
+     */
+    @Deprecated
     public AsyncSinkWriter(
             ElementConverter<InputT, RequestEntryT> elementConverter,
             WriterInitContext context,
             AsyncSinkWriterConfiguration configuration,
             Collection<BufferedRequestState<RequestEntryT>> states) {
+        this(elementConverter, context, configuration, states, null, null);
+    }
+
+    public AsyncSinkWriter(
+            ElementConverter<InputT, RequestEntryT> elementConverter,
+            WriterInitContext context,
+            AsyncSinkWriterConfiguration configuration,
+            Collection<BufferedRequestState<RequestEntryT>> states,
+            @Nullable BatchCreator<RequestEntryT> batchCreator,
+            @Nullable BufferWrapper<RequestEntryT> bufferedRequestEntries) {
         this.elementConverter = elementConverter;
         this.mailboxExecutor = context.getMailboxExecutor();
         this.timeService = context.getProcessingTimeService();
@@ -245,10 +260,17 @@ public AsyncSinkWriter(
         this.rateLimitingStrategy = configuration.getRateLimitingStrategy();
         this.requestTimeoutMS = configuration.getRequestTimeoutMS();
         this.failOnTimeout = configuration.isFailOnTimeout();
-
+        this.bufferedRequestEntries =
+                bufferedRequestEntries == null
+                        ? new DequeBufferWrapper.Builder<RequestEntryT>().build()
+                        : bufferedRequestEntries;
+        this.batchCreator =
+                batchCreator == null
+                        ? new SimpleBatchCreator.Builder<RequestEntryT>()
+                                .setMaxBatchSizeInBytes(maxBatchSizeInBytes)
+                                .build()
+                        : batchCreator;
         this.inFlightRequestsCount = 0;
-        this.bufferedRequestEntriesTotalSizeInBytes = 0;
-
         this.metrics = context.metricGroup();
         this.metrics.setCurrentSendTimeGauge(() -> this.ackTime - this.lastSendTimestamp);
         this.numBytesOutCounter = this.metrics.getIOMetricGroup().getNumBytesOutCounter();
@@ -303,7 +325,7 @@ public void write(InputT element, Context context) throws IOException, Interrupt
     private void nonBlockingFlush() throws InterruptedException {
         while (!rateLimitingStrategy.shouldBlock(createRequestInfo())
                 && (bufferedRequestEntries.size() >= getNextBatchSizeLimit()
-                        || bufferedRequestEntriesTotalSizeInBytes >= maxBatchSizeInBytes)) {
+                        || bufferedRequestEntries.totalSizeInBytes() >= maxBatchSizeInBytes)) {
             flush();
         }
     }
@@ -327,7 +349,12 @@ private void flush() throws InterruptedException {
             requestInfo = createRequestInfo();
         }
 
-        List<RequestEntryT> batch = createNextAvailableBatch(requestInfo);
+        Batch<RequestEntryT> batchCreationResult =
+                batchCreator.createNextBatch(requestInfo, bufferedRequestEntries);
+        List<RequestEntryT> batch = batchCreationResult.getBatchEntries();
+        numBytesOutCounter.inc(batchCreationResult.getSizeInBytes());
+        numRecordsOutCounter.inc(batchCreationResult.getRecordCount());
+
         if (batch.isEmpty()) {
             return;
         }
@@ -344,31 +371,6 @@ private int getNextBatchSize() {
         return Math.min(getNextBatchSizeLimit(), bufferedRequestEntries.size());
     }
 
-    /**
-     * Creates the next batch of request entries while respecting the {@code maxBatchSize} and
-     * {@code maxBatchSizeInBytes}. Also adds these to the metrics counters.
-     */
-    private List<RequestEntryT> createNextAvailableBatch(RequestInfo requestInfo) {
-        List<RequestEntryT> batch = new ArrayList<>(requestInfo.getBatchSize());
-
-        long batchSizeBytes = 0;
-        for (int i = 0; i < requestInfo.getBatchSize(); i++) {
-            long requestEntrySize = bufferedRequestEntries.peek().getSize();
-            if (batchSizeBytes + requestEntrySize > maxBatchSizeInBytes) {
-                break;
-            }
-            RequestEntryWrapper<RequestEntryT> elem = bufferedRequestEntries.remove();
-            batch.add(elem.getRequestEntry());
-            bufferedRequestEntriesTotalSizeInBytes -= requestEntrySize;
-            batchSizeBytes += requestEntrySize;
-        }
-
-        numRecordsOutCounter.inc(batch.size());
-        numBytesOutCounter.inc(batchSizeBytes);
-
-        return batch;
-    }
-
     /**
      * Marks an in-flight request as completed and prepends failed requestEntries back to the
      * internal requestEntry buffer for later retry.
@@ -409,13 +411,7 @@ private void addEntryToBuffer(RequestEntryWrapper<RequestEntryT> entry, boolean
                             entry.getSize(), maxRecordSizeInBytes));
         }
 
-        if (insertAtHead) {
-            bufferedRequestEntries.addFirst(entry);
-        } else {
-            bufferedRequestEntries.add(entry);
-        }
-
-        bufferedRequestEntriesTotalSizeInBytes += entry.getSize();
+        bufferedRequestEntries.add(entry, insertAtHead);
     }
 
     /**
@@ -428,7 +424,7 @@ private void addEntryToBuffer(RequestEntryWrapper<RequestEntryT> entry, boolean
      */
     @Override
     public void flush(boolean flush) throws InterruptedException {
-        while (inFlightRequestsCount > 0 || (bufferedRequestEntries.size() > 0 && flush)) {
+        while (inFlightRequestsCount > 0 || (!bufferedRequestEntries.isEmpty() && flush)) {
             yieldIfThereExistsInFlightRequests();
             if (flush) {
                 flush();
@@ -545,4 +541,14 @@ public void timeout() {
             }
         }
     }
+
+    @VisibleForTesting
+    BufferWrapper<RequestEntryT> getBufferedRequestEntries() {
+        return bufferedRequestEntries;
+    }
+
+    @VisibleForTesting
+    BatchCreator<RequestEntryT> getBatchCreator() {
+        return batchCreator;
+    }
 }

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/Batch.java

@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.flink.connector.base.sink.writer;
+
+import org.apache.flink.annotation.PublicEvolving;
+
+import java.io.Serializable;
+import java.util.List;
+
+/**
+ * A container for the result of creating a batch of request entries, including:
+ *
+ * <ul>
+ *   <li>The actual list of entries forming the batch
+ *   <li>The total size in bytes of those entries
+ *   <li>The total number of entries in the batch
+ * </ul>
+ *
+ * <p>Instances of this class are typically created by a {@link BatchCreator} to summarize which
+ * entries have been selected for sending downstream and to provide any relevant metrics for
+ * tracking, such as the byte size or the record count.
+ *
+ * @param <RequestEntryT> the type of request entry in this batch
+ */
+@PublicEvolving
+public class Batch<RequestEntryT extends Serializable> {
+
+    /** The list of request entries in this batch. */
+    private final List<RequestEntryT> batchEntries;
+
+    /** The total size in bytes of the entire batch. */
+    private final long sizeInBytes;
+
+    /** The total number of entries in the batch. */
+    private final int recordCount;
+
+    /**
+     * Creates a new {@code Batch} with the specified entries, total size, and record count.
+     *
+     * @param requestEntries the list of request entries that form the batch
+     * @param sizeInBytes the total size in bytes of the entire batch
+     * @param recordCount the total number of entries in the batch
+     */
+    public Batch(List<RequestEntryT> requestEntries, long sizeInBytes, int recordCount) {
+        this.batchEntries = requestEntries;
+        this.sizeInBytes = sizeInBytes;
+        this.recordCount = recordCount;
+    }
+
+    /**
+     * Returns the list of request entries in this batch.
+     *
+     * @return a list of request entries for the batch
+     */
+    public List<RequestEntryT> getBatchEntries() {
+        return batchEntries;
+    }
+
+    /**
+     * Returns the total size in bytes of the batch.
+     *
+     * @return the batch's cumulative byte size
+     */
+    public long getSizeInBytes() {
+        return sizeInBytes;
+    }
+
+    /**
+     * Returns the total number of entries in the batch.
+     *
+     * @return the record count in the batch
+     */
+    public int getRecordCount() {
+        return recordCount;
+    }
+}

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/BatchCreator.java

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.flink.connector.base.sink.writer;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.connector.base.sink.writer.strategy.RequestInfo;
+
+import java.io.Serializable;
+import java.util.Deque;
+
+/**
+ * A pluggable interface for forming batches of request entries from a buffer. Implementations
+ * control how many entries are grouped together and in what manner before sending them downstream.
+ *
+ * <p>The {@code AsyncSinkWriter} (or similar sink component) calls {@link
+ * #createNextBatch(RequestInfo, BufferWrapper)} (RequestInfo, Deque)} when it decides to flush or
+ * otherwise gather a new batch of elements. For instance, a batch creator might limit the batch by
+ * the number of elements, total payload size, or any custom partition-based strategy.
+ *
+ * @param <RequestEntryT> the type of the request entries to be batched
+ */
+@PublicEvolving
+public interface BatchCreator<RequestEntryT extends Serializable> {
+
+    /**
+     * Creates the next batch of request entries from the current buffer.
+     *
+     * <p>This method is typically invoked when the sink determines that it's time to flush—e.g.,
+     * based on rate limiting, a buffer-size threshold, or a time-based trigger. The implementation
+     * can select as many entries from {@code bufferedRequestEntries} as it deems appropriate for a
+     * single batch, subject to any internal constraints (for example, a maximum byte size).
+     *
+     * @param requestInfo information about the desired request properties or constraints (e.g., an
+     *     allowed batch size or other relevant hints)
+     * @param bufferedRequestEntries a {@link Deque} of all currently buffered entries waiting to be
+     *     grouped into batches
+     * @return a {@link Batch} containing the new batch of entries along with metadata about the
+     *     batch (e.g., total byte size, record count)
+     */
+    Batch<RequestEntryT> createNextBatch(
+            RequestInfo requestInfo, BufferWrapper<RequestEntryT> bufferedRequestEntries);
+
+    /**
+     * Generic builder interface for creating instances of {@link BatchCreator}.
+     *
+     * @param <R> The type of {@link BatchCreator} that the builder will create.
+     * @param <RequestEntryT> The type of request entries that the batch creator will process.
+     */
+    interface Builder<R extends BatchCreator<RequestEntryT>, RequestEntryT extends Serializable> {
+        /**
+         * Constructs and returns an instance of {@link BatchCreator} with the configured
+         * parameters.
+         *
+         * @return A new instance of {@link BatchCreator}.
+         */
+        R build();
+    }
+}