Add code to wait for segments generated to be loaded on historicals

docs/api-reference/sql-ingestion-api.md

@@ -257,7 +257,15 @@ The response shows an example report for a query.
         "startTime": "2022-09-14T22:12:09.266Z",
         "durationMs": 28227,
         "pendingTasks": 0,
-        "runningTasks": 2
+        "runningTasks": 2,
+        "segmentLoadWaiterStatus": {
+          "state": "SUCCESS",
+          "dataSource": "kttm_simple",
+          "startTime": "2022-09-14T23:12:09.266Z",
+          "duration": 15,
+          "totalSegments": 1,
+          "segmentsLeft": 0
+        }
       },
       "stages": [
         {
@@ -562,6 +570,12 @@ The following table describes the response fields when you retrieve a report for
 | `multiStageQuery.payload.status.durationMs` | Milliseconds elapsed after the query has started running. -1 denotes that the query hasn't started running yet. |
 | `multiStageQuery.payload.status.pendingTasks` | Number of tasks that are not fully started. -1 denotes that the number is currently unknown. |
 | `multiStageQuery.payload.status.runningTasks` | Number of currently running tasks. Should be at least 1 since the controller is included. |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus` | Segment loading waiter container. Only present after the segments have been published. |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus.state` | Either INIT, WAITING, SUCCESS, FAILED or TIMED_OUT. |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus.startTime` | Time since which the controller has been waiting for the segments to finish loading. |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus.duration` | The duration in milliseconds that the controller has been waiting for the segments to load. |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus.totalSegments` | The total number of segments generated by the job. This includes tombstone segments (if any). |
+| `multiStageQuery.payload.status.segmentLoadWaiterStatus.segmentsLeft` | The total number of segments remaining to be loaded. |
 | `multiStageQuery.payload.status.errorReport` | Error object. Only present if there was an error. |
 | `multiStageQuery.payload.status.errorReport.taskId` | The task that reported the error, if known. May be a controller task or a worker task. |
 | `multiStageQuery.payload.status.errorReport.host` | The hostname and port of the task that reported the error, if known. |

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java

@@ -44,6 +44,7 @@
 import org.apache.druid.data.input.impl.DimensionSchema;
 import org.apache.druid.data.input.impl.DimensionsSpec;
 import org.apache.druid.data.input.impl.TimestampSpec;
+import org.apache.druid.discovery.BrokerClient;
 import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
 import org.apache.druid.frame.channel.FrameChannelSequence;
 import org.apache.druid.frame.key.ClusterBy;
@@ -290,6 +291,7 @@ public class ControllerImpl implements Controller
   private WorkerMemoryParameters workerMemoryParameters;
   private boolean isDurableStorageEnabled;
   private boolean isFaultToleranceEnabled;
+  private volatile SegmentLoadWaiter segmentLoadWaiter;
 
   public ControllerImpl(
       final MSQControllerTask task,
@@ -429,6 +431,37 @@ public TaskStatus runTask(final Closer closer)
       }
     }
 
+    if (queryKernel != null && queryKernel.isSuccess()) {
+      // If successful, encourage the tasks to exit successfully.
+      postFinishToAllTasks();
+      workerTaskLauncher.stop(false);
+    } else {
+      // If not successful, cancel running tasks.
+      if (workerTaskLauncher != null) {
+        workerTaskLauncher.stop(true);
+      }
+    }
+
+    // Wait for worker tasks to exit. Ignore their return status. At this point, we've done everything we need to do,
+    // so we don't care about the task exit status.
+    if (workerTaskRunnerFuture != null) {
+      try {
+        workerTaskRunnerFuture.get();
+      }
+      catch (Exception ignored) {
+        // Suppress.
+      }
+    }
+
+    cleanUpDurableStorageIfNeeded();
+
+    if (queryKernel != null && queryKernel.isSuccess()) {
+      if (segmentLoadWaiter != null) {
+        // If successful and there are segments created, segmentLoadWaiter should wait for them to become available.
+        segmentLoadWaiter.waitForSegmentsToLoad();
+      }
+    }
+
     try {
       // Write report even if something went wrong.
       final MSQStagesReport stagesReport;
@@ -480,7 +513,8 @@ public TaskStatus runTask(final Closer closer)
               workerWarnings,
               queryStartTime,
               new Interval(queryStartTime, DateTimes.nowUtc()).toDurationMillis(),
-              workerTaskLauncher
+              workerTaskLauncher,
+              segmentLoadWaiter
           ),
           stagesReport,
           countersSnapshot,
@@ -496,30 +530,6 @@ public TaskStatus runTask(final Closer closer)
       log.warn(e, "Error encountered while writing task report. Skipping.");
     }
 
-    if (queryKernel != null && queryKernel.isSuccess()) {
-      // If successful, encourage the tasks to exit successfully.
-      postFinishToAllTasks();
-      workerTaskLauncher.stop(false);
-    } else {
-      // If not successful, cancel running tasks.
-      if (workerTaskLauncher != null) {
-        workerTaskLauncher.stop(true);
-      }
-    }
-
-    // Wait for worker tasks to exit. Ignore their return status. At this point, we've done everything we need to do,
-    // so we don't care about the task exit status.
-    if (workerTaskRunnerFuture != null) {
-      try {
-        workerTaskRunnerFuture.get();
-      }
-      catch (Exception ignored) {
-        // Suppress.
-      }
-    }
-
-    cleanUpDurableStorageIfNeeded();
-
     if (taskStateForReport == TaskState.SUCCESS) {
       return TaskStatus.success(id());
     } else {
@@ -858,7 +868,8 @@ public Map<String, TaskReport> liveReports()
                     workerWarnings,
                     queryStartTime,
                     queryStartTime == null ? -1L : new Interval(queryStartTime, DateTimes.nowUtc()).toDurationMillis(),
-                    workerTaskLauncher
+                    workerTaskLauncher,
+                    segmentLoadWaiter
                 ),
                 makeStageReport(
                     queryDef,
@@ -1299,17 +1310,36 @@ private void publishAllSegments(final Set<DataSegment> segments) throws IOExcept
       if (segmentsWithTombstones.isEmpty()) {
         // Nothing to publish, only drop. We already validated that the intervalsToDrop do not have any
         // partially-overlapping segments, so it's safe to drop them as intervals instead of as specific segments.
+        // This should not need a segment load wait as segments are marked as unused immediately.
         for (final Interval interval : intervalsToDrop) {
           context.taskActionClient()
                  .submit(new MarkSegmentsAsUnusedAction(task.getDataSource(), interval));
         }
       } else {
+        Set<String> versionsToAwait = segmentsWithTombstones.stream().map(DataSegment::getVersion).collect(Collectors.toSet());
+        segmentLoadWaiter = new SegmentLoadWaiter(
+            context.injector().getInstance(BrokerClient.class),
+            context.jsonMapper(),
+            task.getDataSource(),
+            versionsToAwait,
+            segmentsWithTombstones.size(),
+            true
+        );
         performSegmentPublish(
             context.taskActionClient(),
             SegmentTransactionalInsertAction.overwriteAction(null, null, segmentsWithTombstones)
         );
       }
     } else if (!segments.isEmpty()) {
+      Set<String> versionsToAwait = segments.stream().map(DataSegment::getVersion).collect(Collectors.toSet());
+      segmentLoadWaiter = new SegmentLoadWaiter(
+          context.injector().getInstance(BrokerClient.class),
+          context.jsonMapper(),
+          task.getDataSource(),
+          versionsToAwait,
+          segmentsWithTombstones.size(),
+          true
+      );
       // Append mode.
       performSegmentPublish(
           context.taskActionClient(),
@@ -2055,7 +2085,8 @@ private static MSQStatusReport makeStatusReport(
       final Queue<MSQErrorReport> errorReports,
       @Nullable final DateTime queryStartTime,
       final long queryDuration,
-      MSQWorkerTaskLauncher taskLauncher
+      MSQWorkerTaskLauncher taskLauncher,
+      final SegmentLoadWaiter segmentLoadWaiter
   )
   {
     int pendingTasks = -1;
@@ -2066,14 +2097,18 @@ private static MSQStatusReport makeStatusReport(
       pendingTasks = workerTaskCount.getPendingWorkerCount();
       runningTasks = workerTaskCount.getRunningWorkerCount() + 1; // To account for controller.
     }
+
+    SegmentLoadWaiter.SegmentLoadWaiterStatus status = segmentLoadWaiter == null ? null : segmentLoadWaiter.status();
+
     return new MSQStatusReport(
         taskState,
         errorReport,
         errorReports,
         queryStartTime,
         queryDuration,
         pendingTasks,
-        runningTasks
+        runningTasks,
+        status
     );
   }
 
@@ -2242,6 +2277,7 @@ private Pair<ControllerQueryKernel, ListenableFuture<?>> run() throws IOExceptio
         throwKernelExceptionIfNotUnknown();
       }
 
+      updateLiveReportMaps();
       cleanUpEffectivelyFinishedStages();
       return Pair.of(queryKernel, workerTaskLauncherFuture);
     }