Add utility method to EmbeddedClusterApis to run tasks (#18349)

Description:
Embedded tests currently run into serialization issues when submitting a `Task` or
`CompactionSupervisorSpec` payload to the cluster using any `OverlordClient`.
This can happen because the `ObjectMapper` instance used by the `OverlordClient` might not have the
necessary Jackson modules loaded.
The hack so far has been to ensure that the first server in an embedded cluster is always an Overlord.
But this can quickly run into issues, e.g. when running a cluster with basic auth enabled where a
Coordinator must be started before all other servers.

Fix:
- Fix up methods `EmbeddedClusterApis`: `leaderOverlord`, `leaderCoordinator` and `anyBroker` to use
the appropriate servers. The request is still sent over the wire even if the client bound on a server is used
to call itself.
- Add utility methods in `EmbeddedClusterApis` to `submitTask`, `runTask`, `postSupervisor`

Author: kfaraz

embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.testing.embedded.compact;
 
+import org.apache.druid.common.utils.IdUtils;
 import org.apache.druid.indexing.common.task.TaskBuilder;
 import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
@@ -68,12 +69,11 @@ protected Closeable unloader(String dataSource)
    */
   protected String runTask(TaskBuilder<?, ?, ?> taskBuilder)
   {
-    return cluster.callApi().runTask(
-        (ds, taskId) -> taskBuilder.dataSource(ds).withId(taskId),
-        dataSource,
-        overlord,
-        coordinator
-    );
+    final String taskId = IdUtils.getRandomId();
+    cluster.callApi().runTask(taskBuilder.dataSource(dataSource).withId(taskId), overlord);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(dataSource, coordinator);
+
+    return taskId;
   }
 
   protected void verifySegmentIntervals(List<Interval> expectedIntervals)

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/ConcurrentAppendReplaceTest.java

@@ -68,10 +68,7 @@ public void test_concurrentAppend_toIntervalWithUnusedSegment_usesNewSegmentId()
                    .appendToExisting(true)
                    .context("useConcurrentLocks", true)
                    .dimensions();
-    cluster.callApi().onLeaderOverlord(
-        o -> o.runTask(task1, taskBuilder.withId(task1))
-    );
-    cluster.callApi().waitForTaskToSucceed(task1, overlord);
+    cluster.callApi().runTask(taskBuilder.withId(task1), overlord);
 
     List<DataSegment> usedSegments = getAllUsedSegments();
     Assertions.assertEquals(1, usedSegments.size());
@@ -87,10 +84,7 @@ public void test_concurrentAppend_toIntervalWithUnusedSegment_usesNewSegmentId()
 
     // Run the APPEND task again with a different taskId
     final String task2 = EmbeddedClusterApis.newTaskId(dataSource);
-    cluster.callApi().onLeaderOverlord(
-        o -> o.runTask(task2, taskBuilder.withId(task2))
-    );
-    cluster.callApi().waitForTaskToSucceed(task2, overlord);
+    cluster.callApi().runTask(taskBuilder.withId(task2), overlord);
 
     // Verify that the new segment gets appended with the same version but a different ID
     usedSegments = getAllUsedSegments();

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/IndexParallelTaskTest.java

@@ -214,9 +214,7 @@ public void test_runIndexTask_andReindexIntoAnotherDatasource(PartitionsSpec par
   private String runTask(TaskBuilder.IndexParallel taskBuilder, String dataSource)
   {
     final String taskId = EmbeddedClusterApis.newTaskId(dataSource);
-    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, taskBuilder.withId(taskId)));
-    cluster.callApi().waitForTaskToSucceed(taskId, overlord);
-
+    cluster.callApi().runTask(taskBuilder.withId(taskId), overlord);
     return taskId;
   }
 

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/IndexTaskTest.java

