You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by da...@apache.org on 2021/09/15 19:15:36 UTC

[flink] branch master updated (2da73ed -> 583e20d)

This is an automated email from the ASF dual-hosted git repository.

dannycranmer pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git.


    from 2da73ed  [FLINK-24277][connector/kafka] Add OffsetsInitializerValidator interface for validating offset initializer in KafkaSourceBuilder
     new de36259  [FLINK-24234][connectors] Added byte based flushing for AsyncSinkWriter
     new 583e20d  [FLINK-24234][connectors] Time based flushing for AsyncSinkWriter

The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:
 .../base/sink/writer/AsyncSinkWriter.java          |  99 +++++-
 ...mentConverter.java => RequestEntryWrapper.java} |  32 +-
 .../connector/base/sink/ArrayListAsyncSink.java    |  23 +-
 .../connector/base/sink/AsyncSinkBaseITCase.java   |   2 +-
 .../base/sink/writer/AsyncSinkWriterTest.java      | 394 ++++++++++++++++++++-
 5 files changed, 519 insertions(+), 31 deletions(-)
 copy flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/{ElementConverter.java => RequestEntryWrapper.java} (55%)

[flink] 02/02: [FLINK-24234][connectors] Time based flushing for AsyncSinkWriter

Posted by da...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

dannycranmer pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 583e20dcb97ed56eae76cdc949895086965d54c6
Author: Zichen Liu <zi...@amazon.com>
AuthorDate: Wed Sep 8 18:53:58 2021 +0100

    [FLINK-24234][connectors] Time based flushing for AsyncSinkWriter
---
 .../base/sink/writer/AsyncSinkWriter.java          | 104 +++---
 .../base/sink/writer/RequestEntryWrapper.java      |  48 +++
 .../connector/base/sink/ArrayListAsyncSink.java    |  21 +-
 .../connector/base/sink/AsyncSinkBaseITCase.java   |   3 +-
 .../base/sink/writer/AsyncSinkWriterTest.java      | 390 ++++++++++++++++++++-
 5 files changed, 502 insertions(+), 64 deletions(-)

diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
index 5c6f0e4..b5f6b40 100644
--- a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
+++ b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
@@ -33,6 +33,7 @@ import java.util.Collections;
 import java.util.Deque;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * A generic sink writer that handles the general behaviour of a sink such as batching and flushing,
