This is an automated email from the ASF dual-hosted git repository.
xuang7 pushed a commit to branch release/v1.2
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/release/v1.2 by this push:
new f0a18ffbed fix(amber, v1.2): advance region executions in a later
coordinator round, not inside portCompleted (#7096)
f0a18ffbed is described below
commit f0a18ffbed6484bf59cb5d8788c69a09a022e6a9
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Jul 29 23:28:21 2026 -0700
fix(amber, v1.2): advance region executions in a later coordinator round,
not inside portCompleted (#7096)
### What changes were proposed in this PR?
Automated backport of #6960 to `release/v1.2`.
Source: 226614b527bf51a010355778440e1972ed3130a7 · [automation
run](https://github.com/apache/texera/actions/runs/30506845994)
### Any related issues, documentation, discussions?
Backport of #6960.
### How was this PR tested?
Release-branch CI runs on this branch once the conflicts are resolved
and this PR is marked ready for review.
### Was this PR authored or co-authored using generative AI tooling?
No.
Co-authored-by: Xinyuan Lin <[email protected]>
---
.../architecture/rpc/controllerservice.proto | 1 +
.../ControllerAsyncRPCHandlerInitializer.scala | 1 +
.../AdvanceRegionExecutionsHandler.scala | 64 +++++
.../promisehandlers/PortCompletedHandler.scala | 28 +-
.../scheduling/RegionExecutionCoordinator.scala | 12 +-
.../scheduling/WorkflowExecutionCoordinator.scala | 4 +
.../worker/promisehandlers/EndHandler.scala | 65 ++++-
.../promisehandlers/PortCompletedHandlerSpec.scala | 282 +++++++++++++++++++++
.../worker/promisehandlers/EndHandlerSpec.scala | 69 ++++-
9 files changed, 502 insertions(+), 24 deletions(-)
diff --git
a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
index 0932a7b914..13c916cd8a 100644
---
a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
+++
b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto
@@ -45,6 +45,7 @@ service ControllerService {
rpc JumpToOperatorRegion(JumpToOperatorRegionRequest) returns (EmptyReturn);
rpc LinkWorkers(LinkWorkersRequest) returns (EmptyReturn);
rpc ControllerInitiateQueryStatistics(QueryStatisticsRequest) returns
(EmptyReturn);
+ rpc ControllerInitiateAdvanceRegionExecutions(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/controller/ControllerAsyncRPCHandlerInitializer.scala
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala
index 7e5a904716..19c960da30 100644
---
a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala
+++
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/ControllerAsyncRPCHandlerInitializer.scala
@@ -41,6 +41,7 @@ class ControllerAsyncRPCHandlerInitializer(
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/controller/promisehandlers/AdvanceRegionExecutionsHandler.scala
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/AdvanceRegionExecutionsHandler.scala
new file mode 100644
index 0000000000..305be181a1
--- /dev/null
+++
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/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.controller.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.engine.architecture.controller.{
+ ControllerAsyncRPCHandlerInitializer,
+ 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: controller (itself)
+ */
+trait AdvanceRegionExecutionsHandler {
+ this: ControllerAsyncRPCHandlerInitializer =>
+
+ override def controllerInitiateAdvanceRegionExecutions(
+ request: EmptyRequest,
+ ctx: AsyncRPCContext
+ ): Future[EmptyReturn] = {
+ cp.workflowExecutionCoordinator
+ .coordinateRegionExecutors(cp.actorService)
+ // The requester is the controller 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/controller/promisehandlers/PortCompletedHandler.scala
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala
index 810c098c41..1dfb4b15d7 100644
---
a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala
+++
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala
@@ -20,14 +20,11 @@
package org.apache.texera.amber.engine.architecture.controller.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.controller.{
- ControllerAsyncRPCHandlerInitializer,
- FatalError
-}
+import
org.apache.texera.amber.engine.architecture.controller.ControllerAsyncRPCHandlerInitializer
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.workflowExecutionCoordinator
- .coordinateRegionExecutors(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 controller
addresses to itself is
+ // transmitted and received before it is handled, so the advance
lands behind the
+ // reply below.
+ controllerInterface.controllerInitiateAdvanceRegionExecutions(
+ EmptyRequest(),
+ CONTROLLER
+ )
}
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/RegionExecutionCoordinator.scala
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala
index bc790f3ce5..35e58eed7a 100644
---
a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala
+++
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala
@@ -132,8 +132,16 @@ class RegionExecutionCoordinator(
*
* 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 `ControllerInitiateAdvanceRegionExecutions` instead of advancing
inline, which defers the advance that reaches
+ * here to a later controller round (see `PortCompletedHandler`). That
orders the replies already produced; for one
+ * the controller 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/WorkflowExecutionCoordinator.scala
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala
index deb753beb3..cd29f5732d 100644
---
a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala
+++
b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionCoordinator.scala
@@ -61,6 +61,10 @@ class WorkflowExecutionCoordinator(
* 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
+ * `ControllerInitiateAdvanceRegionExecutions` (see `PortCompletedHandler`)
so the resulting
+ * `EndWorker` cannot overtake the reply that request still owes.
*/
def coordinateRegionExecutors(actorService: PekkoActorService): Future[Unit]
= {
val unfinishedRegionCoordinators =
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 0504e66f52..1ae6efbeee 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
@@ -36,21 +43,67 @@ trait EndHandler {
/**
* The response of endWorker to the controller indicates that this worker
has finished not only
- * the data processing logic, but also , but also the processing of all the
control messages.
+ * 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 controller cannot order all of its
replies before `EndWorker`.
+ * The controller 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
+ * controller when the deferred advance runs — `workerExecutionCompleted`,
which the worker
+ * emits right after its last `portCompleted` — is replied to afterwards,
and the controller
+ * 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-controller 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 controller 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/controller/promisehandlers/PortCompletedHandlerSpec.scala
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/promisehandlers/PortCompletedHandlerSpec.scala
new file mode 100644
index 0000000000..12baaf43cd
--- /dev/null
+++
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/controller/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.controller.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.controller.{
+ ControllerConfig,
+ ControllerProcessor
+}
+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.controllerservice.ControllerServiceGrpc.{
+ METHOD_CONTROLLER_INITIATE_ADVANCE_REGION_EXECUTIONS,
+ METHOD_PORT_COMPLETED
+}
+import
org.apache.texera.amber.engine.architecture.scheduling.RegionCoordinatorTestSupport.{
+ createSingleWorkerRegion,
+ createSourceOp,
+ createWorkerId,
+ seedReusableWorkerExecution
+}
+import org.apache.texera.amber.engine.architecture.scheduling.{
+ RegionCoordinatorTestSupport,
+ 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.CONTROLLER
+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 controller-to-controller control message, which the
controller 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 RegionCoordinatorTestSupport {
+
+ 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 `ControllerProcessor` 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. `coordinateRegionExecutors` 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 controllerWithRegionInFlight()
+ : (ControllerProcessor, mutable.ListBuffer[WorkflowFIFOMessage]) = {
+ val captured = mutable.ListBuffer[WorkflowFIFOMessage]()
+ val cp = new ControllerProcessor(
+ new WorkflowContext(),
+ ControllerConfig(None, None, None, None),
+ CONTROLLER,
+ {
+ case Right(msg) => captured += msg
+ case Left(_) => ()
+ }
+ )
+ val harness = createControllerHarness()
+ registerLiveWorker(harness.actorRefService, workerId)
+ cp.setupActorService(harness.actorService)
+
cp.workflowExecutionCoordinator.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.workflowExecutionCoordinator.schedule = Schedule(Map(0 -> Set(region)))
+ cp.workflowExecutionCoordinator.coordinateRegionExecutors(cp.actorService)
+
+ captured.clear()
+ (cp, captured)
+ }
+
+ /** Deliver a worker's `portCompleted` through the real RPC server, as
production does. */
+ private def receivePortCompleted(cp: ControllerProcessor, portId:
PortIdentity): Unit = {
+ cp.processDCM(
+ ChannelIdentity(workerId, CONTROLLER, isControl = true),
+ ControlInvocation(
+ METHOD_PORT_COMPLETED,
+ PortCompletedRequest(portId, input = false),
+ AsyncRPCContext(workerId, CONTROLLER),
+ 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, CONTROLLER).collect {
+ case WorkflowFIFOMessage(_, _, inv: ControlInvocationPayload) => inv
+ }
+
+ /** The controller-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_CONTROLLER_INITIATE_ADVANCE_REGION_EXECUTIONS.getBareMethodName
+ )
+ )
+
+ /** Resolve the statistics query the handler awaits, so its continuation
runs. */
+ private def resolveStatisticsQuery(
+ cp: ControllerProcessor,
+ 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) = controllerWithRegionInFlight()
+
+ receivePortCompleted(cp, outputPortId)
+ resolveStatisticsQuery(cp, captured)
+
+ // The advance leaves as a controller-to-controller control message, which
the controller
+ // 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) = controllerWithRegionInFlight()
+
+ 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) = controllerWithRegionInFlight()
+
+ 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) = controllerWithRegionInFlight()
+
+ // "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) = controllerWithRegionInFlight()
+
+ // 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 controller — and the
controller 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 controller cannot order it, so
the worker tolerates
+ // it (see EndHandlerSpec, "reply successfully when only a controller
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) = controllerWithRegionInFlight()
+ 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 90e8b817be..9051c06751 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
controller 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
+ * controller when that advance runs — `workerExecutionCompleted`, emitted
right after the last
+ * `portCompleted` — is replied to afterwards, and the controller selects
input channels out of
+ * a `HashMap`, so there is no cross-channel order. Such a reply must not
block termination.
+ */
+ private def controllerReply(seq: Long, commandId: Long): DPInputQueueElement
=
+ FIFOMessageElement(
+ WorkflowFIFOMessage(
+ ChannelIdentity(CONTROLLER, workerId, isControl = true),
+ seq,
+ ReturnInvocation(commandId, EmptyReturn())
+ )
+ )
+
+ it should "reply successfully when only a controller reply is queued" in {
+ val queue = new LinkedBlockingQueue[DPInputQueueElement]()
+ queue.put(controllerReply(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 controller replies are queued" in
{
+ val queue = new LinkedBlockingQueue[DPInputQueueElement]()
+ queue.put(controllerReply(seq = 0, commandId = 3))
+ queue.put(controllerReply(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 controller 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(controllerReply(seq = 0, commandId = 3))
+ queue.put(
+ FIFOMessageElement(
+ WorkflowFIFOMessage(
+ ChannelIdentity(CONTROLLER, 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)
+ }
}
