Core: Fix Async Planning handling in RESTCatalogAdapter for Remote Scan Planning (#14629)

Co-authored-by: Prashant Singh <[email protected]>
Co-authored-by: Amogh Jahagirdar <[email protected]>

Author: 3 people

core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java

@@ -31,6 +31,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -661,7 +662,7 @@ public static PlanTableScanResponse planTableScan(
     tableScan = tableScan.caseSensitive(request.caseSensitive());
 
     if (shouldPlanAsync.test(tableScan)) {
-      String asyncPlanId = UUID.randomUUID().toString();
+      String asyncPlanId = "async-" + UUID.randomUUID();
       asyncPlanFiles(tableScan, asyncPlanId, tasksPerPlanTask.applyAsInt(tableScan));
       return PlanTableScanResponse.builder()
           .withPlanId(asyncPlanId)
@@ -670,11 +671,12 @@ public static PlanTableScanResponse planTableScan(
           .build();
     }
 
-    String planId = UUID.randomUUID().toString();
+    String planId = "sync-" + UUID.randomUUID();
     planFilesFor(tableScan, planId, tasksPerPlanTask.applyAsInt(tableScan));
     Pair<List<FileScanTask>, String> initial = IN_MEMORY_PLANNING_STATE.initialScanTasksFor(planId);
     return PlanTableScanResponse.builder()
         .withPlanStatus(PlanStatus.COMPLETED)
+        .withPlanId(planId)
         .withPlanTasks(IN_MEMORY_PLANNING_STATE.nextPlanTask(initial.second()))
         .withFileScanTasks(initial.first())
         .withDeleteFiles(
@@ -697,6 +699,11 @@ public static PlanTableScanResponse planTableScan(
   public static FetchPlanningResultResponse fetchPlanningResult(
       Catalog catalog, TableIdentifier ident, String planId) {
     Table table = catalog.loadTable(ident);
+    PlanStatus status = IN_MEMORY_PLANNING_STATE.asyncPlanStatus(planId);
+    if (status != PlanStatus.COMPLETED) {
+      return FetchPlanningResultResponse.builder().withPlanStatus(status).build();
+    }
+
     Pair<List<FileScanTask>, String> initial = IN_MEMORY_PLANNING_STATE.initialScanTasksFor(planId);
     return FetchPlanningResultResponse.builder()
         .withPlanStatus(PlanStatus.COMPLETED)
@@ -743,12 +750,11 @@ public static FetchScanTasksResponse fetchScanTasks(
    * @param planId the plan identifier to cancel
    */
   public static void cancelPlanTableScan(String planId) {
-    IN_MEMORY_PLANNING_STATE.removePlan(planId);
+    IN_MEMORY_PLANNING_STATE.cancelPlan(planId);
   }
 
   static void clearPlanningState() {
     InMemoryPlanningState.getInstance().clear();
-    ASYNC_PLANNING_POOL.shutdown();
   }
 
   /**
@@ -775,8 +781,22 @@ private static void planFilesFor(TableScan tableScan, String planId, int tasksPe
     }
   }
 
+  @SuppressWarnings("FutureReturnValueIgnored")
   private static void asyncPlanFiles(
       TableScan tableScan, String asyncPlanId, int tasksPerPlanTask) {
-    ASYNC_PLANNING_POOL.execute(() -> planFilesFor(tableScan, asyncPlanId, tasksPerPlanTask));
+    IN_MEMORY_PLANNING_STATE.addAsyncPlan(asyncPlanId);
+    CompletableFuture.runAsync(
+            () -> {
+              planFilesFor(tableScan, asyncPlanId, tasksPerPlanTask);
+            },
+            ASYNC_PLANNING_POOL)
+        .whenComplete(
+            (result, exception) -> {
+              if (exception != null) {
+                IN_MEMORY_PLANNING_STATE.markAsyncPlanFailed(asyncPlanId);
+              } else {
+                IN_MEMORY_PLANNING_STATE.markAsyncPlanAsComplete(asyncPlanId);
+              }
+            });
   }
 }

core/src/main/java/org/apache/iceberg/rest/InMemoryPlanningState.java

@@ -25,6 +25,7 @@
 import org.apache.iceberg.FileScanTask;
 import org.apache.iceberg.exceptions.NoSuchPlanIdException;
 import org.apache.iceberg.exceptions.NoSuchPlanTaskException;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.base.Splitter;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
@@ -43,10 +44,12 @@ class InMemoryPlanningState {
 
   private final Map<String, List<FileScanTask>> planTaskToFileScanTasks;
   private final Map<String, String> planTaskToNext;
+  private final Map<String, PlanStatus> asyncPlanningStates;
 
   private InMemoryPlanningState() {
     this.planTaskToFileScanTasks = Maps.newConcurrentMap();
     this.planTaskToNext = Maps.newConcurrentMap();
+    this.asyncPlanningStates = Maps.newConcurrentMap();
   }
 
   static InMemoryPlanningState getInstance() {
@@ -68,6 +71,44 @@ void addNextPlanTask(String currentTask, String nextTask) {
     planTaskToNext.put(currentTask, nextTask);
   }
 
+  void addAsyncPlan(String plan) {
+    PlanStatus existingStatus = asyncPlanningStates.get(plan);
+    Preconditions.checkArgument(
+        existingStatus == null, "Plan %s already exists with status %s", plan, existingStatus);
+    asyncPlanningStates.put(plan, PlanStatus.SUBMITTED);
+  }
+
+  PlanStatus asyncPlanStatus(String plan) {
+    PlanStatus existingStatus = asyncPlanningStates.get(plan);
+    if (existingStatus == null) {
+      throw new NoSuchPlanIdException("Cannot find plan with id %s", plan);
+    }
+
+    return asyncPlanningStates.get(plan);
+  }
+
+  void markAsyncPlanAsComplete(String plan) {
+    PlanStatus existingStatus = asyncPlanningStates.get(plan);
+    Preconditions.checkArgument(existingStatus != null, "Cannot find plan %s", plan);
+    Preconditions.checkArgument(
+        existingStatus == PlanStatus.SUBMITTED,
+        "Cannot mark plan %s as completed as it is %s",
+        plan,
+        existingStatus);
+    asyncPlanningStates.put(plan, PlanStatus.COMPLETED);
+  }
+
+  void markAsyncPlanFailed(String plan) {
+    PlanStatus existingStatus = asyncPlanningStates.get(plan);
+    Preconditions.checkArgument(existingStatus != null, "Cannot find plan %s", plan);
+    Preconditions.checkArgument(
+        existingStatus == PlanStatus.SUBMITTED,
+        "Cannot mark plan %s as completed as it is %s",
+        plan,
+        existingStatus);
+    asyncPlanningStates.put(plan, PlanStatus.FAILED);
+  }
+
   List<FileScanTask> fileScanTasksForPlanTask(String planTaskKey) {
     List<FileScanTask> tasks = planTaskToFileScanTasks.get(planTaskKey);
     if (tasks == null) {
@@ -120,13 +161,27 @@ Pair<List<FileScanTask>, String> initialScanTasksFor(String planId) {
     return Pair.of(initialEntry.getValue(), initialEntry.getKey());
   }
 
-  void removePlan(String planId) {
+  void cancelPlan(String planId) {
     planTaskToNext.entrySet().removeIf(entry -> entry.getKey().contains(planId));
     planTaskToFileScanTasks.entrySet().removeIf(entry -> entry.getKey().contains(planId));
+    // Clear the ongoing plan status in case the planID is an async one.
+    if (asyncPlanningStates.containsKey(planId)) {
+      PlanStatus existingStatus = asyncPlanningStates.get(planId);
+      // No need to fail cancellation if the plan could not be found
+      if (existingStatus == null) {
+        return;
+      }
+
+      // No need to fail cancellation if the plan already terminated
+      if (existingStatus == PlanStatus.SUBMITTED) {
+        asyncPlanningStates.put(planId, PlanStatus.CANCELLED);
+      }
+    }
   }
 
   void clear() {
     planTaskToFileScanTasks.clear();
     planTaskToNext.clear();
+    asyncPlanningStates.clear();
   }
 }