CC-29764 Add timeout to sink put cycle

kafka-connect-s3/src/main/java/io/confluent/connect/s3/S3SinkConnectorConfig.java

@@ -162,6 +162,14 @@ public class S3SinkConnectorConfig extends StorageSinkConnectorConfig {
           + "`errors.tolerance` should be set to 'all' for successfully writing into dlq";
   public static final String REPORT_NULL_RECORDS_TO_DLQ_DISPLAY = "Report null value to dlq";
 
+  public static final String MAX_WRITE_DURATION = "max.write.duration.ms";
+  public static final long MAX_WRITE_DURATION_DEFAULT = Long.MAX_VALUE;
+  public static final String MAX_WRITE_DURATION_DOC = "The maximum duration that a task will "
+      + "spend in batching and writing to S3. If the write operation takes longer than this "
+      + "the task will voluntarily return from the put method. This prevents the consumer from "
+      + "being revoked from the group. It also mitigates (but does not eliminate) the risk of "
+      + "a zombie task to continue writing to S3 after it has been revoked.";
+
   /**
    * Maximum back-off time when retrying failed requests.
    */
@@ -631,6 +639,18 @@ public static ConfigDef newConfigDef() {
           Width.SHORT,
           REPORT_NULL_RECORDS_TO_DLQ_DISPLAY
       );
+
+      configDef.define(
+          MAX_WRITE_DURATION,
+          Type.LONG,
+          MAX_WRITE_DURATION_DEFAULT,
+          Importance.LOW,
+          MAX_WRITE_DURATION_DOC,
+          group,
+          ++orderInGroup,
+          Width.SHORT,
+          "Maximum write duration"
+      );
     }
 
     {

kafka-connect-s3/src/main/java/io/confluent/connect/s3/S3SinkTask.java

@@ -30,8 +30,11 @@
 import org.slf4j.LoggerFactory;
 
 import java.lang.reflect.InvocationTargetException;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import io.confluent.common.utils.SystemTime;
@@ -211,6 +214,8 @@ private Partitioner<?> newPartitioner(S3SinkConnectorConfig config)
 
   @Override
   public void put(Collection<SinkRecord> records) throws ConnectException {
+    long putStartTime = time.milliseconds();
+
     for (SinkRecord record : records) {
       String topic = record.topic();
       int partition = record.kafkaPartition();
@@ -228,9 +233,15 @@ public void put(Collection<SinkRecord> records) throws ConnectException {
       log.debug("Read {} records from Kafka", records.size());
     }
 
-    for (TopicPartition tp : topicPartitionWriters.keySet()) {
+    //shuffle the topic partitions as otherwise the last topic partition will
+    //always get the least amount of time to write
+    List<TopicPartition> shuffledList = new ArrayList<>(topicPartitionWriters.keySet());
+    Collections.shuffle(shuffledList);
+
+    for (TopicPartition tp : shuffledList) {
       TopicPartitionWriter writer = topicPartitionWriters.get(tp);
       try {
+        writer.setWriteDeadline(putStartTime);
         writer.write();
         if (log.isDebugEnabled()) {
           log.debug("TopicPartition: {}, SchemaCompatibility:{}, FileRotations: {}",

kafka-connect-s3/src/main/java/io/confluent/connect/s3/TopicPartitionWriter.java

@@ -100,6 +100,11 @@ public class TopicPartitionWriter {
   private static final Time SYSTEM_TIME = new SystemTime();
   private ErrantRecordReporter reporter;
 
+  private final long maxWriteDurationMs;
+  private long writeDeadline;
+
+  boolean isPaused = false;
+
   private final FileRotationTracker fileRotationTracker;
 
   public TopicPartitionWriter(TopicPartition tp,
@@ -173,6 +178,9 @@ public TopicPartitionWriter(TopicPartition tp,
         + "d";
     fileRotationTracker = new FileRotationTracker();
 
+    maxWriteDurationMs = connectorConfig.getLong(S3SinkConnectorConfig.MAX_WRITE_DURATION);
+    writeDeadline = Long.MAX_VALUE;
+
     // Initialize scheduled rotation timer if applicable
     setNextScheduledRotation();
   }
@@ -200,7 +208,8 @@ public void write() {
 
     resetExpiredScheduledRotationIfNoPendingRecords(now);
 
-    while (!buffer.isEmpty()) {
+
+    while (!buffer.isEmpty() && !isWriteDeadlineExceeded()) {
       try {
         executeState(now);
       } catch (IllegalWorkerStateException e) {
@@ -214,7 +223,39 @@ public void write() {
         }
       }
     }
-    commitOnTimeIfNoData(now);
+    if (!isWriteDeadlineExceeded()) {
+      commitOnTimeIfNoData(now);
+    }
+    pauseOrResumeOnBuffer();
+
+  }
+
+  private void pauseOrResumeOnBuffer() {
+    // if the deadline exceeds before all the records in buffer are processed, pause the writer
+    // if the buffered records are greater than flush size, this is to ensure that we don't keep
+    // getting messages from the consumer while we are still processing the buffer which can lead
+    // to memory issues
+    if (buffer.size() >= Math.max(flushSize, 1)) {
+      pause();
+    } else if (isPaused) {
+      resume();
+    }
+  }
+
+  public void setWriteDeadline(long currentTimeMs) {
+    writeDeadline = currentTimeMs + maxWriteDurationMs;
+    //prevent overflow
+    if (writeDeadline < 0) {
+      writeDeadline = Long.MAX_VALUE;
+    }
+  }
+
+  protected boolean isWriteDeadlineExceeded() {
+    boolean isWriteDeadlineExceeded = time.milliseconds() > writeDeadline;
+    if (isWriteDeadlineExceeded) {
+      log.info("Deadline exceeded");
+    }
+    return isWriteDeadlineExceeded;
   }
 
   @SuppressWarnings("fallthrough")
@@ -262,6 +303,11 @@ private void executeState(long now) {
         }
         // fallthrough
       case SHOULD_ROTATE:
+        if (isWriteDeadlineExceeded()) {
+          // note: this is a best-effort attempt to rotate the file before the deadline
+          // this check can pass and the deadline gets exceeded before the rotation is complete
+          break;
+        }
         commitFiles();
         nextState();
         // fallthrough
@@ -502,11 +548,13 @@ private boolean rotateOnSize() {
   private void pause() {
     log.trace("Pausing writer for topic-partition '{}'", tp);
     context.pause(tp);
+    isPaused = true;
   }
 
   private void resume() {
     log.trace("Resuming writer for topic-partition '{}'", tp);
     context.resume(tp);
+    isPaused = false;
   }
 
   private RecordWriter newWriter(SinkRecord record, String encodedPartition)

kafka-connect-s3/src/test/java/io/confluent/connect/s3/TopicPartitionWriterTest.java

@@ -696,6 +696,9 @@ public void testWallclockUsesBatchTimePartitionBoundary() throws Exception {
     // Freeze clock passed into topicPartitionWriter, so we know what time it will use for "now"
     long freezeTime = 3599000L;
     EasyMock.expect(systemTime.milliseconds()).andReturn(freezeTime);
+
+    // Mock system time for deadline checks
+    EasyMock.expect(systemTime.milliseconds()).andReturn(System.currentTimeMillis()).anyTimes();
     EasyMock.replay(systemTime);
 
     String key = "key";