Refactor: Rename DataSourceTaskIdHolder and extract load spec holder to its own class (#18732)

Refactor:

Remove named injections of DataSourceTaskIdHolder in favor of:
- TaskHolder interface: implemented by PeonTaskHolder for peon processes and by NoopTaskHolder for all other servers, where the methods return null.
- LoadSpecHolder interface: implemented by PeonLoadSpecHolder for peon processes and DefaultLoadSpecHolder for all other servers.

This also removes some hacks that were added to avoid cyclic dependencies during Guice init. A module DefaultServerHolderModule that binds the above modules as part of the core injector module for all servers except CliPeon. For CliPeon, the binding to PeonTaskHolder and PeonLoadSpecHolder is explicitly done in CliPeon itself.

Author: abhishekrb19

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java

@@ -35,7 +35,12 @@
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.inject.Injector;
+import org.apache.commons.io.FileUtils;
 import org.apache.curator.test.TestingCluster;
+import org.apache.druid.cli.CliPeon;
+import org.apache.druid.cli.PeonLoadSpecHolder;
+import org.apache.druid.cli.PeonTaskHolder;
 import org.apache.druid.data.input.InputEntity;
 import org.apache.druid.data.input.InputEntityReader;
 import org.apache.druid.data.input.InputFormat;
@@ -52,6 +57,8 @@
 import org.apache.druid.data.input.kafka.KafkaTopicPartition;
 import org.apache.druid.data.input.kafkainput.KafkaInputFormat;
 import org.apache.druid.data.input.kafkainput.KafkaStringHeaderFormat;
+import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.guice.GuiceInjectors;
 import org.apache.druid.indexer.IngestionState;
 import org.apache.druid.indexer.TaskState;
 import org.apache.druid.indexer.TaskStatus;
@@ -112,6 +119,10 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.segment.transform.ExpressionTransform;
 import org.apache.druid.segment.transform.TransformSpec;
+import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec;
+import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
+import org.apache.druid.server.metrics.LoadSpecHolder;
+import org.apache.druid.server.metrics.TaskHolder;
 import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.Resource;
 import org.apache.druid.server.security.ResourceAction;
@@ -127,7 +138,9 @@
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
@@ -141,6 +154,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Properties;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
@@ -153,6 +168,9 @@
 @RunWith(Parameterized.class)
 public class KafkaIndexTaskTest extends SeekableStreamIndexTaskTestBase
 {
+  @Rule
+  public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
   private static final long POLL_RETRY_MS = 100;
   private static final Iterable<Header> SAMPLE_HEADERS = ImmutableList.of(new Header()
   {
@@ -3360,6 +3378,80 @@ public void testCompletionReportMultiplePartitionStats() throws Exception
     Assert.assertTrue(reportData.getRecordsProcessed().values().containsAll(ImmutableSet.of(6L, 2L)));
   }
 
+  @Test(timeout = 60_000L)
+  public void testTaskWithTransformSpecDoesNotCauseCliPeonCyclicDependency()
+      throws IOException, ExecutionException, InterruptedException
+  {
+    insertData();
+
+    final KafkaIndexTask task = createTask(
+        "index_kafka_test_id1",
+        NEW_DATA_SCHEMA.withTransformSpec(
+            new TransformSpec(
+                new SelectorDimFilter("dim1", "b", null),
+                ImmutableList.of(
+                    new ExpressionTransform("dim1t", "concat(dim1,dim1)", ExprMacroTable.nil())
+                )
+            )
+        ),
+        new KafkaIndexTaskIOConfig(
+            0,
+            "sequence0",
+            new SeekableStreamStartSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 0L), ImmutableSet.of()),
+            new SeekableStreamEndSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 5L)),
+            kafkaServer.consumerProperties(),
+            KafkaSupervisorIOConfig.DEFAULT_POLL_TIMEOUT_MILLIS,
+            true,
+            null,
+            null,
+            INPUT_FORMAT,
+            null,
+            Duration.standardHours(2).getStandardMinutes()
+        )
+    );
+
+    File file = temporaryFolder.newFile("task.json");
+
+    FileUtils.write(file, OBJECT_MAPPER.writeValueAsString(task), StandardCharsets.UTF_8);
+
+    final CliPeon peon = new CliPeon();
+    peon.taskAndStatusFile = ImmutableList.of(file.getParent(), "1");
+
+    final Properties properties = new Properties();
+    peon.configure(properties);
+    final Injector baseInjector = GuiceInjectors.makeStartupInjector();
+    peon.configure(properties, baseInjector);
+    final Injector peonInjector = peon.makeInjector(Set.of(NodeRole.PEON));
+
+    final LoadSpecHolder loadSpecHolder = peonInjector.getInstance(LoadSpecHolder.class);
+    Assert.assertTrue(loadSpecHolder instanceof PeonLoadSpecHolder);
+    Assert.assertEquals(LookupLoadingSpec.ALL, loadSpecHolder.getLookupLoadingSpec());
+    Assert.assertEquals(BroadcastDatasourceLoadingSpec.ALL, loadSpecHolder.getBroadcastDatasourceLoadingSpec());
+
+    final TaskHolder taskHolder = peonInjector.getInstance(TaskHolder.class);
+    Assert.assertTrue(taskHolder instanceof PeonTaskHolder);
+    Assert.assertEquals("index_kafka_test_id1", taskHolder.getTaskId());
+    Assert.assertEquals("test_ds", taskHolder.getDataSource());
+
+    final ListenableFuture<TaskStatus> future = runTask(task);
+
+    // Wait for task to exit
+    Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
+    verifyTaskMetrics(task, RowMeters.with().bytes(getTotalSizeOfRecords(0, 5)).thrownAway(4).totalProcessed(1));
+
+    // Check published metadata
+    final List<SegmentDescriptor> publishedDescriptors = publishedDescriptors();
+    assertEqualsExceptVersion(ImmutableList.of(sdd("2009/P1D", 0)), publishedDescriptors);
+    Assert.assertEquals(
+        new KafkaDataSourceMetadata(new SeekableStreamEndSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 5L))),
+        newDataSchemaMetadata()
+    );
+
+    // Check segments in deep storage
+    Assert.assertEquals(ImmutableList.of("b"), readSegmentColumn("dim1", publishedDescriptors.get(0)));
+    Assert.assertEquals(ImmutableList.of("bb"), readSegmentColumn("dim1t", publishedDescriptors.get(0)));
+  }
+
   public static class TestKafkaInputFormat implements InputFormat
   {
     final InputFormat baseInputFormat;

indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java

@@ -1240,6 +1240,8 @@ public static class Builder
     private final String dataSource;
     private final SegmentCacheManagerFactory segmentCacheManagerFactory;
 
+    @Nullable
+    private String id;
     private CompactionIOConfig ioConfig;
     @Nullable
     private DimensionsSpec dimensionsSpec;
@@ -1268,6 +1270,12 @@ public Builder(
       this.segmentCacheManagerFactory = segmentCacheManagerFactory;
     }
 
+    public Builder id(String id)
+    {
+      this.id = id;
+      return this;
+    }
+
     public Builder interval(Interval interval)
     {
       return inputSpec(new CompactionIntervalSpec(interval, null));
@@ -1353,7 +1361,7 @@ public Builder projections(@Nullable List<AggregateProjectionSpec> projections)
     public CompactionTask build()
     {
       return new CompactionTask(
-          null,
+          id,
           null,
           dataSource,
           null,

quidem-ut/src/main/java/org/apache/druid/quidem/ExposedAsBrokerQueryComponentSupplierWrapper.java

@@ -41,6 +41,7 @@
 import org.apache.druid.discovery.DruidNodeDiscoveryProvider;
 import org.apache.druid.guice.AnnouncerModule;
 import org.apache.druid.guice.BrokerServiceModule;
+import org.apache.druid.guice.DefaultServerHolderModule;
 import org.apache.druid.guice.ExpressionModule;
 import org.apache.druid.guice.ExtensionsModule;
 import org.apache.druid.guice.JacksonConfigManagerModule;
@@ -194,6 +195,7 @@ private List<Module> forServerModules()
         new ExternalStorageAccessSecurityModule(),
         new ServiceClientModule(),
         new StorageConnectorModule(),
+        new DefaultServerHolderModule(),
         ServerInjectorBuilder.registerNodeRoleModule(ImmutableSet.of())
     );
   }

server/src/main/java/org/apache/druid/guice/DefaultServerHolderModule.java

@@ -0,0 +1,50 @@
+/*
+ * 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.druid.guice;
+
+import com.google.inject.Binder;
+import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.guice.annotations.ExcludeScope;
+import org.apache.druid.initialization.DruidModule;
+import org.apache.druid.server.metrics.DefaultLoadSpecHolder;
+import org.apache.druid.server.metrics.LoadSpecHolder;
+import org.apache.druid.server.metrics.NoopTaskHolder;
+import org.apache.druid.server.metrics.TaskHolder;
+
+/**
+ * Binds the following holder configs for all servers except {@code CliPeon}:
+ * <ul>
+ * <li>{@link TaskHolder} to {@link NoopTaskHolder}</li>
+ * <li>{@link LoadSpecHolder} to {@link DefaultLoadSpecHolder}</li>
+ * </ul>
+ *
+ * <p>For {@code CliPeon}, these bindings are overridden by the peon-specific module.</p>
+ */
+@ExcludeScope(roles = {NodeRole.PEON_JSON_NAME})
+public class DefaultServerHolderModule implements DruidModule
+{
+  @Override
+  public void configure(Binder binder)
+  {
+    binder.bind(TaskHolder.class).to(NoopTaskHolder.class).in(LazySingleton.class);
+    binder.bind(LoadSpecHolder.class).to(DefaultLoadSpecHolder.class).in(LazySingleton.class);
+  }
+}

server/src/main/java/org/apache/druid/guice/annotations/ExcludeScope.java

@@ -25,7 +25,7 @@
 import java.lang.annotation.Target;
 
 /**
- * An annotation to exlcude specific node types that a {@link com.google.inject.Module} can be loaded on.
+ * An annotation to exclude specific node types that a {@link com.google.inject.Module} can be loaded on.
  * The {@link #roles()} should be the {@link org.apache.druid.discovery.NodeRole#jsonName}. If both {@link ExcludeScope}
  * and {@link LoadScope} are set, {@link ExcludeScope} takes precedence.
  * <p>

server/src/main/java/org/apache/druid/initialization/CoreInjectorBuilder.java

@@ -26,6 +26,7 @@
 import org.apache.druid.guice.AnnouncerModule;
 import org.apache.druid.guice.BuiltInTypesModule;
 import org.apache.druid.guice.CatalogCoreModule;
+import org.apache.druid.guice.DefaultServerHolderModule;
 import org.apache.druid.guice.DruidInjectorBuilder;
 import org.apache.druid.guice.DruidSecondaryModule;
 import org.apache.druid.guice.ExpressionModule;
@@ -142,7 +143,8 @@ public CoreInjectorBuilder forServerWithoutJetty()
         new ExternalStorageAccessSecurityModule(),
         new ServiceClientModule(),
         new StorageConnectorModule(),
-        new CatalogCoreModule()
+        new CatalogCoreModule(),
+        new DefaultServerHolderModule()
     );
     return this;
   }

server/src/main/java/org/apache/druid/query/lookup/LookupListeningAnnouncerConfig.java

@@ -25,23 +25,28 @@
 import com.google.common.base.Preconditions;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
-import org.apache.druid.server.metrics.DataSourceTaskIdHolder;
+import org.apache.druid.server.metrics.LoadSpecHolder;
+import org.apache.druid.server.metrics.TaskHolder;
 
 class LookupListeningAnnouncerConfig
 {
   public static final String DEFAULT_TIER = "__default";
-  private final DataSourceTaskIdHolder dataSourceTaskIdHolder;
+  private final TaskHolder taskHolder;
+  private final LoadSpecHolder loadSpecHolder;
+
   @JsonProperty("lookupTier")
   private String lookupTier = null;
   @JsonProperty("lookupTierIsDatasource")
   private boolean lookupTierIsDatasource = false;
 
   @JsonCreator
   public LookupListeningAnnouncerConfig(
-      @JacksonInject DataSourceTaskIdHolder dataSourceTaskIdHolder
+      @JacksonInject TaskHolder taskHolder,
+      @JacksonInject LoadSpecHolder loadSpecHolder
   )
   {
-    this.dataSourceTaskIdHolder = dataSourceTaskIdHolder;
+    this.taskHolder = taskHolder;
+    this.loadSpecHolder = loadSpecHolder;
   }
 
   public String getLookupTier()
@@ -50,7 +55,8 @@ public String getLookupTier()
         !(lookupTierIsDatasource && null != lookupTier),
         "Cannot specify both `lookupTier` and `lookupTierIsDatasource`"
     );
-    final String lookupTier = lookupTierIsDatasource ? dataSourceTaskIdHolder.getDataSource() : this.lookupTier;
+
+    final String lookupTier = lookupTierIsDatasource ? taskHolder.getDataSource() : this.lookupTier;
 
     return Preconditions.checkNotNull(
         lookupTier == null ? DEFAULT_TIER : StringUtils.emptyToNullNonDruidDataString(lookupTier),
@@ -61,6 +67,6 @@ public String getLookupTier()
 
   public LookupLoadingSpec getLookupLoadingSpec()
   {
-    return dataSourceTaskIdHolder.getLookupLoadingSpec();
+    return loadSpecHolder.getLookupLoadingSpec();
   }
 }

server/src/main/java/org/apache/druid/query/lookup/LookupReferencesManager.java

@@ -42,7 +42,7 @@
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.server.http.ServletResourceUtils;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
-import org.apache.druid.server.metrics.DataSourceTaskIdHolder;
+import org.apache.druid.server.metrics.LoadSpecHolder;
 
 import javax.annotation.Nullable;
 import java.io.File;
@@ -380,7 +380,7 @@ private void takeSnapshot(Map<String, LookupExtractorFactoryContainer> lookupMap
   }
 
   /**
-   * Load a set of lookups based on the injected value in {@link DataSourceTaskIdHolder#getLookupLoadingSpec()}.
+   * Load a set of lookups based on the injected value in {@link LoadSpecHolder#getLookupLoadingSpec()}.
    */
   private void loadLookupsAndInitStateRef()
   {

server/src/main/java/org/apache/druid/segment/realtime/ChatHandlerResource.java

@@ -25,8 +25,9 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.server.initialization.jetty.BadRequestException;
 import org.apache.druid.server.initialization.jetty.ServiceUnavailableException;
-import org.apache.druid.server.metrics.DataSourceTaskIdHolder;
+import org.apache.druid.server.metrics.TaskHolder;
 
+import javax.annotation.Nullable;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.core.Context;
@@ -39,13 +40,14 @@ public class ChatHandlerResource
   public static final String TASK_ID_HEADER = "X-Druid-Task-Id";
 
   private final ChatHandlerProvider handlers;
+  @Nullable
   private final String taskId;
 
   @Inject
-  public ChatHandlerResource(final ChatHandlerProvider handlers, final DataSourceTaskIdHolder taskIdHolder)
+  public ChatHandlerResource(final ChatHandlerProvider handlers, final TaskHolder taskHolder)
   {
     this.handlers = handlers;
-    this.taskId = taskIdHolder.getTaskId();
+    this.taskId = taskHolder.getTaskId();
   }
 
   @Path("/{id}")