Kinesis adaptive memory management

extensions-contrib/kafka-emitter/src/main/java/org/apache/druid/emitter/kafka/KafkaEmitter.java

@@ -23,7 +23,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.druid.emitter.kafka.KafkaEmitterConfig.EventType;
-import org.apache.druid.emitter.kafka.MemoryBoundLinkedBlockingQueue.ObjectContainer;
+import org.apache.druid.java.util.common.MemoryBoundLinkedBlockingQueue;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
 import org.apache.druid.java.util.common.logger.Logger;
@@ -173,7 +173,7 @@ private void sendSegmentMetadataToKafka()
 
   private void sendToKafka(final String topic, MemoryBoundLinkedBlockingQueue<String> recordQueue, Callback callback)
   {
-    ObjectContainer<String> objectToSend;
+    MemoryBoundLinkedBlockingQueue.ObjectContainer<String> objectToSend;
     try {
       while (true) {
         objectToSend = recordQueue.take();
@@ -199,7 +199,7 @@ public void emit(final Event event)
 
         String resultJson = jsonMapper.writeValueAsString(map);
 
-        ObjectContainer<String> objectContainer = new ObjectContainer<>(
+        MemoryBoundLinkedBlockingQueue.ObjectContainer<String> objectContainer = new MemoryBoundLinkedBlockingQueue.ObjectContainer<>(
             resultJson,
             StringUtils.toUtf8(resultJson).length
         );

extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTask.java

@@ -49,6 +49,10 @@
 public class KinesisIndexTask extends SeekableStreamIndexTask<String, String, ByteEntity>
 {
   private static final String TYPE = "index_kinesis";
+
+  // GetRecords returns maximum 10MB per call
+  // (https://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html)
+  private static final long GET_RECORDS_MAX_BYTES_PER_CALL = 10_000_000L;
   private static final Logger log = new Logger(KinesisIndexTask.class);
 
   private final boolean useListShards;
@@ -78,6 +82,14 @@ public KinesisIndexTask(
     );
     this.useListShards = useListShards;
     this.awsCredentialsConfig = awsCredentialsConfig;
+    if (tuningConfig.getRecordBufferSizeConfigured() != null) {
+      log.warn("The 'recordBufferSize' config property of the kinesis tuning config has been deprecated. "
+               + "Please use 'recordBufferSizeBytes'.");
+    }
+    if (tuningConfig.getMaxRecordsPerPollConfigured() != null) {
+      log.warn("The 'maxRecordsPerPoll' config property of the kinesis tuning config has been deprecated. "
+               + "Please use 'maxBytesPerPoll'.");
+    }
   }
 
   @Override
@@ -105,21 +117,18 @@ protected KinesisRecordSupplier newTaskRecordSupplier(final TaskToolbox toolbox)
   {
     KinesisIndexTaskIOConfig ioConfig = ((KinesisIndexTaskIOConfig) super.ioConfig);
     KinesisIndexTaskTuningConfig tuningConfig = ((KinesisIndexTaskTuningConfig) super.tuningConfig);
+    final int recordBufferSizeBytes =
+        tuningConfig.getRecordBufferSizeBytesOrDefault(runtimeInfo.getMaxHeapSizeBytes());
     final int fetchThreads = computeFetchThreads(runtimeInfo, tuningConfig.getFetchThreads());
-    final int recordsPerFetch = ioConfig.getRecordsPerFetchOrDefault(runtimeInfo.getMaxHeapSizeBytes(), fetchThreads);
-    final int recordBufferSize =
-        tuningConfig.getRecordBufferSizeOrDefault(runtimeInfo.getMaxHeapSizeBytes(), ioConfig.isDeaggregate());
-    final int maxRecordsPerPoll = tuningConfig.getMaxRecordsPerPollOrDefault(ioConfig.isDeaggregate());
+    final int maxBytesPerPoll = tuningConfig.getMaxBytesPerPollOrDefault();
 
     log.info(
-        "Starting record supplier with fetchThreads [%d], fetchDelayMillis [%d], recordsPerFetch [%d], "
-        + "recordBufferSize [%d], maxRecordsPerPoll [%d], deaggregate [%s].",
+        "Starting record supplier with fetchThreads [%d], fetchDelayMillis [%d], "
+        + "recordBufferSizeBytes [%d], maxBytesPerPoll [%d]",
         fetchThreads,
         ioConfig.getFetchDelayMillis(),
-        recordsPerFetch,
-        recordBufferSize,
-        maxRecordsPerPoll,
-        ioConfig.isDeaggregate()
+        recordBufferSizeBytes,
+        maxBytesPerPoll
     );
 
     return new KinesisRecordSupplier(
@@ -129,14 +138,12 @@ protected KinesisRecordSupplier newTaskRecordSupplier(final TaskToolbox toolbox)
             ioConfig.getAwsAssumedRoleArn(),
             ioConfig.getAwsExternalId()
         ),
-        recordsPerFetch,
         ioConfig.getFetchDelayMillis(),
         fetchThreads,
-        ioConfig.isDeaggregate(),
-        recordBufferSize,
+        recordBufferSizeBytes,
         tuningConfig.getRecordBufferOfferTimeout(),
         tuningConfig.getRecordBufferFullWait(),
-        maxRecordsPerPoll,
+        maxBytesPerPoll,
         false,
         useListShards
     );
@@ -179,15 +186,36 @@ AWSCredentialsConfig getAwsCredentialsConfig()
   }
 
   @VisibleForTesting
-  static int computeFetchThreads(final RuntimeInfo runtimeInfo, final Integer configuredFetchThreads)
+  static int computeFetchThreads(
+      final RuntimeInfo runtimeInfo,
+      final Integer configuredFetchThreads
+  )
   {
-    final int fetchThreads;
+    int fetchThreads;
     if (configuredFetchThreads != null) {
       fetchThreads = configuredFetchThreads;
     } else {
       fetchThreads = runtimeInfo.getAvailableProcessors() * 2;
     }
 
+    // Each fetchThread can return upto 10MB at a time
+    // (https://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html), cap fetchThreads so that
+    // we don't exceed more than the least of 100MB or 5% of heap at a time. Don't fail if fetchThreads specified
+    // is greater than this as to not cause failure for older configurations, but log warning in this case, and lower
+    // fetchThreads implicitly.
+    final long memoryToUse = Math.min(
+        KinesisIndexTaskIOConfig.MAX_RECORD_FETCH_MEMORY,
+        (long) (runtimeInfo.getMaxHeapSizeBytes() * KinesisIndexTaskIOConfig.RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION)
+    );
+    int maxFetchThreads = Math.max(
+        1,
+        (int) (memoryToUse / GET_RECORDS_MAX_BYTES_PER_CALL)
+    );
+    if (fetchThreads > maxFetchThreads) {
+      log.warn("fetchThreads [%d] being lowered to [%d]", fetchThreads, maxFetchThreads);
+      fetchThreads = maxFetchThreads;
+    }
+
     Preconditions.checkArgument(
         fetchThreads > 0,
         "Must have at least one background fetch thread for the record supplier"

extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskIOConfig.java

@@ -23,7 +23,6 @@
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
-import com.google.common.primitives.Ints;
 import org.apache.druid.data.input.InputFormat;
 import org.apache.druid.indexing.seekablestream.SeekableStreamEndSequenceNumbers;
 import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig;
@@ -41,21 +40,19 @@ public class KinesisIndexTaskIOConfig extends SeekableStreamIndexTaskIOConfig<St
    * Together with {@link KinesisIndexTaskTuningConfig#MAX_RECORD_BUFFER_MEMORY}, don't take up more than 200MB
    * per task.
    */
-  private static final int MAX_RECORD_FETCH_MEMORY = 100_000_000;
+  public static final int MAX_RECORD_FETCH_MEMORY = 100_000_000;
 
   /**
    * Together with {@link KinesisIndexTaskTuningConfig#RECORD_BUFFER_MEMORY_MAX_HEAP_FRACTION}, don't take up more
    * than 15% of the heap.
    */
-  private static final double RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION = 0.05;
+  public static final double RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION = 0.05;
 
   private final String endpoint;
-  private final Integer recordsPerFetch;
   private final int fetchDelayMillis;
 
   private final String awsAssumedRoleArn;
   private final String awsExternalId;
-  private final boolean deaggregate;
 
   @JsonCreator
   public KinesisIndexTaskIOConfig(
@@ -79,11 +76,9 @@ public KinesisIndexTaskIOConfig(
       @JsonProperty("maximumMessageTime") DateTime maximumMessageTime,
       @JsonProperty("inputFormat") @Nullable InputFormat inputFormat,
       @JsonProperty("endpoint") String endpoint,
-      @JsonProperty("recordsPerFetch") Integer recordsPerFetch,
       @JsonProperty("fetchDelayMillis") Integer fetchDelayMillis,
       @JsonProperty("awsAssumedRoleArn") String awsAssumedRoleArn,
-      @JsonProperty("awsExternalId") String awsExternalId,
-      @JsonProperty("deaggregate") boolean deaggregate
+      @JsonProperty("awsExternalId") String awsExternalId
   )
   {
     super(
@@ -105,11 +100,9 @@ public KinesisIndexTaskIOConfig(
     );
 
     this.endpoint = Preconditions.checkNotNull(endpoint, "endpoint");
-    this.recordsPerFetch = recordsPerFetch;
     this.fetchDelayMillis = fetchDelayMillis != null ? fetchDelayMillis : DEFAULT_FETCH_DELAY_MILLIS;
     this.awsAssumedRoleArn = awsAssumedRoleArn;
     this.awsExternalId = awsExternalId;
-    this.deaggregate = deaggregate;
   }
 
   public KinesisIndexTaskIOConfig(
@@ -122,11 +115,9 @@ public KinesisIndexTaskIOConfig(
       DateTime maximumMessageTime,
       InputFormat inputFormat,
       String endpoint,
-      Integer recordsPerFetch,
       Integer fetchDelayMillis,
       String awsAssumedRoleArn,
-      String awsExternalId,
-      boolean deaggregate
+      String awsExternalId
   )
   {
     this(
@@ -142,11 +133,9 @@ public KinesisIndexTaskIOConfig(
         maximumMessageTime,
         inputFormat,
         endpoint,
-        recordsPerFetch,
         fetchDelayMillis,
         awsAssumedRoleArn,
-        awsExternalId,
-        deaggregate
+        awsExternalId
     );
   }
 
@@ -215,32 +204,6 @@ public String getEndpoint()
     return endpoint;
   }
 
-  @Nullable
-  @JsonProperty("recordsPerFetch")
-  @JsonInclude(JsonInclude.Include.NON_NULL)
-  public Integer getRecordsPerFetchConfigured()
-  {
-    return recordsPerFetch;
-  }
-
-  public int getRecordsPerFetchOrDefault(final long maxHeapSize, final int fetchThreads)
-  {
-    if (recordsPerFetch != null) {
-      return recordsPerFetch;
-    } else {
-      final long memoryToUse = Math.min(
-          MAX_RECORD_FETCH_MEMORY,
-          (long) (maxHeapSize * RECORD_FETCH_MEMORY_MAX_HEAP_FRACTION)
-      );
-
-      final int assumedRecordSize = deaggregate
-                                    ? KinesisIndexTaskTuningConfig.ASSUMED_RECORD_SIZE_AGGREGATE
-                                    : KinesisIndexTaskTuningConfig.ASSUMED_RECORD_SIZE;
-
-      return Ints.checkedCast(Math.max(1, memoryToUse / assumedRecordSize / fetchThreads));
-    }
-  }
-
   @JsonProperty
   @JsonInclude(JsonInclude.Include.NON_DEFAULT)
   public int getFetchDelayMillis()
@@ -262,13 +225,6 @@ public String getAwsExternalId()
     return awsExternalId;
   }
 
-  @JsonProperty
-  @JsonInclude(JsonInclude.Include.NON_DEFAULT)
-  public boolean isDeaggregate()
-  {
-    return deaggregate;
-  }
-
   @Override
   public String toString()
   {
@@ -280,11 +236,9 @@ public String toString()
            ", minimumMessageTime=" + getMinimumMessageTime() +
            ", maximumMessageTime=" + getMaximumMessageTime() +
            ", endpoint='" + endpoint + '\'' +
-           ", recordsPerFetch=" + recordsPerFetch +
            ", fetchDelayMillis=" + fetchDelayMillis +
            ", awsAssumedRoleArn='" + awsAssumedRoleArn + '\'' +
            ", awsExternalId='" + awsExternalId + '\'' +
-           ", deaggregate=" + deaggregate +
            '}';
   }
 }