[SPARK-52581] Revise AppCleanUpStep to include cleanup logic for all states

### What changes were proposed in this pull request?

This PR updates AppCleanUpStep so that it can be added as a regular step of each reconciliation, remove the therefore-unused AppTerminatedStep

### Why are the changes needed?

Previously, we have AppTerminatedStep performed at the beginning of each reconciliation, to end the actions early if app has reached terminated state.

However, in some scenarios, we may still want to do "clean up" for applications that are terminated:

* If an app in "terminated without releasing resources" is being deleted, we would like the AppCleanUpStep to be performed, since it covers a few corner cases that might not be covered in general delete flow
* We may also to introduce the concept of "maximal retention duration" for applications along with current "retention policy" in near future to serve the lifecycle management of application secondary resources. That could also require cleaning up secondary resources for already-terminated applications.

The revised AppCleanUpStep includes everything in the previous AppTerminatedStep and takes care of cleaning up resources for those in "TerminatedWithoutReleaseResources" stage as needed.

### Does this PR introduce any user-facing change?

No - internal step(s) changes only

### How was this patch tested?

CIs and E2Es shall cover the clean up logic

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #257 from jiangzho/cleanupForCancel.

Authored-by: Zhou JIANG <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>

Author: jiangzho

spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/SparkAppReconciler.java

@@ -56,7 +56,6 @@
 import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppReconcileStep;
 import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppResourceObserveStep;
 import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppRunningStep;
-import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppTerminatedStep;
 import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppUnknownStateStep;
 import org.apache.spark.k8s.operator.reconciler.reconcilesteps.AppValidateStep;
 import org.apache.spark.k8s.operator.utils.LoggingUtils;
