Fix the potential race between SplittableInputSource.getNumSplits() and SplittableInputSource.createSplits() in TaskMonitor (#8924)

* Fix the potential race SplittableInputSource.getNumSplits() and SplittableInputSource.createSplits() in TaskMonitor

* Fix docs and javadoc

* Add unit tests for large or small estimated num splits

* add override

Author: jihoonson

core/src/main/java/org/apache/druid/data/input/FirehoseFactoryToInputSourceAdaptor.java

@@ -68,7 +68,7 @@ public Stream<InputSplit> createSplits(InputFormat inputFormat, @Nullable SplitH
   }
 
   @Override
-  public int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException
+  public int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException
   {
     if (firehoseFactory.isSplittable()) {
       return firehoseFactory.getNumSplits(splitHintSpec);

core/src/main/java/org/apache/druid/data/input/SplitHintSpec.java

@@ -33,7 +33,7 @@
  * @see FiniteFirehoseFactory#getSplits(SplitHintSpec)
  * @see FiniteFirehoseFactory#getNumSplits(SplitHintSpec)
  * @see SplittableInputSource#createSplits
- * @see SplittableInputSource#getNumSplits
+ * @see SplittableInputSource#estimateNumSplits
  */
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
 @JsonSubTypes(value = {

core/src/main/java/org/apache/druid/data/input/impl/HttpInputSource.java

@@ -86,7 +86,7 @@ public Stream<InputSplit<URI>> createSplits(InputFormat inputFormat, @Nullable S
   }
 
   @Override
-  public int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
+  public int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
   {
     return uris.size();
   }

core/src/main/java/org/apache/druid/data/input/impl/LocalInputSource.java

@@ -77,7 +77,7 @@ public Stream<InputSplit<File>> createSplits(InputFormat inputFormat, @Nullable
   }
 
   @Override
-  public int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
+  public int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
   {
     return Iterators.size(getFileIterator());
   }

core/src/main/java/org/apache/druid/data/input/impl/SplittableInputSource.java

@@ -44,23 +44,27 @@ default boolean isSplittable()
   /**
    * Creates a {@link Stream} of {@link InputSplit}s. The returned stream is supposed to be evaluated lazily to avoid
    * consuming too much memory.
-   * Note that this interface also has {@link #getNumSplits} which is related to this method. The implementations
+   * Note that this interface also has {@link #estimateNumSplits} which is related to this method. The implementations
    * should be careful to <i>NOT</i> cache the created splits in memory.
    *
    * Implementations can consider {@link InputFormat#isSplittable()} and {@link SplitHintSpec} to create splits
-   * in the same way with {@link #getNumSplits}.
+   * in the same way with {@link #estimateNumSplits}.
    */
   Stream<InputSplit<T>> createSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException;
 
   /**
-   * Returns the total number of splits to be created via {@link #createSplits}.
-   * This method can be expensive since it needs to iterate all directories or whatever substructure
-   * to find all input objects.
+   * Returns an estimated total number of splits to be created via {@link #createSplits}. The estimated number of splits
+   * doesn't have to be accurate and can be different from the actual number of InputSplits returned from
+   * {@link #createSplits}. This will be used to estimate the progress of a phase in parallel indexing.
+   * See TaskMonitor for more details of the progress estimation.
+   *
+   * This method can be expensive if an implementation iterates all directories or whatever substructure
+   * to find all input entities.
    *
    * Implementations can consider {@link InputFormat#isSplittable()} and {@link SplitHintSpec} to find splits
    * in the same way with {@link #createSplits}.
    */
-  int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException;
+  int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException;
 
   /**
    * Helper method for ParallelIndexSupervisorTask.

docs/ingestion/native-batch.md

@@ -274,7 +274,7 @@ Returns the name of the current phase if the task running in the parallel mode.
 
 * `http://{PEON_IP}:{PEON_PORT}/druid/worker/v1/chat/{SUPERVISOR_TASK_ID}/progress`
 
-Returns the current progress if the supervisor task is running in the parallel mode.
+Returns the estimated progress of the current phase if the supervisor task is running in the parallel mode.
 
 An example of the result is
 
@@ -285,7 +285,7 @@ An example of the result is
   "failed":0,
   "complete":0,
   "total":10,
-  "expectedSucceeded":10
+  "estimatedExpectedSucceeded":10
 }
 ```
 

extensions-core/google-extensions/src/main/java/org/apache/druid/data/input/google/GoogleCloudStorageInputSource.java

@@ -69,7 +69,7 @@ public Stream<InputSplit<URI>> createSplits(InputFormat inputFormat, @Nullable S
   }
 
   @Override
-  public int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
+  public int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec)
   {
     return uris.size();
   }

extensions-core/hdfs-storage/src/main/java/org/apache/druid/inputsource/hdfs/HdfsInputSource.java

@@ -174,7 +174,7 @@ public Stream<InputSplit<Path>> createSplits(InputFormat inputFormat, @Nullable
   }
 
   @Override
-  public int getNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException
+  public int estimateNumSplits(InputFormat inputFormat, @Nullable SplitHintSpec splitHintSpec) throws IOException
   {
     cachePathsIfNeeded();
     return cachedPaths.size();

extensions-core/hdfs-storage/src/test/java/org/apache/druid/inputsource/hdfs/HdfsInputSourceTest.java

@@ -255,7 +255,7 @@ public void hasCorrectSplits() throws IOException
     @Test
     public void hasCorrectNumberOfSplits() throws IOException
     {
-      int numSplits = target.getNumSplits(null, null);
+      int numSplits = target.estimateNumSplits(null, null);
       Assert.assertEquals(NUM_FILE, numSplits);
     }
   }
@@ -295,7 +295,7 @@ public void hasCorrectSplits() throws IOException
     @Test
     public void hasCorrectNumberOfSplits() throws IOException
     {
-      int numSplits = target.getNumSplits(null, null);
+      int numSplits = target.estimateNumSplits(null, null);
       Assert.assertEquals(0, numSplits);
     }
   }

indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/InputSourceSplitParallelIndexTaskRunner.java

@@ -74,9 +74,9 @@ Iterator<SubTaskSpec<T>> subTaskSpecIterator() throws IOException
   }
 
   @Override
-  final int getTotalNumSubTasks() throws IOException
+  final int estimateTotalNumSubTasks() throws IOException
   {
-    return baseInputSource.getNumSplits(
+    return baseInputSource.estimateNumSplits(
         ingestionSchema.getIOConfig().getInputFormat(),
         getTuningConfig().getSplitHintSpec()
     );