[AMORO-4295] AIP-5 Phase 3+4: idle-driven scale-down with graceful drain for dynamic allocation groups

amoro-ams/src/main/java/org/apache/amoro/server/DefaultOptimizingService.java

@@ -116,6 +116,16 @@ public class DefaultOptimizingService extends StatedPersistentBase
   private final Map<String, OptimizingQueue> optimizingQueueByGroup = new ConcurrentHashMap<>();
   private final Map<String, OptimizingQueue> optimizingQueueByToken = new ConcurrentHashMap<>();
   private final Map<String, OptimizerInstance> authOptimizers = new ConcurrentHashMap<>();
+
+  /**
+   * Tokens of draining optimizers (AIP-5 scale-down): {@link #pollTask} returns {@code null} for
+   * them, blocking new assignments while in-flight tasks complete normally.
+   */
+  private final Set<String> pendingRemovalTokens = ConcurrentHashMap.newKeySet();
+
+  /** Force-removal deadline per draining token, the {@code drain-timeout} safety net. */
+  private final Map<String, Long> drainDeadlines = new ConcurrentHashMap<>();
+
   private final OptimizerKeeper optimizerKeeper = new OptimizerKeeper("optimizer-keeper-thread");
   private final OptimizerGroupKeeper optimizerGroupKeeper =
       new OptimizerGroupKeeper("optimizer-group-keeper-thread");