@@ -144,7 +143,7 @@ public List<EventSource<?, SparkApplication>> prepareEventSources(
   protected List<AppReconcileStep> getReconcileSteps(final SparkApplication app) {
     List<AppReconcileStep> steps = new ArrayList<>();
     steps.add(new AppValidateStep());
-    steps.add(new AppTerminatedStep());
+    steps.add(new AppCleanUpStep());
     switch (app.getStatus().getCurrentState().getCurrentStateSummary()) {
       case Submitted, ScheduledToRestart -> steps.add(new AppInitStep());
       case DriverRequested, DriverStarted -> {
@@ -166,14 +165,6 @@ protected List<AppReconcileStep> getReconcileSteps(final SparkApplication app) {
         steps.add(
             new AppResourceObserveStep(Collections.singletonList(new AppDriverTimeoutObserver())));
       }
-      case DriverReadyTimedOut,
-              DriverStartTimedOut,
-              ExecutorsStartTimedOut,
-              Succeeded,
-              DriverEvicted,
-              Failed,
-              SchedulingFailure ->
-          steps.add(new AppCleanUpStep());
       default -> steps.add(new AppUnknownStateStep());
     }
     return steps;
@@ -196,7 +187,6 @@ public DeleteControl cleanup(
       SparkAppContext ctx = new SparkAppContext(sparkApplication, context, submissionWorker);
       List<AppReconcileStep> cleanupSteps = new ArrayList<>();
       cleanupSteps.add(new AppValidateStep());
-      cleanupSteps.add(new AppTerminatedStep());
       cleanupSteps.add(new AppCleanUpStep(SparkAppStatusUtils::appCancelled));
       for (AppReconcileStep step : cleanupSteps) {
         ReconcileProgress progress = step.reconcile(ctx, sparkAppStatusRecorder);

spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppCleanUpStep.java

@@ -23,6 +23,8 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.function.Supplier;
 
@@ -54,38 +56,78 @@
 @NoArgsConstructor
 @Slf4j
 public class AppCleanUpStep extends AppReconcileStep {
-  private Supplier<ApplicationState> cleanUpSuccessStateSupplier;
+  private Supplier<ApplicationState> onDemandCleanUpReason;
 
+  /**
+   * Cleanup secondary resources for an application if needed and updates application status
+   * accordingly. This step would be performed right after validation step in each reconcile as a
+   * sanity check. It may end the reconciliation if no more actions are needed. In addition, it can
+   * be performed on demand with a reason for cleanup secondary resources.
+   *
+   * <p>An app expects its secondary resources to be released if any of the below is true:
+   *
+   * <ul>
+   *   <li>When the application is being deleted on demand(e.g. being deleted) with a reason
+   *   <li>When the application is stopping
+   *   <li>When the application has terminated without releasing resources, but it has exceeded
+   *       configured retention duration
+   * </ul>
+   *
+   * <p>It would proceed to next steps with no actions for application in other states. Note that
+   * when even the reconciler decides to proceed with clean up, sub-resources may still be retained
+   * based on tolerations.
+   *
+   * @param context context for the app
+   * @param statusRecorder recorder for status updates in this reconcile
+   * @return the reconcile progress
+   */
   @Override
   public ReconcileProgress reconcile(
       SparkAppContext context, SparkAppStatusRecorder statusRecorder) {
-    ApplicationStatus currentStatus = context.getResource().getStatus();
-    ApplicationTolerations tolerations =
-        context.getResource().getSpec().getApplicationTolerations();
-    ResourceRetainPolicy resourceRetainPolicy = tolerations.getResourceRetainPolicy();
+    SparkApplication application = context.getResource();
+    ApplicationStatus currentStatus = application.getStatus();
+    ApplicationState currentState = currentStatus.getCurrentState();
+    ApplicationTolerations tolerations = application.getSpec().getApplicationTolerations();
+    if (ApplicationStateSummary.ResourceReleased.equals(currentState.getCurrentStateSummary())) {
+      statusRecorder.removeCachedStatus(application);
+      return ReconcileProgress.completeAndNoRequeue();
+    }
     String stateMessage = null;
-
-    if (retainReleaseResource(resourceRetainPolicy, currentStatus.getCurrentState())) {
-      if (tolerations.getRestartConfig() != null
-          && !RestartPolicy.Never.equals(tolerations.getRestartConfig().getRestartPolicy())) {
-        stateMessage =
-            "Application is configured to restart, resources created in current "
-                + "attempt would be force released.";
-        log.warn(stateMessage);
+    if (isOnDemandCleanup()) {
+      log.info("Cleaning up application resources on demand");
+    } else {
+      if (ApplicationStateSummary.TerminatedWithoutReleaseResources.equals(
+          currentState.getCurrentStateSummary())) {
+        statusRecorder.removeCachedStatus(application);
+        return ReconcileProgress.completeAndNoRequeue();
+      } else if (currentState.getCurrentStateSummary().isStopping()) {
+        if (retainReleaseResourceForPolicyAndState(
+            tolerations.getResourceRetainPolicy(), currentState)) {
+          if (tolerations.getRestartConfig() != null
+              && !RestartPolicy.Never.equals(tolerations.getRestartConfig().getRestartPolicy())) {
+            stateMessage =
+                "Application is configured to restart, resources created in current "
+                    + "attempt would be force released.";
+            log.warn(stateMessage);
+          } else {
+            ApplicationState terminationState =
+                new ApplicationState(
+                    ApplicationStateSummary.TerminatedWithoutReleaseResources,
+                    "Application is terminated without releasing resources as configured.");
+            long requeueAfterMillis =
+                tolerations.getApplicationTimeoutConfig().getTerminationRequeuePeriodMillis();
+            return appendStateAndRequeueAfter(
+                context, statusRecorder, terminationState, Duration.ofMillis(requeueAfterMillis));
+          }
+        }
       } else {
-        ApplicationState terminationState =
-            new ApplicationState(
-                ApplicationStateSummary.TerminatedWithoutReleaseResources,
-                "Application is terminated without releasing resources as configured.");
-        long requeueAfterMillis =
-            tolerations.getApplicationTimeoutConfig().getTerminationRequeuePeriodMillis();
-        return appendStateAndRequeueAfter(
-            context, statusRecorder, terminationState, Duration.ofMillis(requeueAfterMillis));
+        log.debug("Clean up is not expected for app, proceeding to next step.");
+        return ReconcileProgress.proceed();
       }
     }
+
     List<HasMetadata> resourcesToRemove = new ArrayList<>();
-    if (ApplicationStateSummary.SchedulingFailure.equals(
-        currentStatus.getCurrentState().getCurrentStateSummary())) {
+    if (isReleasingResourcesForSchedulingFailureAttempt(currentStatus)) {
       // if app failed at scheduling, re-compute all spec and delete as they may not be fully
       // owned by driver
       try {
@@ -110,13 +152,13 @@ public ReconcileProgress reconcile(
       Optional<Pod> driver = context.getDriverPod();
       driver.ifPresent(resourcesToRemove::add);
     }
-    boolean forceDelete = enableForceDelete(context.getResource());
+    boolean forceDelete = enableForceDelete(application);
     for (HasMetadata resource : resourcesToRemove) {
       ReconcilerUtils.deleteResourceIfExists(context.getClient(), resource, forceDelete);
     }
     ApplicationStatus updatedStatus;
-    if (cleanUpSuccessStateSupplier != null) {
-      ApplicationState state = cleanUpSuccessStateSupplier.get();
+    if (onDemandCleanUpReason != null) {
+      ApplicationState state = onDemandCleanUpReason.get();
       if (StringUtils.isNotEmpty(stateMessage)) {
         state.setMessage(stateMessage);
       }
@@ -142,7 +184,26 @@ public ReconcileProgress reconcile(
     }
   }
 
-  protected boolean retainReleaseResource(
+  protected boolean isOnDemandCleanup() {
+    return onDemandCleanUpReason != null;
+  }
+
+  protected boolean isReleasingResourcesForSchedulingFailureAttempt(
+      final ApplicationStatus status) {
+    ApplicationState lastObservedState = status.getCurrentState();
+    if (ApplicationStateSummary.TerminatedWithoutReleaseResources.equals(
+        lastObservedState.getCurrentStateSummary())) {
+      // if the app has already terminated, use the last observed state before termination
+      NavigableMap<Long, ApplicationState> navMap =
+          (NavigableMap<Long, ApplicationState>) status.getStateTransitionHistory();
+      Map.Entry<Long, ApplicationState> terminateState = navMap.lastEntry();
+      lastObservedState = navMap.lowerEntry(terminateState.getKey()).getValue();
+    }
+    return ApplicationStateSummary.SchedulingFailure.equals(
+        lastObservedState.getCurrentStateSummary());
+  }
+
+  protected boolean retainReleaseResourceForPolicyAndState(
       ResourceRetainPolicy resourceRetainPolicy, ApplicationState currentState) {
     return switch (resourceRetainPolicy) {
       case Never -> false;