This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6960-7d5276de1e7a6ca6f7b8dc92a316ffa0a4ba320a in repository https://gitbox.apache.org/repos/asf/texera.git
commit 226614b527bf51a010355778440e1972ed3130a7 Author: Xinyuan Lin <[email protected]> AuthorDate: Wed Jul 29 18:37:25 2026 -0700 fix(amber): advance region executions in a later coordinator round, not inside portCompleted (#6960) ### What changes were proposed in this PR? **Root cause of the `LoopIntegrationSpec` retry storm.** `portCompleted`'s handler advanced region executions inline in its own continuation. That advance terminates a completed region, which sends `EndWorker` to each of its workers — so `EndWorker` went out *before* the reply to the very `portCompleted` being handled, and both travel the same FIFO control channel to that worker: ``` Before — all inside one handler round: advanceRegionExecutions() ──> EndWorker to W (seq N on W's control channel) return EmptyReturn() ──> ReturnInvocation to W (seq N+1, same channel) W processes EndWorker while the reply is still queued behind it: EndHandler: queue not empty => IllegalStateException("worker still has unprocessed messages") => attempt 1 of 150 fails => 2 ERRORs + 2 WARNs + 2 stack traces => 200 ms retry succeeds ``` The check's premise — that `EndWorker` is the last message a worker receives — was being violated by the coordinator itself on every teardown. Whether the worker actually observed the trailing reply was a microsecond race, hence the ~1-in-10 flake. **Fix: advance in a later coordinator round instead of inline.** A new coordinator-to-coordinator RPC carries the advance: ``` After: portCompleted round: bookkeeping; send CoordinatorInitiateAdvanceRegionExecutions to self return EmptyReturn() ──> ReturnInvocation to W later round: advance runs ──> EndWorker to W ``` A message the coordinator addresses to itself is transmitted and then received before it is handled, so it lands a whole message behind every reply the requesting round still owed — `EndWorker` can no longer overtake a reply, and the first termination attempt succeeds. **Second half: the worker tolerates a reply it cannot be ordered against.** Deferring orders the replies the coordinator has *already produced*, but region completion keys only on **port** completion, so the advance fires as soon as the last `portCompleted` is booked — while the same worker's `workerExecutionCompleted` (emitted right after it) may still be queued at the coordinator. Since `NetworkInputGateway.tryPickControlChannel` selects channels out of a `HashMap`, there is no cross-channel order: the advance can run first and that reply be emitted afterwards, landing behind `EndWorker`. Ordering cannot close this, so `EndWorker` no longer counts a queued `ReturnInvocation` as unprocessed work — processing one only fulfills a promise for a request the worker already issued, and all four worker→coordinator calls discard their futures, so nothing is pending on it. Every other queued element still fails the request, so the coordinator retries and the worker drains it first. This PR is JVM-only: the Python worker keeps its current check, which fails `EndWorker` whenever anything is queued, and the coordinator's retry covers it exactly as it does today. Bringing it to parity needs a way to classify what is queued without inspecting the next message, which means changing `InternalQueue` — left for a follow-up rather than bundled here. | | Before | After | |---|---|---| | Coordinator send order | `EndWorker`, then the `portCompleted` reply | reply, then `EndWorker` | | First `EndWorker` attempt | fails on a benign race (storm ~1 run in 10) | succeeds | | Normal teardown logs | 2 ERROR + 2 WARN + 2 stack traces | none | | Worker-side check | any queued message fails | a queued **`ReturnInvocation`** no longer fails; all real work still does | | Bookkeeping failure | `ControlError` to the reporting worker | unchanged | | Advance failure | `FatalError` to the client | unchanged | Changes: - **`coordinatorservice.proto`**: adds `CoordinatorInitiateAdvanceRegionExecutions`, alongside the existing self-directed `CoordinatorInitiateQueryStatistics`. - **`AdvanceRegionExecutionsHandler`** (new): handles that self-addressed request by running the advance, and reports failures to the client exactly as the inline call did. `advanceRegionExecutions`'s scaladoc now tells callers handling a worker-initiated request to go through it. - **`PortCompletedHandler`**: keeps its bookkeeping and its reply position; only the advance moves. - **`RegionExecutionManager`**: corrects the scaladoc that claimed `EndWorker` "will be the last message each worker receives". `StartWorkflowHandler` still advances directly: its caller is the client, it awaits the future to answer with the workflow state, and no `EndWorker` is involved at startup. ### Any related issues, documentation, discussions? #6891. ### How was this PR tested? - **`PortCompletedHandlerSpec` (new)**: drives a real `CoordinatorProcessor` through the real RPC server with a captured output gateway, over a single-region schedule already in flight. Pins that the advance leaves as a separate coordinator-addressed control message; that the *only* thing sent to the reporting worker while handling `portCompleted` is its reply (the direct regression guard — on `main` this list also contains `EndWorker`); that the port is recorded completed before the request goes out, since the advance now reads that state in a later round; that a port belonging to no executing region requests nothing; and that a failed continuation still returns a `ControlError` to the reporting worker. - Baseline check: 3 of those tests fail against `main`'s handler, while the `ControlError` test passes on **both** `main` and this branch — which is what verifies the worker-facing error contract is preserved rather than merely asserted. - Untouched and green, as evidence termination semantics are unchanged: `RegionExecutionManagerSpec`, `WorkflowExecutionManagerSpec`, `EndHandlerSpec`, `CoordinatorSpec`, `TrivialControlSpec`, `AsyncRPCClientSpec`. - Lint: `scalafmtCheck`, `scalafixAll --check` pass. - Loop confirmation is CI-side and statistical: `amber-integration` teardown logs should contain zero `Failed to terminate region ... on attempt 1` lines (runs without this fix show them repeatedly). ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Fable 5) --- .../architecture/rpc/coordinatorservice.proto | 1 + .../CoordinatorAsyncRPCHandlerInitializer.scala | 1 + .../AdvanceRegionExecutionsHandler.scala | 64 +++++ .../promisehandlers/PortCompletedHandler.scala | 28 +- .../scheduling/RegionExecutionManager.scala | 12 +- .../scheduling/WorkflowExecutionManager.scala | 4 + .../worker/promisehandlers/EndHandler.scala | 63 ++++- .../promisehandlers/PortCompletedHandlerSpec.scala | 282 +++++++++++++++++++++ .../worker/promisehandlers/EndHandlerSpec.scala | 69 ++++- 9 files changed, 501 insertions(+), 23 deletions(-) diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/coordinatorservice.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/coordinatorservice.proto index b782486c08..2ccaabbe26 100644 --- a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/coordinatorservice.proto +++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/coordinatorservice.proto @@ -45,6 +45,7 @@ service CoordinatorService { rpc JumpToOperatorRegion(JumpToOperatorRegionRequest) returns (EmptyReturn); rpc LinkWorkers(LinkWorkersRequest) returns (EmptyReturn); rpc CoordinatorInitiateQueryStatistics(QueryStatisticsRequest) returns (EmptyReturn); + rpc CoordinatorInitiateAdvanceRegionExecutions(EmptyRequest) returns (EmptyReturn); rpc RetryWorkflow(RetryWorkflowRequest) returns (EmptyReturn); rpc ReconfigureWorkflow(WorkflowReconfigureRequest) returns (EmptyReturn); } diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorAsyncRPCHandlerInitializer.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorAsyncRPCHandlerInitializer.scala index 0f8f9c8bfc..131d0b7a82 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorAsyncRPCHandlerInitializer.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/CoordinatorAsyncRPCHandlerInitializer.scala @@ -41,6 +41,7 @@ class CoordinatorAsyncRPCHandlerInitializer( with ResumeHandler with StartWorkflowHandler with PortCompletedHandler + with AdvanceRegionExecutionsHandler with ConsoleMessageHandler with RetryWorkflowHandler with EvaluatePythonExpressionHandler diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandler.scala new file mode 100644 index 0000000000..ec0e5b0030 --- /dev/null +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/AdvanceRegionExecutionsHandler.scala @@ -0,0 +1,64 @@ +/* + * 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.Future +import org.apache.texera.amber.core.WorkflowRuntimeException +import org.apache.texera.amber.engine.architecture.coordinator.{ + CoordinatorAsyncRPCHandlerInitializer, + FatalError +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + EmptyRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn + +/** Advance the region executions of this workflow. + * + * A handler that needs region executions advanced sends this to itself rather than advancing + * inline (see `PortCompletedHandler`), so that the advance — which terminates completed regions + * and therefore sends `EndWorker` to their workers — runs in its own control round, after every + * reply the requesting round owed. + * + * possible sender: coordinator (itself) + */ +trait AdvanceRegionExecutionsHandler { + this: CoordinatorAsyncRPCHandlerInitializer => + + override def coordinatorInitiateAdvanceRegionExecutions( + request: EmptyRequest, + ctx: AsyncRPCContext + ): Future[EmptyReturn] = { + cp.workflowExecutionManager + .advanceRegionExecutions(cp.actorService) + // The requester is the coordinator itself and discards this reply, so a failure has no + // caller to propagate to. A fatal error is sent to the client, indicating that the region + // cannot be scheduled. + .onFailure { + case err: WorkflowRuntimeException => + sendToClient(FatalError(err, err.relatedWorkerId)) + case other => + sendToClient(FatalError(other, None)) + } + EmptyReturn() + } + +} diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandler.scala index 4815f3a140..458dbc880a 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandler.scala @@ -20,14 +20,11 @@ package org.apache.texera.amber.engine.architecture.coordinator.promisehandlers import com.twitter.util.Future -import org.apache.texera.amber.core.WorkflowRuntimeException import org.apache.texera.amber.core.workflow.GlobalPortIdentity -import org.apache.texera.amber.engine.architecture.coordinator.{ - CoordinatorAsyncRPCHandlerInitializer, - FatalError -} +import org.apache.texera.amber.engine.architecture.coordinator.CoordinatorAsyncRPCHandlerInitializer import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ AsyncRPCContext, + EmptyRequest, PortCompletedRequest, QueryStatisticsRequest, StatisticsUpdateTarget @@ -81,16 +78,17 @@ trait PortCompletedHandler { else operatorExecution.isOutputPortCompleted(msg.portId) if (isPortCompleted) { - cp.workflowExecutionManager - .advanceRegionExecutions(cp.actorService) - // Since this message is sent from a worker, any exception from the above code will be returned to that worker. - // Additionally, a fatal error is sent to the client, indicating that the region cannot be scheduled. - .onFailure { - case err: WorkflowRuntimeException => - sendToClient(FatalError(err, err.relatedWorkerId)) - case other => - sendToClient(FatalError(other, None)) - } + // Advance region executions in a later control round instead of here. Advancing + // inline terminates the completed region and sends `EndWorker` to this very sender + // before this handler's own reply, on the same control channel — the worker would + // then process `EndWorker` with the reply still queued behind it and reject the + // termination (see `EndHandler`). A message the coordinator addresses to itself is + // transmitted and received before it is handled, so the advance lands behind the + // reply below. + coordinatorInterface.coordinatorInitiateAdvanceRegionExecutions( + EmptyRequest(), + COORDINATOR + ) } case None => // currently "start" and "end" ports are not part of a region, thus no region can be found. // do nothing. diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala index 93d085b611..3a87f8124e 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala @@ -134,8 +134,16 @@ class RegionExecutionManager( * * Additionally, this method will also terminate all the workers of this region: * - * 1. An `EndWorker` control message is first sent to all the workers. This will be the last message each worker - * receives. We wait for all workers have replied to indicate they have finished processing all control messages. + * 1. An `EndWorker` control message is first sent to all the workers. We wait for all workers to reply that they + * have finished processing all control messages; a worker that has not fails the request, and + * `terminateWorkersWithRetry` re-sends `EndWorker` after `killRetryDelay`. + * + * Because a worker rejects `EndWorker` while work is still queued for it, termination must not be triggered from + * inside the handler of a request sent by one of these workers — the `EndWorker` would be emitted before that + * handler's own reply and overtake it on their shared FIFO control channel. Such handlers therefore send themselves + * a `CoordinatorInitiateAdvanceRegionExecutions` instead of advancing inline, which defers the advance that reaches + * here to a later coordinator round (see `PortCompletedHandler`). That orders the replies already produced; for one + * the coordinator has not produced yet, `EndHandler` does not count a queued reply as work. * * 2. Only after all workers have processed all control messages do we send a `gracefulStop` (pekko message) to each * worker. JVM workers will be terminated by `gracefulStop`. Python proxy workes will also be terminated by diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala index 65afbea8e1..0b63d1f17c 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala @@ -96,6 +96,10 @@ class WorkflowExecutionManager( * in `Completed` status (phase). * * After the syncs, if there are no running region(s), it will start new regions (if available). + * + * Callers handling a worker-initiated request must not call this directly; they send themselves a + * `CoordinatorInitiateAdvanceRegionExecutions` (see `PortCompletedHandler`) so the resulting + * `EndWorker` cannot overtake the reply that request still owes. */ def advanceRegionExecutions(actorService: PekkoActorService): Future[Unit] = { val unfinishedRegionManagers = diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala index f91a28c627..ff16305c17 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala @@ -24,8 +24,15 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ AsyncRPCContext, EmptyRequest } -import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{ + DPInputQueueElement, + FIFOMessageElement +} /** * The EndWorker control messages is needed to ensure all the other control messages in a worker @@ -37,20 +44,66 @@ trait EndHandler { /** * The response of endWorker to the coordinator indicates that this worker has finished not only * the data processing logic, but also the processing of all the control messages. + * + * A queued reply to one of this worker's own requests does not count: it carries no work (see + * `findUnprocessedWork`), and the coordinator cannot order all of its replies before `EndWorker`. + * The coordinator defers region advancement to a later control round precisely so that + * `EndWorker` follows the replies it owes (see `PortCompletedHandler`), but that only orders + * requests it has already handled. A request of this worker's that is still queued at the + * coordinator when the deferred advance runs — `workerExecutionCompleted`, which the worker + * emits right after its last `portCompleted` — is replied to afterwards, and the coordinator + * picks its input channels out of a `HashMap` (`NetworkInputGateway.tryPickControlChannel`), so + * there is no cross-channel order to rely on. */ override def endWorker( request: EmptyRequest, ctx: AsyncRPCContext ): Future[EmptyReturn] = { - // Ensure this is really the last message. - if (!dp.inputManager.inputMessageQueue.isEmpty) { + // Ensure this is really the last message that asks this worker to do anything. + val pendingWork = findUnprocessedWork + if (pendingWork.isDefined) { logger.warn( - s"Received EndHandler before all messages are processed. Unprocessed messages: " + - s"${dp.inputManager.inputMessageQueue.peek()}" + s"Received EndHandler before all messages are processed. Unprocessed message: " + + s"${describe(pendingWork.get)}" ) return Future.exception(new IllegalStateException("worker still has unprocessed messages")) } // Now we can safely acknowledge that this worker can be terminated. EmptyReturn() } + + /** + * The first queued element that represents work, if any. + * + * A `ReturnInvocation` is excluded: processing one only fulfills a promise for a request this + * worker already issued (`AmberProcessor.processDCM`), and every worker-to-coordinator call + * discards its future, so nothing is pending on it. Everything else — control invocations, data, + * embedded control messages, timer-based controls, actor commands — still blocks termination, so + * the coordinator retries and this worker drains it first. + * + * This is the worker's own arrival queue rather than its actor mailbox, and at termination it + * holds at most a couple of elements. + */ + private def findUnprocessedWork: Option[DPInputQueueElement] = { + val iterator = dp.inputManager.inputMessageQueue.iterator() + while (iterator.hasNext) { + val element = iterator.next() + val isWork = element match { + case FIFOMessageElement(message) => !message.payload.isInstanceOf[ReturnInvocation] + case _ => true + } + if (isWork) { + return Some(element) + } + } + None + } + + /** Identifies a queued element without logging payload contents. */ + private def describe(element: DPInputQueueElement): String = + element match { + case FIFOMessageElement(message) => + s"${message.payload.getClass.getSimpleName} on ${message.channelId}" + case other => other.getClass.getSimpleName + } } diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandlerSpec.scala new file mode 100644 index 0000000000..4898762acf --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/PortCompletedHandlerSpec.scala @@ -0,0 +1,282 @@ +/* + * 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 org.apache.pekko.actor.ActorSystem +import org.apache.pekko.testkit.{ImplicitSender, TestKit} +import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity} +import org.apache.texera.amber.core.workflow.{GlobalPortIdentity, PortIdentity, WorkflowContext} +import org.apache.texera.amber.engine.architecture.coordinator.{ + CoordinatorConfig, + CoordinatorProcessor +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + PortCompletedRequest, + ControlInvocation => ControlInvocationPayload +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + ControlError, + ControlReturn, + EmptyReturn, + ReturnInvocation +} +import org.apache.texera.amber.engine.architecture.rpc.coordinatorservice.CoordinatorServiceGrpc.{ + METHOD_COORDINATOR_INITIATE_ADVANCE_REGION_EXECUTIONS, + METHOD_PORT_COMPLETED +} +import org.apache.texera.amber.engine.architecture.scheduling.RegionExecutionManagerTestSupport.{ + createSingleWorkerRegion, + createSourceOp, + createWorkerId, + seedReusableWorkerExecution +} +import org.apache.texera.amber.engine.architecture.scheduling.{ + RegionExecutionManagerTestSupport, + RegionIdentity, + Schedule +} +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.ControlInvocation +import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpecLike + +import scala.collection.mutable + +/** + * `portCompleted` is the request whose handling completes a region: once the last port of a + * region is done, the region is terminated, which sends `EndWorker` to each of its workers. + * + * The reply to `portCompleted` and that `EndWorker` travel the same FIFO control channel to the + * sending worker, so the region advance must NOT happen inline in this handler — it would put + * `EndWorker` on the wire ahead of the reply, and the worker would then process `EndWorker` with + * the reply still queued behind it and reject the termination (#6891). The advance is instead + * requested as a coordinator-to-coordinator control message, which the coordinator can only + * process in a later round — by which time this round's reply has been sent. + */ +class PortCompletedHandlerSpec + extends TestKit(ActorSystem("PortCompletedHandlerSpec", AmberRuntime.pekkoConfig)) + with ImplicitSender + with AnyFlatSpecLike + with BeforeAndAfterAll + with RegionExecutionManagerTestSupport { + + private val physicalOp = createSourceOp("port-completed-op") + private val workerId = createWorkerId(physicalOp) + private val outputPortId = PortIdentity() + private val portCompletedCommandId = 7L + + override def afterAll(): Unit = { + TestKit.shutdownActorSystem(system) + } + + /** + * A real `CoordinatorProcessor` with a single-region schedule already in flight: the region + * owns the port the worker will report, and its worker execution is pre-seeded so starting the + * region does not create real workers. `advanceRegionExecutions` is called once to move the + * region into the executing set (its `startWorker` RPCs are captured and never answered, so + * the region stays executing); the messages it produced are then dropped so each test only + * observes what handling `portCompleted` emits. + */ + private def coordinatorWithRegionInFlight() + : (CoordinatorProcessor, mutable.ListBuffer[WorkflowFIFOMessage]) = { + val captured = mutable.ListBuffer[WorkflowFIFOMessage]() + val cp = new CoordinatorProcessor( + new WorkflowContext(), + CoordinatorConfig(None, None, None, None), + COORDINATOR, + { + case Right(msg) => captured += msg + case Left(_) => () + } + ) + val harness = createCoordinatorHarness() + registerLiveWorker(harness.actorRefService, workerId) + cp.setupActorService(harness.actorService) + cp.workflowExecutionManager.setupActorRefService(harness.actorRefService) + + val region = createSingleWorkerRegion(1, physicalOp, workerId).copy( + ports = Set(GlobalPortIdentity(physicalOp.id, outputPortId, input = false)) + ) + seedReusableWorkerExecution(cp.workflowExecution, seedRegionId = 101, physicalOp, workerId) + cp.workflowExecutionManager.schedule = Schedule(Map(0 -> Set(region))) + cp.workflowExecutionManager.advanceRegionExecutions(cp.actorService) + + captured.clear() + (cp, captured) + } + + /** Deliver a worker's `portCompleted` through the real RPC server, as production does. */ + private def receivePortCompleted(cp: CoordinatorProcessor, portId: PortIdentity): Unit = { + cp.processDCM( + ChannelIdentity(workerId, COORDINATOR, isControl = true), + ControlInvocation( + METHOD_PORT_COMPLETED, + PortCompletedRequest(portId, input = false), + AsyncRPCContext(workerId, COORDINATOR), + portCompletedCommandId + ) + ) + } + + private def messagesTo( + captured: mutable.ListBuffer[WorkflowFIFOMessage], + receiver: ActorVirtualIdentity + ): Seq[WorkflowFIFOMessage] = + captured.filter(_.channelId.toWorkerId == receiver).toSeq + + private def repliesTo( + captured: mutable.ListBuffer[WorkflowFIFOMessage], + receiver: ActorVirtualIdentity + ): Seq[ReturnInvocation] = + messagesTo(captured, receiver).collect { + case WorkflowFIFOMessage(_, _, ret: ReturnInvocation) => + ret + } + + private def selfInvocations( + captured: mutable.ListBuffer[WorkflowFIFOMessage] + ): Seq[ControlInvocationPayload] = + messagesTo(captured, COORDINATOR).collect { + case WorkflowFIFOMessage(_, _, inv: ControlInvocationPayload) => inv + } + + /** The coordinator-addressed messages asking for region executions to be advanced. */ + private def advanceRequestsIn( + captured: mutable.ListBuffer[WorkflowFIFOMessage] + ): Seq[ControlInvocationPayload] = + // The RPC proxy sends the reflected Java method name, which differs from the generated + // constant only in the leading case; the server matches case-insensitively too. + selfInvocations(captured).filter( + _.methodName.equalsIgnoreCase( + METHOD_COORDINATOR_INITIATE_ADVANCE_REGION_EXECUTIONS.getBareMethodName + ) + ) + + /** Resolve the statistics query the handler awaits, so its continuation runs. */ + private def resolveStatisticsQuery( + cp: CoordinatorProcessor, + captured: mutable.ListBuffer[WorkflowFIFOMessage], + returnValue: ControlReturn = EmptyReturn() + ): Unit = { + val statsCommandId = selfInvocations(captured).head.commandId + cp.asyncRPCClient.fulfillPromise(ReturnInvocation(statsCommandId, returnValue)) + } + + "PortCompletedHandler" should + "request the region advance as a separate control message instead of advancing inline" in { + val (cp, captured) = coordinatorWithRegionInFlight() + + receivePortCompleted(cp, outputPortId) + resolveStatisticsQuery(cp, captured) + + // The advance leaves as a coordinator-to-coordinator control message, which the coordinator + // can only process in a later round — after this round's reply below has been sent. + assert(advanceRequestsIn(captured).size == 1) + assert( + repliesTo(captured, workerId) == Seq( + ReturnInvocation(portCompletedCommandId, EmptyReturn()) + ) + ) + } + + it should "send nothing but the reply to the reporting worker while handling portCompleted" in { + val (cp, captured) = coordinatorWithRegionInFlight() + + receivePortCompleted(cp, outputPortId) + resolveStatisticsQuery(cp, captured) + + // Guards against a regression to inline advancing: an `EndWorker` (or any other control + // invocation) emitted to this worker in this round would be ordered ahead of the reply. + assert(messagesTo(captured, workerId).size == 1) + assert( + !captured.exists(msg => + msg.channelId.toWorkerId == workerId && msg.payload.isInstanceOf[ControlInvocationPayload] + ) + ) + } + + it should "mark the reported port completed before requesting the advance" in { + val (cp, captured) = coordinatorWithRegionInFlight() + + receivePortCompleted(cp, outputPortId) + resolveStatisticsQuery(cp, captured) + + // The advance runs in a later round, so the bookkeeping it depends on must already be + // recorded in the execution state by the time the request goes out. + val operatorExecution = cp.workflowExecution + .getRegionExecution(RegionIdentity(1)) + .getOperatorExecution(physicalOp.id) + assert(operatorExecution.isOutputPortCompleted(outputPortId)) + assert(advanceRequestsIn(captured).size == 1) + } + + it should "not request an advance for a port that belongs to no executing region" in { + val (cp, captured) = coordinatorWithRegionInFlight() + + // "start"/"end" ports are not part of any region, so no region resolves and nothing advances. + receivePortCompleted(cp, PortIdentity(9)) + resolveStatisticsQuery(cp, captured) + + assert(advanceRequestsIn(captured).isEmpty) + assert( + repliesTo(captured, workerId) == Seq( + ReturnInvocation(portCompletedCommandId, EmptyReturn()) + ) + ) + } + + it should "request the advance without waiting for the sender's workerExecutionCompleted" in { + val (cp, captured) = coordinatorWithRegionInFlight() + + // The worker emits `workerExecutionCompleted` right after its last `portCompleted` + // (`OutputManager.finalizeOutput` appends FinalizePort before FinalizeExecutor). Region + // completion keys only on port completion, so booking the last port requests the advance + // while that request may still be queued at the coordinator — and the coordinator selects + // input channels out of a HashMap, so it may run the advance first and reply afterwards. + receivePortCompleted(cp, outputPortId) + resolveStatisticsQuery(cp, captured) + + assert(advanceRequestsIn(captured).size == 1) + // Nothing here has replied to a `workerExecutionCompleted`, so the resulting `EndWorker` can + // legitimately reach the worker before that reply does. `EndHandler` closes this by not + // counting a queued reply as work — the coordinator cannot order it, so the worker tolerates + // it (see EndHandlerSpec, "reply successfully when only a coordinator reply is queued"). + assert(repliesTo(captured, workerId).map(_.commandId) == Seq(portCompletedCommandId)) + } + + it should "return a failed reply to the worker when its bookkeeping fails" in { + val (cp, captured) = coordinatorWithRegionInFlight() + receivePortCompleted(cp, outputPortId) + + // Failing the statistics query fails the handler's continuation. The reply is still chained + // on that continuation, so the error is reported back to the worker that sent + // `portCompleted` — the contract documented for this handler, which deferring the advance + // must not break. + resolveStatisticsQuery(cp, captured, ControlError.defaultInstance) + + val reply = repliesTo(captured, workerId).head + assert(reply.commandId == portCompletedCommandId) + assert(reply.returnValue.isInstanceOf[ControlError]) + assert(advanceRequestsIn(captured).isEmpty) + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala index 3a3c0ecb9e..89fed62709 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala @@ -25,7 +25,10 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ AsyncRPCContext, EmptyRequest } -import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_QUERY_STATISTICS import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{ ActorCommandElement, @@ -110,4 +113,68 @@ class EndHandlerSpec extends AnyFlatSpec { assertEndWorkerFails(handler) } + + /** + * A reply to one of this worker's own requests is not work, and the coordinator cannot order + * every reply it owes before `EndWorker`: it defers region advancement to a later control round + * so `EndWorker` follows the requests it has already handled, but a request still queued at the + * coordinator when that advance runs — `workerExecutionCompleted`, emitted right after the last + * `portCompleted` — is replied to afterwards, and the coordinator selects input channels out of + * a `HashMap`, so there is no cross-channel order. Such a reply must not block termination. + */ + private def coordinatorReply(seq: Long, commandId: Long): DPInputQueueElement = + FIFOMessageElement( + WorkflowFIFOMessage( + ChannelIdentity(COORDINATOR, workerId, isControl = true), + seq, + ReturnInvocation(commandId, EmptyReturn()) + ) + ) + + it should "reply successfully when only a coordinator reply is queued" in { + val queue = new LinkedBlockingQueue[DPInputQueueElement]() + queue.put(coordinatorReply(seq = 0, commandId = 3)) + val handler = createEndHandlerForQueue(queue) + + assert(await(handler.endWorker(EmptyRequest(), rpcContext)) == EmptyReturn()) + // The reply is only inspected, never consumed: the DP thread still fulfills its promise. + assert(queue.size() == 1) + } + + it should "reply successfully when several coordinator replies are queued" in { + val queue = new LinkedBlockingQueue[DPInputQueueElement]() + queue.put(coordinatorReply(seq = 0, commandId = 3)) + queue.put(coordinatorReply(seq = 1, commandId = 4)) + val handler = createEndHandlerForQueue(queue) + + assert(await(handler.endWorker(EmptyRequest(), rpcContext)) == EmptyReturn()) + assert(queue.size() == 2) + } + + it should "fail when work is queued behind a coordinator reply" in { + // The scan must not stop at the head: a reply followed by a real request is still work. + val queue = new LinkedBlockingQueue[DPInputQueueElement]() + queue.put(coordinatorReply(seq = 0, commandId = 3)) + queue.put( + FIFOMessageElement( + WorkflowFIFOMessage( + ChannelIdentity(COORDINATOR, workerId, isControl = true), + 1, + ControlInvocation(METHOD_QUERY_STATISTICS, EmptyRequest(), rpcContext, 5) + ) + ) + ) + val handler = createEndHandlerForQueue(queue) + + assertEndWorkerFails(handler) + assert(queue.size() == 2) + } + + it should "not consume queued messages when it fails" in { + val queue = queueWithFifoControlMessage() + val handler = createEndHandlerForQueue(queue) + + assertEndWorkerFails(handler) + assert(queue.size() == 1) + } }
