chore(amber): bound region termination to a short backoff

amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala

@@ -20,7 +20,7 @@
 package org.apache.texera.amber.engine.architecture.scheduling
 
 import org.apache.pekko.pattern.gracefulStop
-import com.twitter.util.{Duration => TwitterDuration, Future, JavaTimer, Return, Throw, Timer}
+import com.twitter.util.{Future, JavaTimer, Return, Throw, Timer}
 import org.apache.texera.amber.core.state.State
 import org.apache.texera.amber.core.storage.{DocumentFactory, VFSURIFactory}
 import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
@@ -51,7 +51,7 @@ import org.apache.texera.amber.engine.architecture.scheduling.config.{
   ResourceConfig
 }
 import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.Partitioning
-import org.apache.texera.amber.engine.common.AmberLogging
+import org.apache.texera.amber.engine.common.{AmberLogging, Utils}
 import org.apache.texera.amber.engine.common.FutureBijection._
 import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient
 import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR
@@ -64,11 +64,14 @@ import scala.concurrent.duration.{Duration => ScalaDuration}
 
 object RegionExecutionManager {
 
-  // Max EndWorker retries before termination fails. ~30s at DefaultKillRetryDelay (200ms).
-  private[scheduling] val DefaultMaxTerminationAttempts: Int = 150
+  // Terminating a region is deterministic: `EndWorker` plus `gracefulStop` either succeed, or the
+  // engine has a bug that retrying cannot fix. Retries therefore only ride out a transient
+  // failure. With `Utils.retry`'s doubling backoff, 4 attempts from a 200ms base wait
+  // 200 + 400 + 800 ms = 1.4s in the worst case, where the former 150 attempts at a flat 200ms
+  // held a whole region's teardown for ~30s before the failure surfaced.
+  private[scheduling] val DefaultMaxTerminationAttempts: Int = 4
 
-  private[scheduling] val DefaultKillRetryDelay: TwitterDuration =
-    TwitterDuration.fromMilliseconds(200)
+  private[scheduling] val DefaultKillRetryBaseBackoffMs: Long = 200L
 }
 
 /**
@@ -105,8 +108,11 @@ class RegionExecutionManager(
     coordinatorConfig: CoordinatorConfig,
     actorService: PekkoActorService,
     actorRefService: PekkoActorRefMappingService,
+    // Region-termination retry budget, handed to `Utils.retry`: `maxTerminationAttempts`
+    // attempts, waiting `killRetryBaseBackoffMs` before the first retry and doubling after each.
     maxTerminationAttempts: Int = RegionExecutionManager.DefaultMaxTerminationAttempts,
-    killRetryDelay: TwitterDuration = RegionExecutionManager.DefaultKillRetryDelay,
+    killRetryBaseBackoffMs: Long = RegionExecutionManager.DefaultKillRetryBaseBackoffMs,
+    killRetryTimer: Timer = new JavaTimer(true),
     // Loop-back write addresses (Loop Start logical op id -> its input port's
     // state URI), shipped to every worker in InitializeExecutorRequest. See
     // WorkflowExecutionManager.loopStartStateUris.
@@ -125,7 +131,6 @@ class RegionExecutionManager(
     Unexecuted
   )
   private val terminationFutureRef: AtomicReference[Future[Unit]] = new AtomicReference(null)
-  private val killRetryTimer: Timer = new JavaTimer(true)
 
   /**
     * Sync the status of `RegionExecution` and transition this manager's phase to `Completed` only when the
@@ -231,38 +236,42 @@ class RegionExecutionManager(
     }
   }
 
-  private def terminateWorkersWithRetry(
-      regionExecution: RegionExecution,
-      attempt: Int = 1
-  ): Future[Unit] = {
-    terminateWorkers(regionExecution).rescue {
-      case err if attempt >= maxTerminationAttempts =>
-        val workerIds = regionExecution.getAllOperatorExecutions.flatMap {
-          case (_, opExec) => opExec.getWorkerIds
-        }.toSeq
-        val attemptsLabel = if (attempt == 1) "1 attempt" else s"$attempt attempts"
-        logger.error(
-          s"Region ${region.id.id} could not be terminated after $attemptsLabel; giving up. " +
-            s"Workers still not terminated: ${workerIds.mkString(", ")}.",
-          err
-        )
-        Future.exception(
-          new IllegalStateException(
-            s"Region ${region.id.id} could not be terminated after $attemptsLabel " +
-              s"(workers still not terminated: ${workerIds.mkString(", ")}).",
+  private def terminateWorkersWithRetry(regionExecution: RegionExecution): Future[Unit] = {
+    Utils
+      .retry(
+        attempts = maxTerminationAttempts,
+        baseBackoffTimeInMS = killRetryBaseBackoffMs,
+        timer = killRetryTimer,
+        onRetry = (err, attempt, backoffTimeInMS) =>
+          logger.warn(
+            s"Failed to terminate region ${region.id.id} on attempt $attempt of " +
+              s"$maxTerminationAttempts. Retrying in $backoffTimeInMS ms.",
             err
           )
-        )
-      case err =>
-        logger.warn(
-          s"Failed to terminate region ${region.id.id} on attempt $attempt of $maxTerminationAttempts. " +
-            s"Retrying in ${killRetryDelay.inMilliseconds} ms.",
-          err
-        )
-        Future
-          .sleep(killRetryDelay)(killRetryTimer)
-          .flatMap(_ => terminateWorkersWithRetry(regionExecution, attempt + 1))
-    }
+      ) {
+        terminateWorkers(regionExecution)
+      }
+      .rescue {
+        // Every attempt failed. Name the workers that are still alive so the user can act on it.
+        case err =>
+          val workerIds = regionExecution.getAllOperatorExecutions.flatMap {
+            case (_, opExec) => opExec.getWorkerIds
+          }.toSeq
+          val attemptsLabel =
+            if (maxTerminationAttempts <= 1) "1 attempt" else s"$maxTerminationAttempts attempts"
+          logger.error(
+            s"Region ${region.id.id} could not be terminated after $attemptsLabel; giving up. " +
+              s"Workers still not terminated: ${workerIds.mkString(", ")}.",
+            err
+          )
+          Future.exception(
+            new IllegalStateException(
+              s"Region ${region.id.id} could not be terminated after $attemptsLabel " +
+                s"(workers still not terminated: ${workerIds.mkString(", ")}).",
+              err
+            )
+          )
+      }
   }
 
   def isCompleted: Boolean = currentPhaseRef.get == Completed

amber/src/main/scala/org/apache/texera/amber/engine/common/Utils.scala

@@ -19,12 +19,13 @@
 
 package org.apache.texera.amber.engine.common
 
+import com.twitter.util.{Duration, Future, Timer}
 import com.typesafe.scalalogging.LazyLogging
 import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState
 
 import java.nio.file.{Files, Path, Paths}
 import java.util.concurrent.locks.Lock
-import scala.annotation.tailrec
+import scala.util.control.NonFatal
 
 object Utils extends LazyLogging {
 
@@ -60,32 +61,51 @@ object Utils extends LazyLogging {
   val AMBER_HOME_FOLDER_NAME = "amber";
 
   /**
-    * Retry the given logic with a backoff time interval. The attempts are executed sequentially, thus blocking the thread.
-    * Backoff time is doubled after each attempt.
+    * Retry the given logic with a backoff time interval, doubling the backoff after each attempt.
     *
-    * @param attempts            total number of attempts. if n <= 1 then it will not retry at all, decreased by 1 for each recursion.
+    * The waits are scheduled on `timer`, so no thread is blocked between attempts. That matters
+    * for callers on an actor or coordinator thread, where a `Thread.sleep` would also stall
+    * unrelated work queued on that thread.
+    *
+    * A synchronous throw while evaluating `fn` counts as a failed attempt, same as a failed
+    * `Future`. Fatal errors are never retried in either shape: they propagate immediately.
+    *
+    * @param attempts            total number of attempts. if n <= 1 then it will not retry at all.
     * @param baseBackoffTimeInMS time to wait before next attempt, started with the base time, and doubled after each attempt.
-    * @param fn                  the target function to execute.
-    * @tparam T any return type from the provided function fn.
-    * @return the provided function fn's return, or any exception that still being raised after n attempts.
+    * @param timer               schedules the waits between attempts; no thread is blocked.
+    * @param onRetry             invoked before each wait with the failure, the 1-based number of the
+    *                            attempt that just failed, and the backoff about to be waited in ms.
+    *                            Override it to log with caller context.
+    * @param fn                  the target function to execute, re-evaluated on each attempt.
+    * @tparam T any value type the provided function fn's `Future` yields.
+    * @return `fn`'s eventual value, or the last failure once `attempts` is exhausted.
     */
