supervisor: make rejection periods work with stopTasksCount

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskIOConfig.java

@@ -28,6 +28,7 @@
 import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig;
 import org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 import javax.annotation.Nullable;
 import java.util.Map;
@@ -53,7 +54,9 @@ public RabbitStreamIndexTaskIOConfig(
       @JsonProperty("minimumMessageTime") DateTime minimumMessageTime,
       @JsonProperty("maximumMessageTime") DateTime maximumMessageTime,
       @JsonProperty("inputFormat") @Nullable InputFormat inputFormat,
-      @JsonProperty("uri") String uri)
+      @JsonProperty("uri") String uri,
+      @JsonProperty("taskDuration") Duration taskDuration
+  )
   {
     super(
         taskGroupId,
@@ -63,7 +66,9 @@ public RabbitStreamIndexTaskIOConfig(
         useTransaction,
         minimumMessageTime,
         maximumMessageTime,
-        inputFormat);
+        inputFormat,
+        taskDuration
+    );
 
     this.pollTimeout = pollTimeout != null ? pollTimeout : RabbitStreamSupervisorIOConfig.DEFAULT_POLL_TIMEOUT_MILLIS;
     this.uri = uri;

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/supervisor/RabbitStreamSupervisor.java

@@ -59,7 +59,6 @@
 import org.joda.time.DateTime;
 
 import javax.annotation.Nullable;
-
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -202,7 +201,9 @@ protected SeekableStreamIndexTaskIOConfig createTaskIoConfig(
         minimumMessageTime,
         maximumMessageTime,
         ioConfig.getInputFormat(),
-        rabbitConfig.getUri());
+        rabbitConfig.getUri(),
+        ioConfig.getTaskDuration()
+    );
   }
 
   @Override

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaConsumerMonitor.java

@@ -43,16 +43,20 @@ public class KafkaConsumerMonitor extends AbstractMonitor
       ImmutableMap.<String, String>builder()
                   .put("bytes-consumed-total", "kafka/consumer/bytesConsumed")
                   .put("records-consumed-total", "kafka/consumer/recordsConsumed")
+                  .put("io-wait-time-ns-total", "kafka/consumer/io/time")
                   .build();
   private static final String TOPIC_TAG = "topic";
+  private static final String DATASOURCE_TAG = "datasource";
   private static final Set<String> TOPIC_METRIC_TAGS = ImmutableSet.of("client-id", TOPIC_TAG);
 
   private final KafkaConsumer<?, ?> consumer;
+  private final String datasource;
   private final Map<String, AtomicLong> counters = new HashMap<>();
 
-  public KafkaConsumerMonitor(final KafkaConsumer<?, ?> consumer)
+  public KafkaConsumerMonitor(final KafkaConsumer<?, ?> consumer, String datasource)
   {
     this.consumer = consumer;
+    this.datasource = datasource;
   }
 
   @Override
@@ -71,6 +75,9 @@ public boolean doMonitor(final ServiceEmitter emitter)
         if (newValue != priorValue) {
           final ServiceMetricEvent.Builder builder =
               new ServiceMetricEvent.Builder().setDimension(TOPIC_TAG, topic);
+          if (datasource != null) {
+            builder.setDimension(DATASOURCE_TAG, datasource);
+          }
           emitter.emit(builder.setMetric(METRICS.get(metricName.name()), newValue - priorValue));
         }
       }

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java