@@ -63,8 +63,8 @@ public EmbeddedDruidCluster createCluster()
   {
     return EmbeddedDruidCluster.withEmbeddedDerbyAndZookeeper()
                                .useLatchableEmitter()
-                               .addServer(overlord)
                                .addServer(coordinator)
+                               .addServer(overlord)
                                .addServer(indexer)
                                .addServer(historical)
                                .addServer(broker)
@@ -76,13 +76,12 @@ public EmbeddedDruidCluster createCluster()
   public void test_runIndexTask_forInlineDatasource()
   {
     final String taskId = EmbeddedClusterApis.newTaskId(dataSource);
-    final Object task = createIndexTaskForInlineData(
+    final IndexTask task = createIndexTaskForInlineData(
         taskId,
         Resources.InlineData.CSV_10_DAYS
     );
 
-    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
-    cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+    cluster.callApi().runTask(task, overlord);
 
     // Verify that the task created 10 DAY-granularity segments
     final List<DataSegment> segments = new ArrayList<>(
@@ -142,7 +141,7 @@ public void test_run100Tasks_concurrently()
     runTasksConcurrently(100);
   }
 
-  private Object createIndexTaskForInlineData(String taskId, String inlineDataCsv)
+  private IndexTask createIndexTaskForInlineData(String taskId, String inlineDataCsv)
   {
     return TaskBuilder.ofTypeIndex()
                       .dataSource(dataSource)
@@ -169,16 +168,14 @@ private void runTasksConcurrently(int count)
     int index = 0;
     for (String taskId : taskIds) {
       index++;
-      final Object task = createIndexTaskForInlineData(
+      final IndexTask task = createIndexTaskForInlineData(
           taskId,
           StringUtils.format(
               "%s,%s,%d",
               jan1.plusDays(index), "item " + index, index
           )
       );
-      cluster.callApi().onLeaderOverlord(
-          o -> o.runTask(taskId, task)
-      );
+      cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
     }
     for (String taskId : taskIds) {
       cluster.callApi().waitForTaskToSucceed(taskId, overlord);

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/KafkaClusterMetricsTest.java

@@ -156,10 +156,10 @@ public void test_ingest10kRows_ofSelfClusterMetrics_andVerifyValues()
         maxRowsPerSegment
     );
 
-    final Map<String, String> startSupervisorResult = cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec)
+    Assertions.assertEquals(
+        supervisorId,
+        cluster.callApi().postSupervisor(kafkaSupervisorSpec)
     );
-    Assertions.assertEquals(Map.of("id", supervisorId), startSupervisorResult);
 
     // Wait for segments to be handed off
     indexer.latchableEmitter().waitForEventAggregate(
@@ -183,9 +183,7 @@ public void test_ingest10kRows_ofSelfClusterMetrics_andVerifyValues()
     verifyIngestedMetricCountMatchesEmittedCount("coordinator/time", coordinator);
 
     // Suspend the supervisor
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec.createSuspendedSpec())
-    );
+    cluster.callApi().postSupervisor(kafkaSupervisorSpec.createSuspendedSpec());
   }
 
   @Test
@@ -208,9 +206,7 @@ public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andSkip
         taskCompletionTimeoutMillis,
         maxRowsPerSegment
     );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec)
-    );
+    cluster.callApi().postSupervisor(kafkaSupervisorSpec);
 
     // Wait for a task to succeed
     overlord.latchableEmitter().waitForEventAggregate(
@@ -248,9 +244,7 @@ public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andSkip
         false,
         null
     );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(compactionSupervisorSpec)
-    );
+    cluster.callApi().postSupervisor(compactionSupervisorSpec);
 
     // Wait until some compaction tasks have finished
     overlord.latchableEmitter().waitForEventAggregate(
@@ -309,12 +303,8 @@ public void test_ingestClusterMetrics_withConcurrentCompactionSupervisor_andSkip
     );
 
     // Suspend the supervisors
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(compactionSupervisorSpec.createSuspendedSpec())
-    );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec.createSuspendedSpec())
-    );
+    cluster.callApi().postSupervisor(compactionSupervisorSpec.createSuspendedSpec());
+    cluster.callApi().postSupervisor(kafkaSupervisorSpec.createSuspendedSpec());
   }
 
   @Test
@@ -337,9 +327,7 @@ public void test_ingestClusterMetrics_compactionSkipsLockedIntervals()
         taskCompletionTimeoutMillis,
         maxRowsPerSegment
     );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec)
