a simple back pressure mechanism

Mike Skells · Mike Skells · commit de087199adcf · 2024-11-14T14:06:16.000Z
diff --git a/azure-sink-connector/src/main/java/io/aiven/kafka/connect/azure/sink/AzureBlobSinkConfig.java b/azure-sink-connector/src/main/java/io/aiven/kafka/connect/azure/sink/AzureBlobSinkConfig.java
@@ -80,6 +80,7 @@ public static ConfigDef configDef() {
         addKafkaBackoffPolicy(configDef);
         addAzureRetryPolicies(configDef);
         addUserAgentConfig(configDef);
+        addOverloadConfigGroup(configDef);
         return configDef;
     }
 
diff --git a/azure-sink-connector/src/main/java/io/aiven/kafka/connect/azure/sink/AzureBlobSinkTask.java b/azure-sink-connector/src/main/java/io/aiven/kafka/connect/azure/sink/AzureBlobSinkTask.java
@@ -57,6 +57,9 @@ public final class AzureBlobSinkTask extends SinkTask {
     private BlobContainerClient containerClient;
     private final Map<String, BlockBlobClient> blobClientMap = new ConcurrentHashMap<>();
 
+    private long backPressureHardLimit;
+    private long backPressureCurrentBuffer;
+
     // required by Connect
     public AzureBlobSinkTask() {
         super();
@@ -104,6 +107,7 @@ public void start(final Map<String, String> props) {
         if (Objects.nonNull(config.getKafkaRetryBackoffMs())) {
             context.timeout(config.getKafkaRetryBackoffMs());
         }
+        backPressureHardLimit = config.getBackPressureHardLimit();
     }
 
     private void initRecordGrouper() {
@@ -122,6 +126,11 @@ public void put(final Collection<SinkRecord> records) {
         for (final SinkRecord record : records) {
             recordGrouper.put(record);
         }
+        backPressureCurrentBuffer += records.size();
+        if (backPressureCurrentBuffer >= backPressureHardLimit) {
+            LOG.warn("Back pressure limit reached, requesting flush");
+            context.requestCommit();
+        }
     }
 
     @Override
@@ -131,6 +140,7 @@ public void flush(final Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
         } finally {
             recordGrouper.clear();
         }
+        backPressureCurrentBuffer = 0;
     }
 
     private void flushFile(final String filename, final List<SinkRecord> records) {
diff --git a/commons/src/main/java/io/aiven/kafka/connect/common/config/AivenCommonConfig.java b/commons/src/main/java/io/aiven/kafka/connect/common/config/AivenCommonConfig.java
@@ -53,6 +53,9 @@ public class AivenCommonConfig extends AbstractConfig {
     private static final String GROUP_RETRY_BACKOFF_POLICY = "Retry backoff policy";
     public static final String KAFKA_RETRY_BACKOFF_MS_CONFIG = "kafka.retry.backoff.ms";
 
+    private static final String GROUP_OVERLOAD = "Overload Control";
+    private static final String OVERLOAD_MAX_RECORDS_HARD_LIMIT = "overload.hard.record.limit";
+
     protected AivenCommonConfig(final ConfigDef definition, final Map<?, ?> originals) {
         super(definition, originals);
         // TODO: calls getOutputFields, can be overridden in subclasses.
@@ -74,23 +77,23 @@ private void validate() {
     protected static void addKafkaBackoffPolicy(final ConfigDef configDef) {
         configDef.define(KAFKA_RETRY_BACKOFF_MS_CONFIG, ConfigDef.Type.LONG, null, new ConfigDef.Validator() {
 
-            final long maximumBackoffPolicy = TimeUnit.HOURS.toMillis(24);
-
-            @Override
-            public void ensureValid(final String name, final Object value) {
-                if (Objects.isNull(value)) {
-                    return;
-                }
-                assert value instanceof Long;
-                final var longValue = (Long) value;
-                if (longValue < 0) {
-                    throw new ConfigException(name, value, "Value must be at least 0");
-                } else if (longValue > maximumBackoffPolicy) {
-                    throw new ConfigException(name, value,
-                            "Value must be no more than " + maximumBackoffPolicy + " (24 hours)");
-                }
-            }
-        }, ConfigDef.Importance.MEDIUM,
+                    final long maximumBackoffPolicy = TimeUnit.HOURS.toMillis(24);
+
+                    @Override
+                    public void ensureValid(final String name, final Object value) {
+                        if (Objects.isNull(value)) {
+                            return;
+                        }
+                        assert value instanceof Long;
+                        final var longValue = (Long) value;
+                        if (longValue < 0) {
+                            throw new ConfigException(name, value, "Value must be at least 0");
+                        } else if (longValue > maximumBackoffPolicy) {
+                            throw new ConfigException(name, value,
+                                    "Value must be no more than " + maximumBackoffPolicy + " (24 hours)");
+                        }
+                    }
+                }, ConfigDef.Importance.MEDIUM,
                 "The retry backoff in milliseconds. "
                         + "This config is used to notify Kafka Connect to retry delivering a message batch or "
                         + "performing recovery in case of transient exceptions. Maximum value is "
@@ -103,7 +106,7 @@ public Long getKafkaRetryBackoffMs() {
     }
 
     protected static void addOutputFieldsFormatConfigGroup(final ConfigDef configDef,
-            final OutputFieldType defaultFieldType) {
+                                                           final OutputFieldType defaultFieldType) {
         int formatGroupCounter = 0;
 
         addFormatTypeConfig(configDef, formatGroupCounter);
@@ -145,7 +148,7 @@ public FormatType getFormatType() {
     }
 
     protected static void addCompressionTypeConfig(final ConfigDef configDef,
-            final CompressionType defaultCompressionType) {
+                                                   final CompressionType defaultCompressionType) {
         configDef.define(FILE_COMPRESSION_TYPE_CONFIG, ConfigDef.Type.STRING,
                 Objects.isNull(defaultCompressionType) ? null : defaultCompressionType.name, // NOPMD NullAssignment
                 new FileCompressionTypeValidator(), ConfigDef.Importance.MEDIUM,
@@ -229,4 +232,16 @@ private Boolean isKeyBased(final String groupType) {
         return RecordGrouperFactory.KEY_RECORD.equals(groupType)
                 || RecordGrouperFactory.KEY_TOPIC_PARTITION_RECORD.equals(groupType);
     }
+
+    public long getBackPressureHardLimit() {
+        return getLong(OVERLOAD_MAX_RECORDS_HARD_LIMIT);
+    }
+
+    protected static void addOverloadConfigGroup(final ConfigDef configDef) {
+        int groupCounter = 0;
+
+        configDef.define(OVERLOAD_MAX_RECORDS_HARD_LIMIT, ConfigDef.Type.LONG, 1000000L, ConfigDef.Importance.MEDIUM,
+                "The maximum number of records to buffer before requesting a flush.", GROUP_OVERLOAD, groupCounter++,
+                ConfigDef.Width.NONE, OVERLOAD_MAX_RECORDS_HARD_LIMIT);
+    }
 }
diff --git a/gcs-sink-connector/src/main/java/io/aiven/kafka/connect/gcs/GcsSinkConfig.java b/gcs-sink-connector/src/main/java/io/aiven/kafka/connect/gcs/GcsSinkConfig.java
@@ -107,6 +107,7 @@ public static ConfigDef configDef() {
         addKafkaBackoffPolicy(configDef);
         addGcsRetryPolicies(configDef);
         addUserAgentConfig(configDef);
+        addOverloadConfigGroup(configDef);
         return configDef;
     }
 
diff --git a/gcs-sink-connector/src/main/java/io/aiven/kafka/connect/gcs/GcsSinkTask.java b/gcs-sink-connector/src/main/java/io/aiven/kafka/connect/gcs/GcsSinkTask.java
@@ -49,6 +49,8 @@ public final class GcsSinkTask extends SinkTask {
     private GcsSinkConfig config;
 
     private Storage storage;
+    private long backPressureHardLimit;
+    private long backPressureCurrentBuffer;
 
     // required by Connect
     public GcsSinkTask() {
@@ -88,6 +90,7 @@ public void start(final Map<String, String> props) {
         if (Objects.nonNull(config.getKafkaRetryBackoffMs())) {
             context.timeout(config.getKafkaRetryBackoffMs());
         }
+        backPressureHardLimit = config.getBackPressureHardLimit();
     }
 
     private void initRest() {
@@ -106,6 +109,11 @@ public void put(final Collection<SinkRecord> records) {
         for (final SinkRecord record : records) {
             recordGrouper.put(record);
         }
+        backPressureCurrentBuffer += records.size();
+        if (backPressureCurrentBuffer >= backPressureHardLimit) {
+            LOG.warn("Back pressure limit reached, requesting flush");
+            context.requestCommit();
+        }
     }
 
     @Override
@@ -115,6 +123,7 @@ public void flush(final Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
         } finally {
             recordGrouper.clear();
         }
+        backPressureCurrentBuffer = 0;
     }
 
     private void flushFile(final String filename, final List<SinkRecord> records) {
diff --git a/s3-sink-connector/src/main/java/io/aiven/kafka/connect/s3/S3SinkTask.java b/s3-sink-connector/src/main/java/io/aiven/kafka/connect/s3/S3SinkTask.java
@@ -66,6 +66,9 @@ public final class S3SinkTask extends SinkTask {
 
     AwsCredentialProviderFactory credentialFactory = new AwsCredentialProviderFactory();
 
+    private long backPressureHardLimit;
+    private long backPressureCurrentBuffer;
+
     @SuppressWarnings("PMD.UnnecessaryConstructor") // required by Connect
     public S3SinkTask() {
         super();
@@ -84,6 +87,7 @@ public void start(final Map<String, String> props) {
         if (Objects.nonNull(config.getKafkaRetryBackoffMs())) {
             context.timeout(config.getKafkaRetryBackoffMs());
         }
+        backPressureHardLimit = config.getBackPressureHardLimit();
     }
 
     private AmazonS3 createAmazonS3Client(final S3SinkConfig config) {
@@ -110,6 +114,11 @@ public void put(final Collection<SinkRecord> records) {
         Objects.requireNonNull(records, "records cannot be null");
         LOGGER.info("Processing {} records", records.size());
         records.forEach(recordGrouper::put);
+        backPressureCurrentBuffer += records.size();
+        if (backPressureCurrentBuffer >= backPressureHardLimit) {
+            LOGGER.warn("Back pressure limit reached, requesting flush");
+            context.requestCommit();
+        }
     }
 
     @Override
@@ -119,6 +128,7 @@ public void flush(final Map<TopicPartition, OffsetAndMetadata> offsets) {
         } finally {
             recordGrouper.clear();
         }
+        backPressureCurrentBuffer = 0;
     }
 
     private void flushFile(final String filename, final List<SinkRecord> records) {
diff --git a/s3-sink-connector/src/main/java/io/aiven/kafka/connect/s3/config/S3SinkConfig.java b/s3-sink-connector/src/main/java/io/aiven/kafka/connect/s3/config/S3SinkConfig.java
@@ -189,6 +189,7 @@ public static ConfigDef configDef() {
         addDeprecatedConfiguration(configDef);
         addKafkaBackoffPolicy(configDef);
         addS3RetryPolicies(configDef);
+        addOverloadConfigGroup(configDef);
         return configDef;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -80,6 +80,7 @@ public static ConfigDef configDef() {`
`80`	`80`	`addKafkaBackoffPolicy(configDef);`
`81`	`81`	`addAzureRetryPolicies(configDef);`
`82`	`82`	`addUserAgentConfig(configDef);`
	`83`	`+ addOverloadConfigGroup(configDef);`
`83`	`84`	`return configDef;`
`84`	`85`	`}`
`85`	`86`
Original file line number	Diff line number	Diff line change
`@@ -57,6 +57,9 @@ public final class AzureBlobSinkTask extends SinkTask {`
`57`	`57`	`private BlobContainerClient containerClient;`
`58`	`58`	`private final Map<String, BlockBlobClient> blobClientMap = new ConcurrentHashMap<>();`
`59`	`59`
	`60`	`+ private long backPressureHardLimit;`
	`61`	`+ private long backPressureCurrentBuffer;`
	`62`	`+`
`60`	`63`	`// required by Connect`
`61`	`64`	`public AzureBlobSinkTask() {`
`62`	`65`	`super();`
`@@ -104,6 +107,7 @@ public void start(final Map<String, String> props) {`
`104`	`107`	`if (Objects.nonNull(config.getKafkaRetryBackoffMs())) {`
`105`	`108`	`context.timeout(config.getKafkaRetryBackoffMs());`
`106`	`109`	`}`
	`110`	`+ backPressureHardLimit = config.getBackPressureHardLimit();`
`107`	`111`	`}`
`108`	`112`
`109`	`113`	`private void initRecordGrouper() {`
`@@ -122,6 +126,11 @@ public void put(final Collection<SinkRecord> records) {`
`122`	`126`	`for (final SinkRecord record : records) {`
`123`	`127`	`recordGrouper.put(record);`
`124`	`128`	`}`
	`129`	`+ backPressureCurrentBuffer += records.size();`
	`130`	`+ if (backPressureCurrentBuffer >= backPressureHardLimit) {`
	`131`	`+ LOG.warn("Back pressure limit reached, requesting flush");`
	`132`	`+ context.requestCommit();`
	`133`	`+ }`
`125`	`134`	`}`
`126`	`135`
`127`	`136`	`@Override`
`@@ -131,6 +140,7 @@ public void flush(final Map<TopicPartition, OffsetAndMetadata> currentOffsets) {`
`131`	`140`	`} finally {`
`132`	`141`	`recordGrouper.clear();`
`133`	`142`	`}`
	`143`	`+ backPressureCurrentBuffer = 0;`
`134`	`144`	`}`
`135`	`145`
`136`	`146`	`private void flushFile(final String filename, final List<SinkRecord> records) {`
Original file line number	Diff line number	Diff line change
`@@ -107,6 +107,7 @@ public static ConfigDef configDef() {`
`107`	`107`	`addKafkaBackoffPolicy(configDef);`
`108`	`108`	`addGcsRetryPolicies(configDef);`
`109`	`109`	`addUserAgentConfig(configDef);`
	`110`	`+ addOverloadConfigGroup(configDef);`
`110`	`111`	`return configDef;`
`111`	`112`	`}`
`112`	`113`
Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,8 @@ public final class GcsSinkTask extends SinkTask {`
`49`	`49`	`private GcsSinkConfig config;`
`50`	`50`
`51`	`51`	`private Storage storage;`
	`52`	`+ private long backPressureHardLimit;`
	`53`	`+ private long backPressureCurrentBuffer;`
`52`	`54`
`53`	`55`	`// required by Connect`
`54`	`56`	`public GcsSinkTask() {`
`@@ -88,6 +90,7 @@ public void start(final Map<String, String> props) {`
`88`	`90`	`if (Objects.nonNull(config.getKafkaRetryBackoffMs())) {`
`89`	`91`	`context.timeout(config.getKafkaRetryBackoffMs());`
`90`	`92`	`}`
	`93`	`+ backPressureHardLimit = config.getBackPressureHardLimit();`
`91`	`94`	`}`
`92`	`95`
`93`	`96`	`private void initRest() {`
`@@ -106,6 +109,11 @@ public void put(final Collection<SinkRecord> records) {`
`106`	`109`	`for (final SinkRecord record : records) {`
`107`	`110`	`recordGrouper.put(record);`
`108`	`111`	`}`
	`112`	`+ backPressureCurrentBuffer += records.size();`
	`113`	`+ if (backPressureCurrentBuffer >= backPressureHardLimit) {`
	`114`	`+ LOG.warn("Back pressure limit reached, requesting flush");`
	`115`	`+ context.requestCommit();`
	`116`	`+ }`
`109`	`117`	`}`
`110`	`118`
`111`	`119`	`@Override`
`@@ -115,6 +123,7 @@ public void flush(final Map<TopicPartition, OffsetAndMetadata> currentOffsets) {`
`115`	`123`	`} finally {`
`116`	`124`	`recordGrouper.clear();`
`117`	`125`	`}`
	`126`	`+ backPressureCurrentBuffer = 0;`
`118`	`127`	`}`
`119`	`128`
`120`	`129`	`private void flushFile(final String filename, final List<SinkRecord> records) {`
Original file line number	Diff line number	Diff line change
`@@ -189,6 +189,7 @@ public static ConfigDef configDef() {`
`189`	`189`	`addDeprecatedConfiguration(configDef);`
`190`	`190`	`addKafkaBackoffPolicy(configDef);`
`191`	`191`	`addS3RetryPolicies(configDef);`
	`192`	`+ addOverloadConfigGroup(configDef);`
`192`	`193`	`return configDef;`
`193`	`194`	`}`
`194`	`195`