@@ -111,7 +111,8 @@ protected KafkaRecordSupplier newTaskRecordSupplier(final TaskToolbox toolbox)
 
       final KafkaRecordSupplier recordSupplier =
           new KafkaRecordSupplier(props, configMapper, kafkaIndexTaskIOConfig.getConfigOverrides(),
-                                  kafkaIndexTaskIOConfig.isMultiTopic());
+                                  kafkaIndexTaskIOConfig.isMultiTopic(), dataSchema.getDataSource()
+          );
 
       if (toolbox.getMonitorScheduler() != null) {
         toolbox.getMonitorScheduler().addMonitor(recordSupplier.monitor());

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskIOConfig.java

@@ -30,6 +30,7 @@
 import org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers;
 import org.apache.druid.indexing.seekablestream.extension.KafkaConfigOverrides;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 import javax.annotation.Nullable;
 import java.util.Map;
@@ -63,7 +64,8 @@ public KafkaIndexTaskIOConfig(
       @JsonProperty("maximumMessageTime") DateTime maximumMessageTime,
       @JsonProperty("inputFormat") @Nullable InputFormat inputFormat,
       @JsonProperty("configOverrides") @Nullable KafkaConfigOverrides configOverrides,
-      @JsonProperty("multiTopic") @Nullable Boolean multiTopic
+      @JsonProperty("multiTopic") @Nullable Boolean multiTopic,
+      @JsonProperty("taskDuration") Duration taskDuration
   )
   {
     super(
@@ -76,7 +78,8 @@ public KafkaIndexTaskIOConfig(
         useTransaction,
         minimumMessageTime,
         maximumMessageTime,
-        inputFormat
+        inputFormat,
+        taskDuration
     );
 
     this.consumerProperties = Preconditions.checkNotNull(consumerProperties, "consumerProperties");
@@ -107,7 +110,8 @@ public KafkaIndexTaskIOConfig(
       DateTime minimumMessageTime,
       DateTime maximumMessageTime,
       InputFormat inputFormat,
-      KafkaConfigOverrides configOverrides
+      KafkaConfigOverrides configOverrides,
+      Duration taskDuration
   )
   {
     this(
@@ -124,7 +128,8 @@ public KafkaIndexTaskIOConfig(
         maximumMessageTime,
         inputFormat,
         configOverrides,
-        KafkaSupervisorIOConfig.DEFAULT_IS_MULTI_TOPIC
+        KafkaSupervisorIOConfig.DEFAULT_IS_MULTI_TOPIC,
+        taskDuration
     );
   }
 

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java

@@ -45,6 +45,7 @@
 import org.apache.kafka.common.serialization.Deserializer;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Type;
@@ -78,21 +79,23 @@ public KafkaRecordSupplier(
       Map<String, Object> consumerProperties,
       ObjectMapper sortingMapper,
       KafkaConfigOverrides configOverrides,
-      boolean multiTopic
+      boolean multiTopic,
+      @Nullable String datasource
   )
   {
-    this(getKafkaConsumer(sortingMapper, consumerProperties, configOverrides), multiTopic);
+    this(getKafkaConsumer(sortingMapper, consumerProperties, configOverrides), multiTopic, datasource);
   }
 
   @VisibleForTesting
   public KafkaRecordSupplier(
       KafkaConsumer<byte[], byte[]> consumer,
-      boolean multiTopic
+      boolean multiTopic,
+      @Nullable String datasource
   )
   {
     this.consumer = consumer;
     this.multiTopic = multiTopic;
-    this.monitor = new KafkaConsumerMonitor(consumer);
+    this.monitor = new KafkaConsumerMonitor(consumer, datasource);
   }
 
   @Override

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java

@@ -74,7 +74,7 @@ protected KafkaRecordSupplier createRecordSupplier()
       KafkaSupervisorIOConfig kafkaSupervisorIOConfig = (KafkaSupervisorIOConfig) ioConfig;
 
       return new KafkaRecordSupplier(props, objectMapper, kafkaSupervisorIOConfig.getConfigOverrides(),
-                                     kafkaSupervisorIOConfig.isMultiTopic()
+                                     kafkaSupervisorIOConfig.isMultiTopic(), null
       );
     }
     finally {

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java

@@ -137,7 +137,8 @@
         spec.getIoConfig().getConsumerProperties(),
         sortingMapper,
         spec.getIoConfig().getConfigOverrides(),
-        spec.getIoConfig().isMultiTopic()
+        spec.getIoConfig().isMultiTopic(),
+        spec.getId()
     );
   }
 
@@ -218,7 +219,8 @@
         maximumMessageTime,
         ioConfig.getInputFormat(),
         kafkaIoConfig.getConfigOverrides(),
-        kafkaIoConfig.isMultiTopic()
+        kafkaIoConfig.isMultiTopic(),
+        ioConfig.getTaskDuration()
     );
   }