-    );
+    cluster.callApi().postSupervisor(kafkaSupervisorSpec);
 
     // Wait for some segments to be published
     overlord.latchableEmitter().waitForEvent(
@@ -371,9 +359,7 @@ public void test_ingestClusterMetrics_compactionSkipsLockedIntervals()
         false,
         null
     );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(compactionSupervisorSpec)
-    );
+    cluster.callApi().postSupervisor(compactionSupervisorSpec);
 
     // Wait until some skipped metrics have been emitted
     overlord.latchableEmitter().waitForEventAggregate(
@@ -388,12 +374,8 @@ public void test_ingestClusterMetrics_compactionSkipsLockedIntervals()
     );
 
     // Suspend the supervisors
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(compactionSupervisorSpec.createSuspendedSpec())
-    );
-    cluster.callApi().onLeaderOverlord(
-        o -> o.postSupervisor(kafkaSupervisorSpec.createSuspendedSpec())
-    );
+    cluster.callApi().postSupervisor(compactionSupervisorSpec.createSuspendedSpec());
+    cluster.callApi().postSupervisor(kafkaSupervisorSpec.createSuspendedSpec());
   }
 
   /**

embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/BaseRealtimeQueryTest.java

@@ -40,6 +40,7 @@
 import org.apache.druid.segment.indexing.DataSchema;
 import org.apache.druid.testing.embedded.EmbeddedClusterApis;
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.druid.testing.embedded.EmbeddedOverlord;
 import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.joda.time.Period;
@@ -84,13 +85,14 @@ void setupCreateKafkaTopic()
   /**
    * Submits a supervisor spec to the Overlord.
    */
-  protected void submitSupervisor()
+  protected void submitSupervisor(EmbeddedOverlord overlord)
   {
     // Submit a supervisor.
     final KafkaSupervisorSpec kafkaSupervisorSpec = createKafkaSupervisor();
-    final Map<String, String> startSupervisorResult =
-        cluster.callApi().onLeaderOverlord(o -> o.postSupervisor(kafkaSupervisorSpec));
-    Assertions.assertEquals(Map.of("id", dataSource), startSupervisorResult);
+    Assertions.assertEquals(
+        dataSource,
+        cluster.callApi().postSupervisor(kafkaSupervisorSpec)
+    );
   }
 
   /**

embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java

@@ -167,7 +167,7 @@ protected void setupLookups() throws Exception
   void setUpEach()
   {
     msqApis = new EmbeddedMSQApis(cluster, overlord);
-    submitSupervisor();
+    submitSupervisor(overlord);
     publishToKafka(TestIndex.getMMappedWikipediaIndex());
 
     // Wait for it to be loaded.

embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeUnnestQueryTest.java

@@ -113,7 +113,7 @@ void setUpEach()
 
     QueryableIndex index = TestIndex.getMMappedTestIndex();
 
-    submitSupervisor();
+    submitSupervisor(overlord);
     publishToKafka(index);
 
     final int totalRows = index.getNumRows();

embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/CoordinatorClientTest.java

@@ -192,8 +192,7 @@ private void runIndexTask()
                                  .dimensions()
                                  .withId(taskId);
 
-    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
-    cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+    cluster.callApi().runTask(task, overlord);
     cluster.callApi().waitForAllSegmentsToBeAvailable(dataSource, coordinator);
   }
 }

embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/HighAvailabilityTest.java

@@ -121,8 +121,7 @@ public void test_switchLeader_andVerifyUsingSysTables()
         .inlineInputSourceWithData(Resources.InlineData.CSV_10_DAYS)
         .dimensions()
         .withId(taskId);
-    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
-    cluster.callApi().waitForTaskToSucceed(taskId, overlord1);
+    cluster.callApi().runTask(task, overlord1);
     coordinator1.latchableEmitter().waitForEvent(
         event -> event.hasMetricName("segment/metadataCache/used/count")
                       .hasDimension(DruidMetrics.DATASOURCE, dataSource)