diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandlerSpec.scala new file mode 100644 index 00000000000..a942c615aa4 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandlerSpec.scala @@ -0,0 +1,247 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.amber.engine.architecture.coordinator.promisehandlers + +import com.twitter.util.{Await, Duration, Future, Promise} +import org.apache.pekko.actor.ActorSystem +import org.apache.pekko.testkit.TestKit +import org.apache.texera.amber.core.WorkflowRuntimeException +import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity} +import org.apache.texera.amber.core.workflow.WorkflowContext +import org.apache.texera.amber.engine.architecture.common.PekkoActorService +import org.apache.texera.amber.engine.architecture.coordinator.execution.WorkflowExecution +import org.apache.texera.amber.engine.architecture.coordinator.{ + ClientEvent, + CoordinatorAsyncRPCHandlerInitializer, + CoordinatorConfig, + CoordinatorProcessor, + FatalError +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + EmptyRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} +import org.apache.texera.amber.engine.architecture.rpc.coordinatorservice.CoordinatorServiceGrpc.METHOD_COORDINATOR_INITIATE_ADVANCE_REGION_EXECUTIONS +import org.apache.texera.amber.engine.architecture.scheduling.{ + RegionExecutionManagerTestSupport, + WorkflowExecutionManager +} +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage +import org.apache.texera.amber.engine.common.AmberRuntime +import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage +import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient +import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation +import org.apache.texera.amber.engine.common.virtualidentity.util.{CLIENT, COORDINATOR} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpecLike + +import scala.collection.mutable.ArrayBuffer + +/** + * `coordinatorInitiateAdvanceRegionExecutions` is how the coordinator advances its region + * executions in a control round of its own: a handler that must not advance inline (see + * `PortCompletedHandler`) sends this request to the coordinator itself, so that the `EndWorker` + * messages a completed region produces cannot overtake the replies that round still owed. + * + * The behaviors this spec pins down: + * - the advance is delegated to the coordinator's own `WorkflowExecutionManager`; + * - the reply is produced synchronously, WITHOUT awaiting the advance. The requester is the + * coordinator itself and discards the reply, so awaiting it would only hold the control round + * open; + * - because nothing awaits the reply, a failing advance has no caller to propagate to and is + * reported to the client as a `FatalError` instead — carrying the related worker id when the + * failure came from a worker RPC (`WorkflowRuntimeException`) and `None` for anything else. + */ +class AdvanceRegionExecutionsHandlerSpec + extends TestKit(ActorSystem("AdvanceRegionExecutionsHandlerSpec", AmberRuntime.pekkoConfig)) + with AnyFlatSpecLike + with BeforeAndAfterAll + with RegionExecutionManagerTestSupport { + + private val awaitTimeout = Duration.fromSeconds(5) + private val coordinatorConfig = CoordinatorConfig(None, None, None, None) + private val ctx: AsyncRPCContext = AsyncRPCContext(COORDINATOR, COORDINATOR) + private val advanceCommandId = 11L + private val relatedWorker = ActorVirtualIdentity("Worker:unschedulable-region-worker") + + override def afterAll(): Unit = { + TestKit.shutdownActorSystem(system) + } + + /** + * Stands in for the real `WorkflowExecutionManager`: driving its advance to a *failed* future + * from the outside means staging a whole region launch or termination, and the causes that + * reaches the handler with are `IllegalStateException`s (unresolved output schema, termination + * retries exhausted) or the bare `Throwable` that `ErrorUtils.reconstructThrowable` builds for a + * worker RPC error — never a `WorkflowRuntimeException` carrying a worker id. This stub records + * the advance calls and hands back a future the test controls, so the handler can be observed + * both while the advance is still pending and after it has failed with either kind of cause. + */ + private class PendingAdvanceExecutionManager( + workflowExecution: WorkflowExecution, + asyncRPCClient: AsyncRPCClient + ) extends WorkflowExecutionManager(workflowExecution, coordinatorConfig, asyncRPCClient) { + val advanceCalls: ArrayBuffer[PekkoActorService] = ArrayBuffer() + private val advance: Promise[Unit] = Promise[Unit]() + + override def advanceRegionExecutions(actorService: PekkoActorService): Future[Unit] = { + advanceCalls += actorService + advance + } + + /** Twitter futures run their continuations on the thread that satisfies them, so the + * handler's `onFailure` has already run once this returns — no polling needed. + */ + def failAdvance(cause: Throwable): Unit = advance.setException(cause) + + def finishAdvance(): Unit = advance.setValue(()) + } + + private class TestCoordinatorProcessor( + outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit + ) extends CoordinatorProcessor( + new WorkflowContext(), + coordinatorConfig, + COORDINATOR, + outputHandler + ) { + override val workflowExecutionManager: PendingAdvanceExecutionManager = + new PendingAdvanceExecutionManager(workflowExecution, asyncRPCClient) + } + + private case class Fixture( + cp: TestCoordinatorProcessor, + init: CoordinatorAsyncRPCHandlerInitializer, + sent: ArrayBuffer[WorkflowFIFOMessage] + ) { + def manager: PendingAdvanceExecutionManager = cp.workflowExecutionManager + } + + /** + * A real `CoordinatorProcessor` (and therefore the real RPC layer) with the execution manager + * replaced by the controllable stub, and a real `PekkoActorService` set up so the argument the + * handler forwards is something other than `null`. + */ + private def newFixture(): Fixture = { + val sent = ArrayBuffer[WorkflowFIFOMessage]() + val cp = new TestCoordinatorProcessor({ + case Right(msg) => sent += msg + case Left(_) => () + }) + cp.setupActorService(createCoordinatorHarness().actorService) + Fixture(cp, new CoordinatorAsyncRPCHandlerInitializer(cp), sent) + } + + private def clientEvents(sent: ArrayBuffer[WorkflowFIFOMessage]): Seq[ClientEvent] = + sent.toSeq.filter(_.channelId.toWorkerId == CLIENT).map(_.payload).collect { + case event: ClientEvent => event + } + + private def repliesToCoordinator(sent: ArrayBuffer[WorkflowFIFOMessage]): Seq[ReturnInvocation] = + sent.toSeq.filter(_.channelId.toWorkerId == COORDINATOR).map(_.payload).collect { + case ret: ReturnInvocation => ret + } + + behavior of "AdvanceRegionExecutionsHandler" + + it should "advance the region executions with the coordinator's own actor service" in { + val fixture = newFixture() + + fixture.init.coordinatorInitiateAdvanceRegionExecutions(EmptyRequest(), ctx) + + // The actor service is what the advance needs to create the next region's workers on, so the + // coordinator's own one must be forwarded; guarded against a vacuous null == null match. + assert(fixture.cp.actorService != null) + assert(fixture.manager.advanceCalls.toSeq == Seq(fixture.cp.actorService)) + } + + it should "reply immediately instead of awaiting the advance" in { + val fixture = newFixture() + + val response = fixture.init.coordinatorInitiateAdvanceRegionExecutions(EmptyRequest(), ctx) + + // The advance was started but is still pending, and the reply is already satisfied: this + // handler deliberately does not chain its reply on the advance. Chaining it would keep the + // control round open for the whole advance, which itself waits on region termination RPCs. + assert(fixture.manager.advanceCalls.size == 1) + assert(response.isDefined) + assert(Await.result(response, awaitTimeout) == EmptyReturn()) + } + + it should "reply to the requesting round while the advance is still running" in { + val fixture = newFixture() + + // Delivered the way production does it: as a coordinator-to-coordinator control message, + // dispatched by the real RPC server. No other spec drives this handler that way — + // `PortCompletedHandlerSpec` only asserts that the invocation is sent. + fixture.cp.processDCM( + ChannelIdentity(COORDINATOR, COORDINATOR, isControl = true), + ControlInvocation( + METHOD_COORDINATOR_INITIATE_ADVANCE_REGION_EXECUTIONS, + EmptyRequest(), + ctx, + advanceCommandId + ) + ) + + assert(fixture.manager.advanceCalls.size == 1) + assert( + repliesToCoordinator(fixture.sent) == Seq( + ReturnInvocation(advanceCommandId, EmptyReturn()) + ) + ) + } + + it should "notify the client of a failed advance with the related worker id" in { + val fixture = newFixture() + fixture.init.coordinatorInitiateAdvanceRegionExecutions(EmptyRequest(), ctx) + + val failure = new WorkflowRuntimeException("region cannot be scheduled", Some(relatedWorker)) + fixture.manager.failAdvance(failure) + + // A `WorkflowRuntimeException` from a worker RPC knows which worker it came from, and the + // client needs that id to attribute the error to an operator. + assert(clientEvents(fixture.sent) == Seq(FatalError(failure, Some(relatedWorker)))) + } + + it should "notify the client of a failed advance without a worker id for other failures" in { + val fixture = newFixture() + fixture.init.coordinatorInitiateAdvanceRegionExecutions(EmptyRequest(), ctx) + + val failure = new IllegalStateException("no resource config for the next region") + fixture.manager.failAdvance(failure) + + // Not every advance failure comes from a worker, so there is no id to attribute. + assert(clientEvents(fixture.sent) == Seq(FatalError(failure, None))) + } + + it should "not notify the client when the advance succeeds" in { + val fixture = newFixture() + fixture.init.coordinatorInitiateAdvanceRegionExecutions(EmptyRequest(), ctx) + + fixture.manager.finishAdvance() + + assert(clientEvents(fixture.sent).isEmpty) + } +}