@@ -50,15 +51,14 @@ import java.util.function.Consumer;
 public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable>
         implements SinkWriter<InputT, Void, Collection<RequestEntryT>> {
 
-    private static final int BYTES_IN_MB = 1024 * 1024;
-
     private final MailboxExecutor mailboxExecutor;
     private final Sink.ProcessingTimeService timeService;
 
     private final int maxBatchSize;
     private final int maxInFlightRequests;
     private final int maxBufferedRequests;
-    private final double flushOnBufferSizeMB;
+    private final long flushOnBufferSizeInBytes;
+    private final long maxTimeInBufferMS;
 
     /**
      * The ElementConverter provides a mapping between for the elements of a stream to request
@@ -70,7 +70,8 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
     private final ElementConverter<InputT, RequestEntryT> elementConverter;
 
     /**
-     * Buffer to hold request entries that should be persisted into the destination.
+     * Buffer to hold request entries that should be persisted into the destination, along with its
+     * 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.
@@ -81,7 +82,8 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      * construct a new (retry) request entry from the response and add that back to the queue for
      * later retry.
      */
-    private final Deque<RequestEntryT> bufferedRequestEntries = new ArrayDeque<>();
+    private final Deque<RequestEntryWrapper<RequestEntryT>> bufferedRequestEntries =
+            new ArrayDeque<>();
 
     /**
      * Tracks all pending async calls that have been executed since the last checkpoint. Calls that
@@ -100,16 +102,11 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
 
     /**
      * Tracks the cumulative size of all elements in {@code bufferedRequestEntries} to facilitate
-     * the criterion for flushing after {@code flushOnBufferSizeMB} is reached.
+     * the criterion for flushing after {@code flushOnBufferSizeInBytes} is reached.
      */
-    private double bufferedRequestEntriesTotalSizeMB;
+    private double bufferedRequestEntriesTotalSizeInBytes;
 
-    /**
-     * Tracks the size of each element in {@code bufferedRequestEntries}. The sizes are stored in MB
-     * and the position in the deque reflects the position of the corresponding element in {@code
-     * bufferedRequestEntries}.
-     */
-    private final Deque<Double> bufferedRequestEntriesSizeMB = new ArrayDeque<>();
+    private boolean existsActiveTimerCallback = false;
 
     /**
      * This method specifies how to persist buffered request entries into the destination. It is
@@ -143,7 +140,7 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      * @param requestEntry the requestEntry for which we want to know the size
      * @return the size of the requestEntry, as defined previously
      */
-    protected abstract int getSizeInBytes(RequestEntryT requestEntry);
+    protected abstract long getSizeInBytes(RequestEntryT requestEntry);
 
     public AsyncSinkWriter(
             ElementConverter<InputT, RequestEntryT> elementConverter,
@@ -151,7 +148,8 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
             int maxBatchSize,
             int maxInFlightRequests,
             int maxBufferedRequests,
-            double flushOnBufferSizeMB) {
+            long flushOnBufferSizeInBytes,
+            long maxTimeInBufferMS) {
         this.elementConverter = elementConverter;
         this.mailboxExecutor = context.getMailboxExecutor();
         this.timeService = context.getProcessingTimeService();
@@ -160,6 +158,8 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
         Preconditions.checkArgument(maxBatchSize > 0);
         Preconditions.checkArgument(maxBufferedRequests > 0);
         Preconditions.checkArgument(maxInFlightRequests > 0);
+        Preconditions.checkArgument(flushOnBufferSizeInBytes > 0);
+        Preconditions.checkArgument(maxTimeInBufferMS > 0);
         Preconditions.checkArgument(
                 maxBufferedRequests > maxBatchSize,
                 "The maximum number of requests that may be buffered should be strictly"
@@ -167,30 +167,40 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
         this.maxBatchSize = maxBatchSize;
         this.maxInFlightRequests = maxInFlightRequests;
         this.maxBufferedRequests = maxBufferedRequests;
-        this.flushOnBufferSizeMB = flushOnBufferSizeMB;
+        this.flushOnBufferSizeInBytes = flushOnBufferSizeInBytes;
+        this.maxTimeInBufferMS = maxTimeInBufferMS;
 
         this.inFlightRequestsCount = 0;
-        this.bufferedRequestEntriesTotalSizeMB = 0;
+        this.bufferedRequestEntriesTotalSizeInBytes = 0;
+    }
+
+    private void registerCallback() {
+        Sink.ProcessingTimeService.ProcessingTimeCallback ptc =
+                instant -> {
+                    existsActiveTimerCallback = false;
+                    while (!bufferedRequestEntries.isEmpty()) {
+                        flush();
+                    }
+                };
+        timeService.registerProcessingTimer(
+                timeService.getCurrentProcessingTime() + maxTimeInBufferMS, ptc);
+        existsActiveTimerCallback = true;
     }
 
     @Override
     public void write(InputT element, Context context) throws IOException, InterruptedException {
         while (bufferedRequestEntries.size() >= maxBufferedRequests) {
-            mailboxExecutor.yield();
+            mailboxExecutor.tryYield();
         }
 
-        RequestEntryT requestEntry = elementConverter.apply(element, context);
-        double requestEntrySizeMB = getSizeInMB(requestEntry);
-        bufferedRequestEntries.add(requestEntry);
-        bufferedRequestEntriesSizeMB.add(requestEntrySizeMB);
-        bufferedRequestEntriesTotalSizeMB += requestEntrySizeMB;
+        addEntryToBuffer(elementConverter.apply(element, context), false);
 
         flushIfAble();
     }
 
-    private void flushIfAble() throws InterruptedException {
+    private void flushIfAble() {
         while (bufferedRequestEntries.size() >= maxBatchSize
-                || bufferedRequestEntriesTotalSizeMB >= flushOnBufferSizeMB) {
+                || bufferedRequestEntriesTotalSizeInBytes >= flushOnBufferSizeInBytes) {
             flush();
         }
     }
@@ -201,18 +211,18 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      *
      * <p>The method blocks if too many async requests are in flight.
      */
-    private void flush() throws InterruptedException {
+    private void flush() {
         while (inFlightRequestsCount >= maxInFlightRequests) {
-            mailboxExecutor.yield();
+            mailboxExecutor.tryYield();
         }
 
         List<RequestEntryT> batch = new ArrayList<>(maxBatchSize);
 
         int batchSize = Math.min(maxBatchSize, bufferedRequestEntries.size());
         for (int i = 0; i < batchSize; i++) {
-            batch.add(bufferedRequestEntries.remove());
-            double elementSizeMB = bufferedRequestEntriesSizeMB.remove();
-            bufferedRequestEntriesTotalSizeMB -= elementSizeMB;
+            RequestEntryWrapper<RequestEntryT> elem = bufferedRequestEntries.remove();
+            batch.add(elem.getRequestEntry());
+            bufferedRequestEntriesTotalSizeInBytes -= elem.getSize();
         }
 
         if (batch.size() == 0) {
@@ -238,17 +248,24 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      */
     private void completeRequest(Collection<RequestEntryT> failedRequestEntries) {
         inFlightRequestsCount--;
-        failedRequestEntries.forEach(
-                failedEntry -> {
-                    bufferedRequestEntries.addFirst(failedEntry);
-                    double sizeOfFailedEntry = getSizeInMB(failedEntry);
-                    bufferedRequestEntriesSizeMB.addFirst(sizeOfFailedEntry);
-                    bufferedRequestEntriesTotalSizeMB += sizeOfFailedEntry;
-                });
+        failedRequestEntries.forEach(failedEntry -> addEntryToBuffer(failedEntry, true));
     }
 
-    private double getSizeInMB(RequestEntryT requestEntry){
-        return getSizeInBytes(requestEntry) / (double) BYTES_IN_MB;
+    private void addEntryToBuffer(RequestEntryT entry, boolean insertAtHead) {
+        if (bufferedRequestEntries.isEmpty() && !existsActiveTimerCallback) {
+            registerCallback();
+        }
+
+        RequestEntryWrapper<RequestEntryT> wrappedEntry =
+                new RequestEntryWrapper<>(entry, getSizeInBytes(entry));
+
+        if (insertAtHead) {
+            bufferedRequestEntries.addFirst(wrappedEntry);
+        } else {
+            bufferedRequestEntries.add(wrappedEntry);
+        }
+
+        bufferedRequestEntriesTotalSizeInBytes += wrappedEntry.getSize();
     }
 
     /**
@@ -260,9 +277,9 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      * <p>To this end, all in-flight requests need to completed before proceeding with the commit.
      */
     @Override
-    public List<Void> prepareCommit(boolean flush) throws InterruptedException {
+    public List<Void> prepareCommit(boolean flush) {
         while (inFlightRequestsCount > 0 || bufferedRequestEntries.size() > 0) {
-            mailboxExecutor.yield();
+            mailboxExecutor.tryYield();
             if (flush) {
                 flush();
             }
@@ -279,7 +296,10 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      */
     @Override
     public List<Collection<RequestEntryT>> snapshotState() {
-        return Arrays.asList(bufferedRequestEntries);
+        return Arrays.asList(
+                bufferedRequestEntries.stream()
+                        .map(RequestEntryWrapper::getRequestEntry)
+                        .collect(Collectors.toList()));
     }
 
     @Override
diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/RequestEntryWrapper.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/RequestEntryWrapper.java
new file mode 100644
index 0000000..2a44c10
--- /dev/null
+++ b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/RequestEntryWrapper.java
@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+/**
+ * A wrapper that contains a {@code RequestEntryT} ready to be written by the Sink Writer class,
+ * along with the size of that entry as defined by the method {@code getSizeInBytes(RequestEntryT)}
+ * of the {@code AsyncSinkWriter}.
+ *
+ * @param <RequestEntryT> Corresponds to the type parameter of the same name in {@code
+ *     AsyncSinkWriter}
+ */
+@PublicEvolving
+public class RequestEntryWrapper<RequestEntryT> {
+
+    private final RequestEntryT requestEntry;
+    private final long size;
+
+    public RequestEntryWrapper(RequestEntryT requestEntry, long size) {
+        this.requestEntry = requestEntry;
+        this.size = size;
+    }
+
+    public RequestEntryT getRequestEntry() {
+        return requestEntry;
+    }
+
+    public long getSize() {
+        return size;
+    }
+}
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
index 4661110..29c8c56 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
@@ -33,18 +33,24 @@ public class ArrayListAsyncSink extends AsyncSinkBase<String, Integer> {
     private final int maxBatchSize;
     private final int maxInFlightRequests;
     private final int maxBufferedRequests;
-    private final double flushOnBufferSizeMB;
+    private final long flushOnBufferSizeInBytes;
+    private final long maxTimeInBufferMS;
 
     public ArrayListAsyncSink() {
-        this(25, 1, 100, 0.1);
+        this(25, 1, 100, 100000, 1000);
     }
 
-    public ArrayListAsyncSink(int maxBatchSize, int maxInFlightRequests, int maxBufferedRequests,
-                              double flushOnBufferSizeMB) {
+    public ArrayListAsyncSink(
+            int maxBatchSize,
+            int maxInFlightRequests,
+            int maxBufferedRequests,
+            long flushOnBufferSizeInBytes,
+            long maxTimeInBufferMS) {
         this.maxBatchSize = maxBatchSize;
         this.maxInFlightRequests = maxInFlightRequests;
         this.maxBufferedRequests = maxBufferedRequests;
-        this.flushOnBufferSizeMB = flushOnBufferSizeMB;
+        this.flushOnBufferSizeInBytes = flushOnBufferSizeInBytes;
+        this.maxTimeInBufferMS = maxTimeInBufferMS;
     }
 
     @Override
@@ -59,7 +65,8 @@ public class ArrayListAsyncSink extends AsyncSinkBase<String, Integer> {
                 maxBatchSize,
                 maxInFlightRequests,
                 maxBufferedRequests,
-                flushOnBufferSizeMB) {
+                flushOnBufferSizeInBytes,
+                maxTimeInBufferMS) {
 
             @Override
             protected void submitRequestEntries(
@@ -69,7 +76,7 @@ public class ArrayListAsyncSink extends AsyncSinkBase<String, Integer> {
             }
 
             @Override
-            protected int getSizeInBytes(Integer requestEntry) {
+            protected long getSizeInBytes(Integer requestEntry) {
                 return 4;
             }
         };
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
index 4f6d5b1..a7ed47d 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
@@ -21,6 +21,7 @@ import org.apache.flink.api.common.restartstrategy.RestartStrategies;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.runtime.client.JobExecutionException;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
@@ -41,7 +42,7 @@ public class AsyncSinkBaseITCase {
     public void testFailuresOnPersistingToDestinationAreCaughtAndRaised() {
         env.fromSequence(999_999, 1_000_100)
                 .map(Object::toString)
-                .sinkTo(new ArrayListAsyncSink(1, 1, 2, 10));
+                .sinkTo(new ArrayListAsyncSink(1, 1, 2, 10, 1000));
         Exception e =
                 assertThrows(
                         JobExecutionException.class,
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
index 8db0aec..b3754d0 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
@@ -21,11 +21,13 @@ import org.apache.flink.api.common.operators.MailboxExecutor;
 import org.apache.flink.api.connector.sink.Sink;
 import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
 import org.apache.flink.streaming.runtime.tasks.StreamTaskActionExecutor;
+import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
 import org.apache.flink.streaming.runtime.tasks.mailbox.MailboxExecutorImpl;
 import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailboxImpl;
 import org.apache.flink.util.UserCodeClassLoader;
 import org.apache.flink.util.function.RunnableWithException;
 import org.apache.flink.util.function.ThrowingRunnable;
+
 import org.junit.Before;
 import org.junit.Test;
 
@@ -38,12 +40,18 @@ import java.util.List;
 import java.util.OptionalLong;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 /**
  * Unit Tests the functionality of AsyncSinkWriter without any assumptions of what a concrete
@@ -51,7 +59,6 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
  */
 public class AsyncSinkWriterTest {
 
-    private static final int BYTES_IN_MB = 1024 * 1024;
     private final List<Integer> res = new ArrayList<>();
     private final SinkInitContext sinkInitContext = new SinkInitContext();
 
@@ -93,6 +100,15 @@ public class AsyncSinkWriterTest {
     }
 
     @Test
+    public void testThatMailboxYieldDoesNotBlockWhileATimerIsRegisteredAndHasYetToElapse()
+            throws Exception {
+        AsyncSinkWriterImpl sink = new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100, false);
+        sink.write(String.valueOf(0));
+        sink.prepareCommit(true);
+        assertEquals(1, res.size());
+    }
+
+    @Test
     public void testThatSnapshotsAreTakenOfBufferCorrectlyBeforeAndAfterAutomaticFlush()
             throws IOException, InterruptedException {
         AsyncSinkWriterImpl sink = new AsyncSinkWriterImpl(sinkInitContext, 3, 1, 100, false);
@@ -132,7 +148,8 @@ public class AsyncSinkWriterTest {
         sink.write("95");
         sink.write("35");
         Exception e = assertThrows(RuntimeException.class, () -> sink.write("135"));
-        assertEquals("Deliberate runtime exception occurred in SinkWriterImplementation.",
+        assertEquals(
+                "Deliberate runtime exception occurred in SinkWriterImplementation.",
                 e.getMessage());
         assertEquals(3, res.size());
     }
@@ -241,10 +258,10 @@ public class AsyncSinkWriterTest {
     }
 
     @Test
-    public void flushThresholdMetBeforeBatchLimitWillCreateASmallerBatchOfSizeAboveThreshold()
+    public void testFlushThresholdMetBeforeBatchLimitWillCreateASmallerBatchOfSizeAboveThreshold()
             throws IOException, InterruptedException {
-        AsyncSinkWriterImpl sink = new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100,
-                (double) 30 / BYTES_IN_MB, true);
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100, 30, 1000, false);
 
         /* Sink has flush threshold of 30 bytes, each integer is 4 bytes, therefore, flushing
          * should occur once 8 elements have been written.
@@ -253,6 +270,256 @@ public class AsyncSinkWriterTest {
             sink.write(String.valueOf(i));
         }
         assertEquals(8, res.size());
+        sink.write(String.valueOf(15));
+        assertEquals(16, res.size());
+    }
+
+    @Test
+    public void testThatWhenNumberOfItemAndSizeOfRecordThresholdsAreMetSimultaneouslyAFlushOccurs()
+            throws IOException, InterruptedException {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 8, 1, 100, 32, 1000, false);
+
+        for (int i = 0; i < 8; i++) {
+            sink.write(String.valueOf(i));
+        }
+        assertEquals(8, res.size());
+        for (int i = 8; i < 16; i++) {
+            sink.write(String.valueOf(i));
+        }
+        assertEquals(16, res.size());
+    }
+
+    @Test
+    public void testThatIntermittentlyFailingEntriesAreEnqueuedOnToTheBufferWithCorrectSize()
+            throws IOException, InterruptedException {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100, 110, 1000, true);
+
+        sink.write(String.valueOf(225)); // Buffer: 100/110B; 1/10 elements; 0 inflight
+        sink.write(String.valueOf(1)); //   Buffer: 104/110B; 2/10 elements; 0 inflight
+        sink.write(String.valueOf(2)); //   Buffer: 108/110B; 3/10 elements; 0 inflight
+        sink.write(String.valueOf(3)); //   Buffer: 112/110B; 4/10 elements; 0 inflight -- flushing
+        assertEquals(3, res.size()); // Element 225 failed on first attempt
+        sink.write(String.valueOf(4)); //   Buffer:   4/110B; 1/10 elements; 1 inflight
+        sink.write(String.valueOf(5)); //   Buffer:   8/110B; 2/10 elements; 1 inflight
+        sink.write(String.valueOf(6)); //   Buffer:  12/110B; 3/10 elements; 1 inflight
+        sink.write(String.valueOf(325)); // Buffer: 112/110B; 4/10 elements; 1 inflight -- flushing
+
+        assertEquals(Arrays.asList(1, 2, 3, 225, 4, 5, 6), res);
+    }
+
+    @Test
+    public void testThatABatchWithSizeSmallerThanMaxBatchSizeIsFlushedOnTimeoutExpiry()
+            throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 20, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        tpts.setCurrentTime(0L);
+        for (int i = 0; i < 8; i++) {
+            sink.write(String.valueOf(i));
+        }
+
+        tpts.setCurrentTime(99L);
+        assertEquals(0, res.size());
+        tpts.setCurrentTime(100L);
+        assertEquals(8, res.size());
+    }
+
+    @Test
+    public void testThatTimeBasedBatchPicksUpAllRelevantItemsUpUntilExpiryOfTimer()
+            throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 20, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        for (int i = 0; i < 98; i++) {
+            tpts.setCurrentTime(i);
+            sink.write(String.valueOf(i));
+        }
+        tpts.setCurrentTime(99L);
+        assertEquals(90, res.size());
+        tpts.setCurrentTime(100L);
+        assertEquals(98, res.size());
+    }
+
+    @Test
+    public void testThatOneAndOnlyOneCallbackIsEverRegistered() throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 20, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        tpts.setCurrentTime(0L);
+        sink.write("1"); // A timer is registered here to elapse at t=100
+        assertEquals(0, res.size());
+        tpts.setCurrentTime(10L);
+        sink.prepareCommit(true);
+        assertEquals(1, res.size());
+        tpts.setCurrentTime(20L); // At t=20, we write a new element that should not trigger another
+        sink.write("2"); // timer to be registered. If it is, it should elapse at t=120s.
+        assertEquals(1, res.size());
+        tpts.setCurrentTime(100L);
+        assertEquals(2, res.size());
+        sink.write("3");
+        tpts.setCurrentTime(199L); // At t=199s, our third element has not been written
+        assertEquals(2, res.size()); // therefore, no timer fired at 120s.
+        tpts.setCurrentTime(200L);
+        assertEquals(3, res.size());
+    }
+
+    @Test
+    public void testThatIntermittentlyFailingEntriesShouldBeFlushedWithMainBatchInTimeBasedFlush()
+            throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        tpts.setCurrentTime(0L);
+        sink.write("1");
+        sink.write("2");
+        sink.write("225");
+        tpts.setCurrentTime(100L);
+        assertEquals(2, res.size());
+        sink.write("3");
+        sink.write("4");
+        tpts.setCurrentTime(199L);
+        assertEquals(2, res.size());
+        tpts.setCurrentTime(200L);
+        assertEquals(5, res.size());
+    }
+
+    @Test
+    public void testThatFlushingAnEmptyBufferDoesNotResultInErrorOrFailure() throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 20, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        tpts.setCurrentTime(0L);
+        sink.write("1");
+        tpts.setCurrentTime(50L);
+        sink.prepareCommit(true);
+        assertEquals(1, res.size());
+        tpts.setCurrentTime(200L);
+    }
+
+    @Test
+    public void testThatOnExpiryOfAnOldTimeoutANewOneMayBeRegisteredImmediately() throws Exception {
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkWriterImpl(sinkInitContext, 10, 20, 100, 10000, 100, true);
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        tpts.setCurrentTime(0L);
+        sink.write("1");
+        tpts.setCurrentTime(100L);
+        assertEquals(1, res.size());
+        sink.write("2");
+        tpts.setCurrentTime(200L);
+        assertEquals(2, res.size());
+    }
+
+    /**
+     * This test considers what could happen if the timer elapses, triggering a flush, while a
+     * long-running call to {@code submitRequestEntries} remains uncompleted for some time. We have
+     * a countdown latch with an expiry of 500ms installed in the call to {@code
+     * submitRequestEntries} that blocks if the batch size received is 3 and subsequently accepts
+     * and succeeds with any value.
+     *
+     * <p>Let us call the thread writing "3" thread3 and the thread writing "4" thread4. Thread3
+     * will enter {@code submitRequestEntries} with 3 entries and release thread4. Thread3 becomes
+     * blocked for 500ms. Thread4 writes "4" to the buffer and is flushed when the timer triggers
+     * (timer was first set when "1" was written). Thread4 then is blocked during the flush phase
+     * since thread3 is in-flight and maxInFlightRequests=1. After 500ms elapses, thread3 is revived
+     * and proceeds, which also unblocks thread4. This results in 1, 2, 3 being written prior to 4.
+     *
+     * <p>This test also implicitly asserts that any thread in the SinkWriter must be the mailbox
+     * thread if it enters {@code mailbox.tryYield()}.
+     */
+    @Test
+    public void testThatInterleavingThreadsMayBlockEachOtherButDoNotCauseRaceConditions()
+            throws Exception {
+        CountDownLatch blockedWriteLatch = new CountDownLatch(1);
+        CountDownLatch delayedStartLatch = new CountDownLatch(1);
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkReleaseAndBlockWriterImpl(
+                        sinkInitContext,
+                        3,
+                        1,
+                        20,
+                        100,
+                        100,
+                        blockedWriteLatch,
+                        delayedStartLatch,
+                        true);
+
+        writeTwoElementsAndInterleaveTheNextTwoElements(sink, blockedWriteLatch, delayedStartLatch);
+        assertEquals(Arrays.asList(1, 2, 3, 4), res);
+    }
+
+    /**
+     * This test considers what could happen if the timer elapses, triggering a flush, while a
+     * long-running call to {@code submitRequestEntries} remains blocked. We have a countdown latch
+     * that blocks permanently until freed once the timer based flush is complete.
+     *
+     * <p>Let us call the thread writing "3" thread3 and the thread writing "4" thread4. Thread3
+     * will enter {@code submitRequestEntries} with 3 entries and release thread4. Thread3 becomes
+     * blocked. Thread4 writes "4" to the buffer and is flushed when the timer triggers (timer was
+     * first set when "1" was written). Thread4 completes and frees thread3. Thread3 is revived and
+     * proceeds. This results in 4 being written prior to 1, 2, 3.
+     *
+     * <p>This test also implicitly asserts that any thread in the SinkWriter must be the mailbox
+     * thread if it enters {@code mailbox.tryYield()}.
+     */
+    @Test
+    public void testThatIfOneInterleavedThreadIsBlockedTheOtherThreadWillContinueAndCorrectlyWrite()
+            throws Exception {
+        CountDownLatch blockedWriteLatch = new CountDownLatch(1);
+        CountDownLatch delayedStartLatch = new CountDownLatch(1);
+        AsyncSinkWriterImpl sink =
+                new AsyncSinkReleaseAndBlockWriterImpl(
+                        sinkInitContext,
+                        3,
+                        2,
+                        20,
+                        100,
+                        100,
+                        blockedWriteLatch,
+                        delayedStartLatch,
+                        false);
+
+        writeTwoElementsAndInterleaveTheNextTwoElements(sink, blockedWriteLatch, delayedStartLatch);
+        assertEquals(new ArrayList<>(Arrays.asList(4, 1, 2, 3)), res);
+    }
+
+    private void writeTwoElementsAndInterleaveTheNextTwoElements(
+            AsyncSinkWriterImpl sink,
+            CountDownLatch blockedWriteLatch,
+            CountDownLatch delayedStartLatch)
+            throws Exception {
+
+        TestProcessingTimeService tpts = sinkInitContext.getTestProcessingTimeService();
+        ExecutorService es = Executors.newFixedThreadPool(4);
+
+        tpts.setCurrentTime(0L);
+        sink.write("1");
+        sink.write("2");
+        es.submit(
+                () -> {
+                    try {
+                        sink.write("3");
+                    } catch (IOException | InterruptedException e) {
+                        e.printStackTrace();
+                    }
+                });
+
+        delayedStartLatch.await();
+        sink.write("4");
+        tpts.setCurrentTime(100L);
+        blockedWriteLatch.countDown();
+        es.shutdown();
+        assertTrue(
+                es.awaitTermination(500, TimeUnit.MILLISECONDS),
+                "Executor Service stuck at termination, not terminated after 500ms!");
     }
 
     private class AsyncSinkWriterImpl extends AsyncSinkWriter<String, Integer> {
@@ -261,14 +528,19 @@ public class AsyncSinkWriterTest {
         private final boolean simulateFailures;
 
         public AsyncSinkWriterImpl(
-                Sink.InitContext context, int maxBatchSize, int maxInFlightRequests,
-                int maxBufferedRequests, boolean simulateFailures) {
+                Sink.InitContext context,
+                int maxBatchSize,
+                int maxInFlightRequests,
+                int maxBufferedRequests,
+                boolean simulateFailures) {
             super(
                     (elem, ctx) -> Integer.parseInt(elem),
                     context,
                     maxBatchSize,
                     maxInFlightRequests,
-                    maxBufferedRequests, 100);
+                    maxBufferedRequests,
+                    10000000,
+                    1000);
             this.simulateFailures = simulateFailures;
         }
 
@@ -277,10 +549,17 @@ public class AsyncSinkWriterTest {
                 int maxBatchSize,
                 int maxInFlightRequests,
                 int maxBufferedRequests,
-                double flushOnBufferSizeMB,
+                long flushOnBufferSizeInBytes,
+                long maxTimeInBufferMS,
                 boolean simulateFailures) {
-            super((elem, ctx) -> Integer.parseInt(elem), context, maxBatchSize, maxInFlightRequests,
-                    maxBufferedRequests, flushOnBufferSizeMB);
+            super(
+                    (elem, ctx) -> Integer.parseInt(elem),
+                    context,
+                    maxBatchSize,
+                    maxInFlightRequests,
+                    maxBufferedRequests,
+                    flushOnBufferSizeInBytes,
+                    maxTimeInBufferMS);
             this.simulateFailures = simulateFailures;
         }
 
@@ -329,14 +608,24 @@ public class AsyncSinkWriterTest {
             }
         }
 
+        /**
+         * @return If we're simulating failures and the requestEntry value is greater than 200, then
+         *     the entry is size 100 bytes, otherwise each entry is 4 bytes.
+         */
         @Override
-        protected int getSizeInBytes(Integer requestEntry) {
-            return 4;
+        protected long getSizeInBytes(Integer requestEntry) {
+            return requestEntry > 200 && simulateFailures ? 100 : 4;
         }
     }
 
     private static class SinkInitContext implements Sink.InitContext {
 
+        private static final TestProcessingTimeService processingTimeService;
+
+        static {
+            processingTimeService = new TestProcessingTimeService();
+        }
+
         @Override
         public UserCodeClassLoader getUserCodeClassLoader() {
             return null;
@@ -370,7 +659,19 @@ public class AsyncSinkWriterTest {
 
         @Override
         public Sink.ProcessingTimeService getProcessingTimeService() {
-            return null;
+            return new Sink.ProcessingTimeService() {
+                @Override
+                public long getCurrentProcessingTime() {
+                    return processingTimeService.getCurrentProcessingTime();
+                }
+
+                @Override
+                public void registerProcessingTimer(
+                        long time, ProcessingTimeCallback processingTimerCallback) {
+                    processingTimeService.registerTimer(
+                            time, processingTimerCallback::onProcessingTime);
+                }
+            };
         }
 
         @Override
@@ -392,5 +693,66 @@ public class AsyncSinkWriterTest {
         public OptionalLong getRestoredCheckpointId() {
             return OptionalLong.empty();
         }
+
+        public TestProcessingTimeService getTestProcessingTimeService() {
+            return processingTimeService;
+        }
+    }
+
+    /**
+     * This SinkWriter releases the lock on existing threads blocked by {@code delayedStartLatch}
+     * and blocks itself until {@code blockedThreadLatch} is unblocked.
+     */
+    private class AsyncSinkReleaseAndBlockWriterImpl extends AsyncSinkWriterImpl {
+
+        private final CountDownLatch blockedThreadLatch;
+        private final CountDownLatch delayedStartLatch;
+        private final boolean blockForLimitedTime;
+
+        public AsyncSinkReleaseAndBlockWriterImpl(
+                Sink.InitContext context,
+                int maxBatchSize,
+                int maxInFlightRequests,
+                int maxBufferedRequests,
+                long flushOnBufferSizeInBytes,
+                long maxTimeInBufferMS,
+                CountDownLatch blockedThreadLatch,
+                CountDownLatch delayedStartLatch,
+                boolean blockForLimitedTime) {
+            super(
+                    context,
+                    maxBatchSize,
+                    maxInFlightRequests,
+                    maxBufferedRequests,
+                    flushOnBufferSizeInBytes,
+                    maxTimeInBufferMS,
+                    false);
+            this.blockedThreadLatch = blockedThreadLatch;
+            this.delayedStartLatch = delayedStartLatch;
+            this.blockForLimitedTime = blockForLimitedTime;
+        }
+
+        @Override
+        protected void submitRequestEntries(
+                List<Integer> requestEntries, Consumer<Collection<Integer>> requestResult) {
+            if (requestEntries.size() == 3) {
+                try {
+                    delayedStartLatch.countDown();
+                    if (blockForLimitedTime) {
+                        assertFalse(
+                                blockedThreadLatch.await(500, TimeUnit.MILLISECONDS),
+                                "The countdown latch was released before the full amount"
+                                        + "of time was reached.");
+                    } else {
+                        blockedThreadLatch.await();
+                    }
+                } catch (InterruptedException e) {
+                    fail("The unit test latch must not have been interrupted by another thread.");
+                }
+            }
+
+            res.addAll(requestEntries);
+            requestResult.accept(new ArrayList<>());
+        }
     }
 }

[flink] 01/02: [FLINK-24234][connectors] Added byte based flushing for AsyncSinkWriter

Posted by da...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

dannycranmer pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit de362595f215d442c937e6aae93048480e55ea25
Author: Zichen Liu <zi...@amazon.com>
AuthorDate: Wed Sep 8 11:30:12 2021 +0100

    [FLINK-24234][connectors] Added byte based flushing for AsyncSinkWriter
---
 .../base/sink/writer/AsyncSinkWriter.java          | 61 +++++++++++++++++++---
 .../connector/base/sink/ArrayListAsyncSink.java    | 16 ++++--
 .../connector/base/sink/AsyncSinkBaseITCase.java   |  3 +-
 .../base/sink/writer/AsyncSinkWriterTest.java      | 46 ++++++++++++----
 4 files changed, 106 insertions(+), 20 deletions(-)

diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
index 1d17d2f..5c6f0e4 100644
--- a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
+++ b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriter.java
@@ -50,12 +50,15 @@ import java.util.function.Consumer;
 public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable>
         implements SinkWriter<InputT, Void, Collection<RequestEntryT>> {
 
+    private static final int BYTES_IN_MB = 1024 * 1024;
+
     private final MailboxExecutor mailboxExecutor;
     private final Sink.ProcessingTimeService timeService;
 
     private final int maxBatchSize;
     private final int maxInFlightRequests;
     private final int maxBufferedRequests;
+    private final double flushOnBufferSizeMB;
 
     /**
      * The ElementConverter provides a mapping between for the elements of a stream to request
@@ -96,6 +99,19 @@ 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 flushOnBufferSizeMB} is reached.
+     */
+    private double bufferedRequestEntriesTotalSizeMB;
+
+    /**
+     * Tracks the size of each element in {@code bufferedRequestEntries}. The sizes are stored in MB
+     * and the position in the deque reflects the position of the corresponding element in {@code
+     * bufferedRequestEntries}.
+     */
+    private final Deque<Double> bufferedRequestEntriesSizeMB = new ArrayDeque<>();
+
+    /**
      * This method specifies how to persist buffered request entries into the destination. It is
      * implemented when support for a new destination is added.
      *
@@ -118,12 +134,24 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
     protected abstract void submitRequestEntries(
             List<RequestEntryT> requestEntries, Consumer<Collection<RequestEntryT>> requestResult);
 
+    /**
+     * This method allows the getting of the size of a {@code RequestEntryT} in bytes. The size in
+     * this case is measured as the total bytes that is written to the destination as a result of
+     * persisting this particular {@code RequestEntryT} rather than the serialized length (which may
+     * be the same).
+     *
+     * @param requestEntry the requestEntry for which we want to know the size
+     * @return the size of the requestEntry, as defined previously
+     */
+    protected abstract int getSizeInBytes(RequestEntryT requestEntry);
+
     public AsyncSinkWriter(
             ElementConverter<InputT, RequestEntryT> elementConverter,
             Sink.InitContext context,
             int maxBatchSize,
             int maxInFlightRequests,
-            int maxBufferedRequests) {
+            int maxBufferedRequests,
+            double flushOnBufferSizeMB) {
         this.elementConverter = elementConverter;
         this.mailboxExecutor = context.getMailboxExecutor();
         this.timeService = context.getProcessingTimeService();
@@ -139,6 +167,10 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
         this.maxBatchSize = maxBatchSize;
         this.maxInFlightRequests = maxInFlightRequests;
         this.maxBufferedRequests = maxBufferedRequests;
+        this.flushOnBufferSizeMB = flushOnBufferSizeMB;
+
+        this.inFlightRequestsCount = 0;
+        this.bufferedRequestEntriesTotalSizeMB = 0;
     }
 
     @Override
@@ -147,13 +179,18 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
             mailboxExecutor.yield();
         }
 
-        bufferedRequestEntries.add(elementConverter.apply(element, context));
+        RequestEntryT requestEntry = elementConverter.apply(element, context);
+        double requestEntrySizeMB = getSizeInMB(requestEntry);
+        bufferedRequestEntries.add(requestEntry);
+        bufferedRequestEntriesSizeMB.add(requestEntrySizeMB);
+        bufferedRequestEntriesTotalSizeMB += requestEntrySizeMB;
 
-        flushIfFull();
+        flushIfAble();
     }
 
-    private void flushIfFull() throws InterruptedException {
-        while (bufferedRequestEntries.size() >= maxBatchSize) {
+    private void flushIfAble() throws InterruptedException {
+        while (bufferedRequestEntries.size() >= maxBatchSize
+                || bufferedRequestEntriesTotalSizeMB >= flushOnBufferSizeMB) {
             flush();
         }
     }
@@ -174,6 +211,8 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
         int batchSize = Math.min(maxBatchSize, bufferedRequestEntries.size());
         for (int i = 0; i < batchSize; i++) {
             batch.add(bufferedRequestEntries.remove());
+            double elementSizeMB = bufferedRequestEntriesSizeMB.remove();
+            bufferedRequestEntriesTotalSizeMB -= elementSizeMB;
         }
 
         if (batch.size() == 0) {
@@ -199,7 +238,17 @@ public abstract class AsyncSinkWriter<InputT, RequestEntryT extends Serializable
      */
     private void completeRequest(Collection<RequestEntryT> failedRequestEntries) {
         inFlightRequestsCount--;
-        failedRequestEntries.forEach(bufferedRequestEntries::addFirst);
+        failedRequestEntries.forEach(
+                failedEntry -> {
+                    bufferedRequestEntries.addFirst(failedEntry);
+                    double sizeOfFailedEntry = getSizeInMB(failedEntry);
+                    bufferedRequestEntriesSizeMB.addFirst(sizeOfFailedEntry);
+                    bufferedRequestEntriesTotalSizeMB += sizeOfFailedEntry;
+                });
+    }
+
+    private double getSizeInMB(RequestEntryT requestEntry){
+        return getSizeInBytes(requestEntry) / (double) BYTES_IN_MB;
     }
 
     /**
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
index c4d6639..4661110 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/ArrayListAsyncSink.java
@@ -33,15 +33,18 @@ public class ArrayListAsyncSink extends AsyncSinkBase<String, Integer> {
     private final int maxBatchSize;
     private final int maxInFlightRequests;
     private final int maxBufferedRequests;
+    private final double flushOnBufferSizeMB;
 
     public ArrayListAsyncSink() {
-        this(25, 1, 100);
+        this(25, 1, 100, 0.1);
     }
 
-    public ArrayListAsyncSink(int maxBatchSize, int maxInFlightRequests, int maxBufferedRequests) {
+    public ArrayListAsyncSink(int maxBatchSize, int maxInFlightRequests, int maxBufferedRequests,
+                              double flushOnBufferSizeMB) {
         this.maxBatchSize = maxBatchSize;
         this.maxInFlightRequests = maxInFlightRequests;
         this.maxBufferedRequests = maxBufferedRequests;
+        this.flushOnBufferSizeMB = flushOnBufferSizeMB;
     }
 
     @Override
@@ -55,13 +58,20 @@ public class ArrayListAsyncSink extends AsyncSinkBase<String, Integer> {
                 context,
                 maxBatchSize,
                 maxInFlightRequests,
-                maxBufferedRequests) {
+                maxBufferedRequests,
+                flushOnBufferSizeMB) {
+
             @Override
             protected void submitRequestEntries(
                     List<Integer> requestEntries, Consumer<Collection<Integer>> requestResult) {
                 ArrayListDestination.putRecords(requestEntries);
                 requestResult.accept(Arrays.asList());
             }
+
+            @Override
+            protected int getSizeInBytes(Integer requestEntry) {
+                return 4;
+            }
         };
     }
 
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
index 63d51e3..4f6d5b1 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/AsyncSinkBaseITCase.java
@@ -21,7 +21,6 @@ import org.apache.flink.api.common.restartstrategy.RestartStrategies;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.runtime.client.JobExecutionException;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
-
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
@@ -42,7 +41,7 @@ public class AsyncSinkBaseITCase {
     public void testFailuresOnPersistingToDestinationAreCaughtAndRaised() {
         env.fromSequence(999_999, 1_000_100)
                 .map(Object::toString)
-                .sinkTo(new ArrayListAsyncSink(1, 1, 2));
+                .sinkTo(new ArrayListAsyncSink(1, 1, 2, 10));
         Exception e =
                 assertThrows(
                         JobExecutionException.class,
diff --git a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
index 378198d..8db0aec 100644
--- a/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
+++ b/flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/sink/writer/AsyncSinkWriterTest.java
@@ -26,7 +26,6 @@ import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailboxImpl;
 import org.apache.flink.util.UserCodeClassLoader;
 import org.apache.flink.util.function.RunnableWithException;
 import org.apache.flink.util.function.ThrowingRunnable;
-
 import org.junit.Before;
 import org.junit.Test;
 
@@ -52,6 +51,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
  */
 public class AsyncSinkWriterTest {
 
+    private static final int BYTES_IN_MB = 1024 * 1024;
     private final List<Integer> res = new ArrayList<>();
     private final SinkInitContext sinkInitContext = new SinkInitContext();
 
@@ -132,8 +132,7 @@ public class AsyncSinkWriterTest {
         sink.write("95");
         sink.write("35");
         Exception e = assertThrows(RuntimeException.class, () -> sink.write("135"));
-        assertEquals(
-                "Deliberate runtime exception occurred in SinkWriterImplementation.",
+        assertEquals("Deliberate runtime exception occurred in SinkWriterImplementation.",
                 e.getMessage());
         assertEquals(3, res.size());
     }
@@ -241,23 +240,47 @@ public class AsyncSinkWriterTest {
         assertEquals(z, new ArrayList<>(sink.snapshotState().get(0)));
     }
 
+    @Test
+    public void flushThresholdMetBeforeBatchLimitWillCreateASmallerBatchOfSizeAboveThreshold()
+            throws IOException, InterruptedException {
+        AsyncSinkWriterImpl sink = new AsyncSinkWriterImpl(sinkInitContext, 10, 1, 100,
+                (double) 30 / BYTES_IN_MB, true);
+
+        /* Sink has flush threshold of 30 bytes, each integer is 4 bytes, therefore, flushing
+         * should occur once 8 elements have been written.
+         */
+        for (int i = 0; i < 15; i++) {
+            sink.write(String.valueOf(i));
+        }
+        assertEquals(8, res.size());
+    }
+
     private class AsyncSinkWriterImpl extends AsyncSinkWriter<String, Integer> {
 
         private final Set<Integer> failedFirstAttempts = new HashSet<>();
         private final boolean simulateFailures;
 
         public AsyncSinkWriterImpl(
-                Sink.InitContext context,
-                int maxBatchSize,
-                int maxInFlightRequests,
-                int maxBufferedRequests,
-                boolean simulateFailures) {
+                Sink.InitContext context, int maxBatchSize, int maxInFlightRequests,
+                int maxBufferedRequests, boolean simulateFailures) {
             super(
                     (elem, ctx) -> Integer.parseInt(elem),
                     context,
                     maxBatchSize,
                     maxInFlightRequests,
-                    maxBufferedRequests);
+                    maxBufferedRequests, 100);
+            this.simulateFailures = simulateFailures;
+        }
+
+        public AsyncSinkWriterImpl(
+                Sink.InitContext context,
+                int maxBatchSize,
+                int maxInFlightRequests,
+                int maxBufferedRequests,
+                double flushOnBufferSizeMB,
+                boolean simulateFailures) {
+            super((elem, ctx) -> Integer.parseInt(elem), context, maxBatchSize, maxInFlightRequests,
+                    maxBufferedRequests, flushOnBufferSizeMB);
             this.simulateFailures = simulateFailures;
         }
 
@@ -305,6 +328,11 @@ public class AsyncSinkWriterTest {
                 requestResult.accept(new ArrayList<>());
             }
         }
+
+        @Override
+        protected int getSizeInBytes(Integer requestEntry) {
+            return 4;
+        }
     }
 
     private static class SinkInitContext implements Sink.InitContext {