@@ -242,9 +252,18 @@ private void unregisterOptimizer(String token) {
     doAs(OptimizerMapper.class, mapper -> mapper.deleteOptimizer(token));
     OptimizingQueue optimizingQueue = optimizingQueueByToken.remove(token);
     OptimizerInstance optimizer = authOptimizers.remove(token);
-    if (optimizingQueue != null) {
-      optimizingQueue.removeOptimizer(optimizer);
+    if (optimizer != null) {
+      if (optimizingQueue == null) {
+        optimizingQueue = optimizingQueueByGroup.get(optimizer.getGroupName());
+      }
+      if (optimizingQueue != null) {
+        optimizingQueue.removeOptimizer(optimizer);
+      }
     }
+    // An optimizer that dies mid-drain is unregistered here by heartbeat expiry; its token can
+    // never be matched again, so leftover drain state would sit in the pending-removal set
+    // forever (its replacement pod registers under a fresh token).
+    cancelDrain(token);
   }
 
   @Override
@@ -269,17 +288,118 @@ private OptimizerInstance getAuthenticatedOptimizer(String authToken) {
 
   @Override
   public OptimizingTask pollTask(String authToken, int threadId) {
+    if (pendingRemovalTokens.contains(authToken)) {
+      return null;
+    }
     LOG.debug("Optimizer {} (threadId {}) try polling task", authToken, threadId);
     OptimizerThread optimizerThread = getAuthenticatedOptimizer(authToken).getThread(threadId);
     OptimizingQueue queue = getQueueByToken(authToken);
-    TaskRuntime<?> task = queue.pollTask(optimizerThread, pollingTimeout, breakQuotaLimit);
+    TaskRuntime<?> task =
+        guardDrainedPoll(
+            authToken, queue.pollTask(optimizerThread, pollingTimeout, breakQuotaLimit));
     if (task != null) {
       LOG.info("OptimizerThread {} polled task {}", optimizerThread, task.getTaskId());
       return task.extractProtocolTask();
     }
     return null;
   }
 
+  /**
+   * Close the long-poll race on drain start: the entry check above cannot stop a thread already
+   * parked inside the queue's poll, which may fetch a task after its token entered the
+   * pending-removal set. Hand such a task back instead of assigning it to a draining optimizer.
+   */
+  @VisibleForTesting
+  TaskRuntime<?> guardDrainedPoll(String authToken, TaskRuntime<?> task) {
+    if (task == null || !pendingRemovalTokens.contains(authToken)) {
+      return task;
+    }
+    OptimizingQueue queue = optimizingQueueByToken.get(authToken);
+    if (queue != null) {
+      try {
+        queue.retryTask(task);
+      } catch (Exception e) {
+        // The existing suspending-task safety net will still reclaim it after the removal.
+        LOG.warn(
+            "Failed to hand back task {} from draining optimizer {}",
+            task.getTaskId(),
+            authToken,
+            e);
+      }
+    }
+    return null;
+  }
+
+  /** Block new task assignments to the token; in-flight tasks keep completing normally. */
+  void beginGracefulDrain(String token, long deadlineMs) {
+    drainDeadlines.put(token, deadlineMs);
+    pendingRemovalTokens.add(token);
+    LOG.info("Optimizer {} begins graceful drain", token);
+  }
+
+  /** Re-admit the token to task assignment, e.g. when dynamic allocation is disabled mid-drain. */
+  void cancelDrain(String token) {
+    pendingRemovalTokens.remove(token);
+    drainDeadlines.remove(token);
+  }
+
+  @VisibleForTesting
+  boolean isDraining(String token) {
+    return pendingRemovalTokens.contains(token);
+  }
+
+  /**
+   * Run one dynamic-allocation round for the group at an injected time. The production cadence is
+   * driven by the scale keeper's delay queue with the wall clock; tests inject times because the
+   * validated minimum {@code executor-idle-timeout} (30s) puts real idle waits beyond sane test
+   * durations.
+   */
+  @VisibleForTesting
+  void evaluateDynamicAllocation(String groupName, long nowMs) {
+    ResourceGroup resourceGroup = optimizerManager.getResourceGroup(groupName);
+    OptimizingQueue queue = optimizingQueueByGroup.get(groupName);
+    optimizerScaleKeeper.scaleIfNeeded(
+        resourceGroup, queue, DynamicAllocationConfig.parse(resourceGroup), nowMs);
+  }
+
+  /**
+   * Remove a drained optimizer: release the container resource, delete the persisted resource row,
+   * and unregister. A missing resource row (the pod self-registered after a persist failure, or a
+   * manual release raced the row away) is not an error — the instance itself carries the
+   * container-side identity, so release through it and only skip the row delete; treating this as a
+   * retryable failure would loop forever on a pod whose row can never reappear. A container release
+   * failure keeps the drain state so a later round retries the idempotent deletion.
+   */
+  void executeRemoval(String token) {
+    OptimizerInstance optimizer = authOptimizers.get(token);
+    if (optimizer == null || optimizer.getResourceId() == null) {
+      // Already unregistered, or externally launched: nothing for AMS to release.
+      cancelDrain(token);
+      return;
+    }
+    try {
+      Resource resource = optimizerManager.getResource(optimizer.getResourceId());
+      if (resource != null) {
+        resource.getProperties().putAll(optimizer.getProperties());
+        ((AbstractOptimizerContainer) Containers.get(resource.getContainerName()))
+            .releaseResource(resource);
+        optimizerManager.deleteResource(optimizer.getResourceId());
+      } else {
+        ((AbstractOptimizerContainer) Containers.get(optimizer.getContainerName()))
+            .releaseResource(optimizer);
+      }
+    } catch (Throwable t) {
+      LOG.warn(
+          "Failed to release optimizer {} (resource {}), will retry",
+          token,
+          optimizer.getResourceId(),
+          t);
+      return;
+    }
+    unregisterOptimizer(token);
+    LOG.info("Optimizer {} (resource {}) removed by scale-down", token, optimizer.getResourceId());
+  }
+
   @Override
   public void ackTask(String authToken, int threadId, OptimizingTaskId taskId) {
     LOG.info("Ack task {} by optimizer {} (threadId {})", taskId, authToken, threadId);
@@ -1044,6 +1164,12 @@ private void unwatch(String groupName) {
       watchedGroups.remove(groupName);
       scaleStates.remove(groupName);
       planningBoundStreaks.remove(groupName);
+      // A drain block left behind would starve the group's pods forever once the legacy floor
+      // keeper resumes duty for the disabled group: re-admit them to task assignment.
+      authOptimizers.values().stream()
+          .filter(optimizer -> groupName.equals(optimizer.getGroupName()))
+          .map(OptimizerInstance::getToken)
+          .forEach(DefaultOptimizingService.this::cancelDrain);
       // pendingRegistrations is deliberately kept: a pod requested before a disable survives its
       // boot window, so re-enabling within it does not re-request the same capacity. Entries
       // self-prune past their deadline.
@@ -1094,7 +1220,7 @@ protected void processTask(DraScaleTask task) {
       }
       DynamicAllocationConfig config = DynamicAllocationConfig.parse(resourceGroup);
       try {
-        scaleIfNeeded(resourceGroup, queue, config);
+        scaleIfNeeded(resourceGroup, queue, config, System.currentTimeMillis());
       } catch (Throwable t) {
         LOG.error("Dynamic allocation scale evaluation failed for group {}", task.groupName, t);
       } finally {
@@ -1109,6 +1235,81 @@ private int pendingThreads(String groupName) {
       return pending == null ? 0 : pending.pendingThreads(System.currentTimeMillis());
     }
 
+    /**
+     * Advance this group's drains: an entry whose in-flight count reached zero, or whose {@code
+     * drain-timeout} deadline passed, executes its removal now (a force-removed instance's orphaned
+     * tasks are reclaimed by the existing suspending-task safety net). Returns the thread and
+     * busy-task counts of instances still draining afterwards — a failed release keeps its instance
+     * in both, since it remains registered.
+     */
+    private int[] processDrainProgress(
+        String groupName, DynamicAllocationState.GroupLoad load, long now) {
+      int drainingThreads = 0;
+      int drainingBusy = 0;
+      for (String token : pendingRemovalTokens) {
+        OptimizerInstance optimizer = authOptimizers.get(token);
+        if (optimizer == null) {
+          // Unregistered mid-drain (e.g. its heartbeat expired): nothing left to remove.
+          cancelDrain(token);
+          continue;
+        }
+        if (!groupName.equals(optimizer.getGroupName())) {
+          continue;
+        }
+        int inFlight = load.getInFlightByToken().getOrDefault(token, 0);
+        Long deadline = drainDeadlines.get(token);
+        if (inFlight == 0 || (deadline != null && now >= deadline)) {
+          executeRemoval(token);
+          if (!authOptimizers.containsKey(token)) {
+            continue;
+          }
+        }
+        drainingThreads += optimizer.getThreadCount();
+        drainingBusy += inFlight;
+      }
+      return new int[] {drainingThreads, drainingBusy};
+    }
+
+    private Set<String> registeredTokens(String groupName) {
+      return authOptimizers.values().stream()
+          .filter(optimizer -> groupName.equals(optimizer.getGroupName()))
+          .map(OptimizerInstance::getToken)
+          .collect(Collectors.toSet());
+    }
+
+    private void evaluateScaleDown(
+        String groupName,
+        OptimizingQueue queue,
+        DynamicAllocationState state,
+        DynamicAllocationConfig config,
+        int registeredThreads,
+        int drainingThreads,
+        long now) {
+      List<DynamicAllocationState.RemovalCandidate> candidates =
+          authOptimizers.values().stream()
+              // Externally-registered optimizers (no resourceId) are not AMS's to remove.
+              .filter(optimizer -> groupName.equals(optimizer.getGroupName()))
+              .filter(optimizer -> optimizer.getResourceId() != null)
+              .filter(optimizer -> !pendingRemovalTokens.contains(optimizer.getToken()))
+              .map(
+                  optimizer ->
+                      new DynamicAllocationState.RemovalCandidate(
+                          optimizer.getToken(), optimizer.getThreadCount()))
+              .collect(Collectors.toList());
+      String victim =
+          state.computeScaleDown(candidates, registeredThreads, drainingThreads, config, now);
+      if (victim == null) {
+        return;
+      }
+      beginGracefulDrain(victim, now + config.getDrainTimeout().toMillis());
+      // Only a snapshot taken after the token entered the pending-removal set can prove idleness:
+      // the pre-insert one may miss a task fetched by a long-poll racing the drain start.
+      DynamicAllocationState.GroupLoad fresh = queue.collectDynamicAllocationLoad();
+      if (fresh.getInFlightByToken().getOrDefault(victim, 0) == 0) {
+        executeRemoval(victim);
+      }
+    }
+
     private void recheckAfterUnwatch(String groupName) {
       try {
         ResourceGroup fresh = optimizerManager.getResourceGroup(groupName);
@@ -1155,27 +1356,46 @@ private void warnOnPlanningBoundTransition(
     }
 
     private void scaleIfNeeded(
-        ResourceGroup resourceGroup, OptimizingQueue queue, DynamicAllocationConfig config) {
+        ResourceGroup resourceGroup,
+        OptimizingQueue queue,
+        DynamicAllocationConfig config,
+        long now) {
       String groupName = resourceGroup.getName();
-      long now = System.currentTimeMillis();
       PendingRegistrations pending =
           pendingRegistrations.computeIfAbsent(
               groupName, name -> new PendingRegistrations(BOOT_TIMEOUT_MS));
       DynamicAllocationState state =
           scaleStates.computeIfAbsent(groupName, name -> new DynamicAllocationState());
-      int registeredThreads = getTotalQuota(groupName);
-      int effectiveThreads = registeredThreads + pending.pendingThreads(now);
       DynamicAllocationState.GroupLoad load = queue.collectDynamicAllocationLoad();
+      // Drain progress runs before anything else and unconditionally: a completed or expired
+      // drain must convert to a removal even in rounds that scale up, or a busy drain would
+      // linger to its full timeout while backlog persists.
+      int[] draining = processDrainProgress(groupName, load, now);
+      int drainingThreads = draining[0];
+      int drainingBusy = draining[1];
+      int registeredThreads = getTotalQuota(groupName);
+      // A draining instance takes no new work, so it is accounted as already gone on both sides:
+      // leaving its threads in the capacity undercounts demand by up to their count, and leaving
+      // its tasks in the load keeps future demand (busy >= effective) from ever firing mid-drain.
+      int effectiveThreads = registeredThreads - drainingThreads + pending.pendingThreads(now);
+      int busyThreads = load.getBusyThreads() - drainingBusy;
+      // Observed every round, including scale-up ones: an instance busy through a burst must not
+      // come out of it looking idle since before the burst began.
+      state.observe(registeredTokens(groupName), load.getInFlightByToken(), now);
       warnOnPlanningBoundTransition(groupName, registeredThreads, load);
       int addInstances =
           state.computeScaleUp(
               effectiveThreads,
-              load.getBusyThreads(),
+              busyThreads,
               load.getServiceablePlanned(),
               load.getPendingTables(),
               config,
               now);
       if (addInstances <= 0) {
+        if (!state.wasDemandActive()) {
+          evaluateScaleDown(
+              groupName, queue, state, config, registeredThreads, drainingThreads, now);
+        }
         return;
       }
       int threadsPerInstance = config.getExecutorParallelism();

amoro-ams/src/main/java/org/apache/amoro/server/optimizing/OptimizingQueue.java

@@ -428,11 +428,13 @@ public List<TaskRuntime<?>> collectTasks(Predicate<TaskRuntime<?>> predicate) {
   public DynamicAllocationState.GroupLoad collectDynamicAllocationLoad() {
     Map<Long, Integer> plannedByTable = Maps.newHashMap();
     Map<Long, Integer> occupiedByTable = Maps.newHashMap();
+    Map<String, Integer> inFlightByToken = Maps.newHashMap();
     int busyThreads = 0;
     for (TaskRuntime<?> task : collectTasks()) {
       if (DynamicAllocationState.occupiesThread(task.getStatus())) {
         busyThreads++;
         occupiedByTable.merge(task.getTableId(), 1, Integer::sum);
+        inFlightByToken.merge(task.getToken(), 1, Integer::sum);
       } else if (task.getStatus() == Status.PLANNED) {
         plannedByTable.merge(task.getTableId(), 1, Integer::sum);
       }
@@ -457,7 +459,10 @@ public DynamicAllocationState.GroupLoad collectDynamicAllocationLoad() {
                     targetQuotaByTable.getOrDefault(tableId, 1.0),
                     occupiedByTable.getOrDefault(tableId, 0))));
     return new DynamicAllocationState.GroupLoad(
-        busyThreads, DynamicAllocationState.serviceablePlannedCount(demands), pendingTables);
+        busyThreads,
+        DynamicAllocationState.serviceablePlannedCount(demands),
+        pendingTables,
+        inFlightByToken);
   }
 
   public void retryTask(TaskRuntime<?> taskRuntime) {

amoro-ams/src/main/java/org/apache/amoro/server/optimizing/dra/DynamicAllocationConfig.java

@@ -338,6 +338,19 @@ public void validate() {
         OptimizerProperties.DYNAMIC_ALLOCATION_SUSTAINED_BACKLOG_TIMEOUT, sustainedBacklogTimeout);
     requirePositive(OptimizerProperties.DYNAMIC_ALLOCATION_SCALE_DOWN_COOLDOWN, scaleDownCooldown);
     requirePositive(OptimizerProperties.DYNAMIC_ALLOCATION_DRAIN_TIMEOUT, drainTimeout);
+    // The scale keeper evaluates each group at sustained-backlog-timeout cadence, which is also
+    // the idle observation resolution: sampling slower than half the idle timeout lets an
+    // instance that worked between samples be misjudged as continuously idle and drained.
+    if (sustainedBacklogTimeout.toMillis() * 2 > executorIdleTimeout.toMillis()) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Resource group:%s '%s'(%s) must be <= half of '%s'(%s).",
+              groupName,
+              OptimizerProperties.DYNAMIC_ALLOCATION_SUSTAINED_BACKLOG_TIMEOUT,
+              sustainedBacklogTimeout,
+              OptimizerProperties.DYNAMIC_ALLOCATION_EXECUTOR_IDLE_TIMEOUT,
+              executorIdleTimeout));
+    }
   }
 
   private void requirePositive(String property, Duration value) {