-  @tailrec
-  def retry[T](attempts: Int, baseBackoffTimeInMS: Long)(fn: => T): T = {
-    try {
-      fn
-    } catch {
-      case e: Throwable =>
-        if (attempts > 1) {
-          logger.warn(
-            "retrying after " + baseBackoffTimeInMS + "ms, number of attempts left: " + (attempts - 1),
-            e
-          )
-          Thread.sleep(baseBackoffTimeInMS)
-          retry(attempts - 1, baseBackoffTimeInMS * 2)(fn)
-        } else throw e
-    }
+  def retry[T](
+      attempts: Int,
+      baseBackoffTimeInMS: Long,
+      timer: Timer,
+      onRetry: (Throwable, Int, Long) => Unit = logRetryAttempt
+  )(fn: => Future[T]): Future[T] = {
+    def attempt(attemptNumber: Int, backoffTimeInMS: Long): Future[T] =
+      Future(fn).flatten.rescue {
+        // `NonFatal` so that a fatal handed back as a failed `Future` is not retried either, which
+        // also matches how the blocking backoff loops elsewhere in the repo catch `Exception`.
+        case NonFatal(e) if attemptNumber < attempts =>
+          onRetry(e, attemptNumber, backoffTimeInMS)
+          Future
+            .sleep(Duration.fromMilliseconds(backoffTimeInMS))(timer)
+            .flatMap(_ => attempt(attemptNumber + 1, backoffTimeInMS * 2))
+      }
+
+    attempt(attemptNumber = 1, backoffTimeInMS = baseBackoffTimeInMS)
   }
 
+  private def logRetryAttempt(failure: Throwable, attempt: Int, backoffTimeInMS: Long): Unit =
+    logger.warn(
+      "retrying after " + backoffTimeInMS + "ms, attempts made so far: " + attempt,
+      failure
+    )
+
   private def isAmberHomePath(path: Path): Boolean = {
     path.toRealPath().endsWith(AMBER_HOME_FOLDER_NAME)
   }