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-7118-e6b000ec93214d0a8a9db644081e32bec6b91286 in repository https://gitbox.apache.org/repos/asf/texera.git
commit 648d54e4e5fae0d34218e8aeb3bdc3cf561955a1 Author: Xinyuan Lin <[email protected]> AuthorDate: Fri Jul 31 11:44:06 2026 -0700 test(amber): add specs for the checkpoint and ECM promise handlers (#7118) ### What changes were proposed in this PR? The checkpoint / ECM promise handlers were the engine's lowest-coverage cluster: | Handler | Before | New spec | |---|---|---| | `TakeGlobalCheckpointHandler` | 4.5% | `TakeGlobalCheckpointHandlerSpec` (6 tests) | | `FinalizeCheckpointHandler` | 4.5% | `FinalizeCheckpointHandlerSpec` (5) | | `PrepareCheckpointHandler` | 4.8% | `PrepareCheckpointHandlerSpec` (4) | | `EmbeddedControlMessageHandler` | 8.3% | `EmbeddedControlMessageHandlerSpec` (6) | | `EvaluatePythonExpressionHandler` | 8.3% | `EvaluatePythonExpressionHandlerSpec` (6) | All five follow the existing precedent in this layer (`RetrieveWorkflowStateHandlerSpec`): a real `CoordinatorProcessor` plus a second handler initializer, capturing what the output handler emits — so the tests drive the **real** handler methods, not a re-implementation. Storage uses `ram://` in-memory VFS per the `LoggingSpec` / `DPThreadSpec` precedent, with a distinct folder per case since `VFS.getManager` is a JVM-wide singleton. Covered: the single-physical-op guard; the already-taken short-circuit that downgrades a checkpoint to an estimate; the per-checkpoint subfolder that `WorkflowActor.setupReplay` later resolves; channel scoping and the "latest execution" worker selection; the estimation branches on both worker handlers; the `CheckpointSupport` hand-off and the observable consequence of its skip branch; and the fan-in that must not answer before every targeted worker has replied. **Assertion strength was measured, not asserted.** Four rounds of production mutation, 21 of 22 non-trivial cases shown to go red, every mutation then reverted (`git diff -- amber/src/main` is empty): | Mutation | Result | |---|---| | `physicalOps.size != 1` → `> 2` | 2 tests fail | | `msg.targetOps.flatMap` → `msg.scope.flatMap` | 1 | | `if (!msg.estimationOnly)` → inverted | 3 | | `findLast` → `find` (latest execution) | 1 | | channel-scope filter `&&` → `||` | 1 | | `markAsReplayDestination` hoisted before `collect` | 1 | | drop the already-taken short-circuit | 1 | | `checkpoints.contains` → `!contains` | 5 | | fan-in → fire-and-forget | 1 | | …and 11 more | | Also removes the two commented-out blocks in `CheckpointSpec` (−73 lines), verified dead rather than assumed: the first matches `OpExecInitInfoWithCode` / `OpExecInitInfoWithFunc`, names that a repo-wide grep finds **only inside those comment lines** (`OpExecInitInfo` is now a ScalaPB sealed oneof), and casts to `CheckpointSupport`, which no production operator implements — plus it has zero assertions and one branch is literally `???`. Its intent is already covered by `SerializationManagerSpec` and `CheckpointSubsystemSpec`. The second calls `AmberClient` with six constructor arguments where the signature takes five, and starts a real two-operator workflow with 30s awaits. **Two production bugs surfaced while writing these; neither is fixed here, and neither is pinned as current behaviour.** 1. **`TakeGlobalCheckpointHandler` under-reports `totalSize`.** It accumulates into a captured `var` via a per-future `.onSuccess { resp => totalSize += resp.size }` registered *before* `Future.collect` wraps those futures. A `com.twitter.util.Promise` runs continuations in **reverse** registration order, so for the last-satisfied future the `collect` continuation — and the `.map` that reads the counter — runs before that future's `onSuccess` has contributed. Observed directly: two workers reporting 4000 and 56 produced `TakeGlobalCheckpointResponse(4000)`; a single worker reporting 7 produced `0`. Fix direction is to sum from the collected `Seq` inside the `.map`. The spec case is named for the fan-in it does verify and carries a comment explaining why the total is deliberately not asserted. 2. **A `CheckpointState` can never be written to storage with the shipped config.** `SequentialRecordWriter.writeRecord` does `AmberRuntime.serde.serialize(obj).get`, but `CheckpointState` is a plain `class` that does not extend `java.io.Serializable`, while `cluster.conf` sets `allow-java-serialization = off` with `"java.io.Serializable" = kryo` as the only binding — so the write throws `NotSerializableException`. That kills both persist paths, which means global checkpointing cannot produce a checkpoint on disk and the checkpoint-restore branch of `WorkflowActor.setupReplay` has nothing to read. Existing specs missed it because they only exercise `chkpt.save` / `chkpt.load` (`SerializedState` *is* a case class). Consequence for this PR: the two write-branch cases wrap the call in `Try` and assert only the state the handler mutates beforehand — assertions that hold both before and after a fix, so they are not inverted change-detectors. The class scaladoc says so explicitly. A third, smaller hazard is recorded in the notes but not tested: `EvaluatePythonExpressionHandler` statically types worker results as `EvaluatedValue`, which is not a member of the `ControlReturn` sealed oneof. Erasure hides it on the JVM side and the Python side sets a stray attribute, so an empty `ControlReturn` comes back. That is why the spec asserts dispatch and fan-in shape but never fulfils a worker promise with a fake `EvaluatedValue` — doing so would cement a type confusion. ### Any related issues, documentation, discussions? Closes #7117 ### How was this PR tested? Five new specs, 27 tests. Run together with the three adjacent checkpoint specs they must not duplicate or break — 49 tests, Java 17: ``` sbt "WorkflowExecutionService/testOnly *EvaluatePythonExpressionHandlerSpec *EmbeddedControlMessageHandlerSpec *TakeGlobalCheckpointHandlerSpec *PrepareCheckpointHandlerSpec *FinalizeCheckpointHandlerSpec *CheckpointSpec *SerializationManagerSpec *CheckpointSubsystemSpec" ``` ``` [info] Suites: completed 8, aborted 0 [info] Tests: succeeded 49, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` Isolation was checked too: the new specs were re-run interleaved with `LoggingSpec`, `WorkflowWorkerSpec`, `DPThreadSpec` and the sibling promise-handler specs (16 suites, concurrent in one JVM), and then a second time in the same JVM to confirm the `ram://` folders leave no residue. No new spec mutates `AmberRuntime`'s private `_serde` / `_actorSystem`. `Test/scalafmtCheck` and `Test/scalafix --check` both `[success]`. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../EmbeddedControlMessageHandlerSpec.scala | 341 +++++++++++++++++++++ .../EvaluatePythonExpressionHandlerSpec.scala | 237 ++++++++++++++ .../TakeGlobalCheckpointHandlerSpec.scala | 338 ++++++++++++++++++++ .../FinalizeCheckpointHandlerSpec.scala | 276 +++++++++++++++++ .../PrepareCheckpointHandlerSpec.scala | 237 ++++++++++++++ .../engine/faulttolerance/CheckpointSpec.scala | 79 +---- 6 files changed, 1435 insertions(+), 73 deletions(-) diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EmbeddedControlMessageHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EmbeddedControlMessageHandlerSpec.scala new file mode 100644 index 0000000000..05e83fa1d2 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EmbeddedControlMessageHandlerSpec.scala @@ -0,0 +1,341 @@ +/* + * 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} +import org.apache.texera.amber.core.executor.OpExecInitInfo +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + ChannelIdentity, + EmbeddedControlMessageIdentity, + OperatorIdentity, + PhysicalOpIdentity +} +import org.apache.texera.amber.core.workflow.WorkflowContext.{ + DEFAULT_EXECUTION_ID, + DEFAULT_WORKFLOW_ID +} +import org.apache.texera.amber.core.workflow.{ + GlobalPortIdentity, + PhysicalLink, + PhysicalOp, + PortIdentity, + WorkflowContext +} +import org.apache.texera.amber.engine.architecture.coordinator.{ + CoordinatorAsyncRPCHandlerInitializer, + CoordinatorConfig, + CoordinatorProcessor +} +import org.apache.texera.amber.engine.architecture.logreplay.{ReplayLogManager, ReplayLogRecord} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.{ + NO_ALIGNMENT, + PORT_ALIGNMENT +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + EmbeddedControlMessage, + EmbeddedControlMessageType, + EmptyRequest, + PropagateEmbeddedControlMessageRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + PropagateEmbeddedControlMessageResponse, + ReturnInvocation, + StringResponse +} +import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_RETRIEVE_STATE +import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity} +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage +import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage +import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage.SequentialRecordWriter +import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR +import org.apache.texera.amber.util.VirtualIdentityUtils +import org.scalatest.flatspec.AnyFlatSpec + +import scala.collection.mutable.ArrayBuffer + +/** + * `propagateEmbeddedControlMessage` turns a "run this command on these operators" request into a + * single embedded control message (ECM) that rides the data channels: it packs one + * `ControlInvocation` per target worker into the ECM's command mapping, computes the set of + * channels the ECM must traverse, and injects it on the control channel of every source operator's + * workers. + * + * The breakages this spec catches: + * - collapsing `targetOps` and `scope` into one another. They are deliberately different sets: + * `targetOps` decides who runs the command, `scope` decides how far the ECM travels. + * - admitting a data channel whose endpoints are not both in `scope`, which would make the ECM + * wait for alignment on a channel that never receives it. + * - dropping either direction of the source workers' control channels from the scope. + * - answering (or marking the replay destination) before every target worker has replied — the + * mark is what a later replay uses as its stop point, so it must not be written while the + * propagation is still in flight. + * - dispatching to workers of a region execution that has already completed. + */ +class EmbeddedControlMessageHandlerSpec extends AnyFlatSpec { + + private val awaitTimeout = Duration.fromSeconds(1) + private val ecmId = EmbeddedControlMessageIdentity("ecm-under-test") + private val retrieveState = METHOD_RETRIEVE_STATE.getBareMethodName + + private val sourceOp = mkPhysicalOp("source", "main") + private val middleOp = mkPhysicalOp("middle", "main") + private val sinkOp = mkPhysicalOp("sink", "main") + private val sourceWorker = mkWorkerId(sourceOp, 0) + private val middleWorker = mkWorkerId(middleOp, 0) + private val sinkWorker = mkWorkerId(sinkOp, 0) + private val sourceToMiddle = + PhysicalLink(sourceOp.id, PortIdentity(), middleOp.id, PortIdentity()) + private val middleToSink = PhysicalLink(middleOp.id, PortIdentity(), sinkOp.id, PortIdentity()) + + private def mkPhysicalOp(logicalOpId: String, layerName: String): PhysicalOp = + PhysicalOp( + PhysicalOpIdentity(OperatorIdentity(logicalOpId), layerName), + DEFAULT_WORKFLOW_ID, + DEFAULT_EXECUTION_ID, + OpExecInitInfo.Empty + ) + + private def mkWorkerId(physicalOp: PhysicalOp, index: Int): ActorVirtualIdentity = + VirtualIdentityUtils.createWorkerIdentity(DEFAULT_WORKFLOW_ID, physicalOp.id, index) + + /** Records the ECM ids the handler marks, so completion ordering can be asserted. */ + private class RecordingLogManager extends ReplayLogManager { + val marked: ArrayBuffer[EmbeddedControlMessageIdentity] = ArrayBuffer() + + override def setupWriter(logWriter: SequentialRecordWriter[ReplayLogRecord]): Unit = () + + override def sendCommitted( + msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] + ): Unit = () + + override def terminate(): Unit = () + + override def markAsReplayDestination(id: EmbeddedControlMessageIdentity): Unit = marked += id + } + + private case class Fixture( + init: CoordinatorAsyncRPCHandlerInitializer, + sent: ArrayBuffer[WorkflowFIFOMessage], + logManager: RecordingLogManager + ) + + private def newFixture(): Fixture = { + val sent = ArrayBuffer[WorkflowFIFOMessage]() + val cp = new CoordinatorProcessor( + new WorkflowContext(), + CoordinatorConfig(None, None, None, None), + COORDINATOR, + { + case Right(m) => sent += m + case _ => () + } + ) + val logManager = new RecordingLogManager() + cp.setupLogManager(logManager) + Fixture(new CoordinatorAsyncRPCHandlerInitializer(cp), sent, logManager) + } + + /** + * Seed one running region execution spanning source -> middle -> sink, with one worker each and + * a channel execution per data link. The region's only port is the sink's input port and it is + * left incomplete, which is what keeps the region execution in the running set. + */ + private def seedRunningRegion(init: CoordinatorAsyncRPCHandlerInitializer): Unit = { + val region = Region( + RegionIdentity(1), + physicalOps = Set(sourceOp, middleOp, sinkOp), + physicalLinks = Set(sourceToMiddle, middleToSink), + ports = Set(GlobalPortIdentity(sinkOp.id, PortIdentity(), input = true)) + ) + val regionExecution = init.cp.workflowExecution.initRegionExecution(region) + regionExecution.initOperatorExecution(sourceOp.id).initWorkerExecution(sourceWorker) + regionExecution.initOperatorExecution(middleOp.id).initWorkerExecution(middleWorker) + regionExecution.initOperatorExecution(sinkOp.id).initWorkerExecution(sinkWorker) + regionExecution + .initLinkExecution(sourceToMiddle) + .initChannelExecution(ChannelIdentity(sourceWorker, middleWorker, isControl = false)) + regionExecution + .initLinkExecution(middleToSink) + .initChannelExecution(ChannelIdentity(middleWorker, sinkWorker, isControl = false)) + } + + /** + * The request used by most cases: `scope` deliberately stops before the sink operator and + * `targetOps` is narrower still, so a handler that confused the two sets would be caught. + */ + private def request( + ecmType: EmbeddedControlMessageType = PORT_ALIGNMENT + ): PropagateEmbeddedControlMessageRequest = + PropagateEmbeddedControlMessageRequest( + sourceOpToStartProp = Seq(sourceOp.id), + id = ecmId, + ecmType = ecmType, + scope = Seq(sourceOp.id, middleOp.id), + targetOps = Seq(middleOp.id), + command = EmptyRequest(), + methodName = retrieveState + ) + + private def embeddedMessages( + sent: ArrayBuffer[WorkflowFIFOMessage] + ): Seq[(ChannelIdentity, EmbeddedControlMessage)] = + sent.toSeq.collect { + case WorkflowFIFOMessage(channelId, _, ecm: EmbeddedControlMessage) => (channelId, ecm) + } + + private def onlyEmbeddedMessage( + sent: ArrayBuffer[WorkflowFIFOMessage] + ): (ChannelIdentity, EmbeddedControlMessage) = { + val messages = embeddedMessages(sent) + assert(messages.size == 1, s"expected exactly one ECM, got: $messages") + messages.head + } + + behavior of "EmbeddedControlMessageHandler" + + it should "inject the ECM on the control channel of every source operator's workers" in { + val fixture = newFixture() + seedRunningRegion(fixture.init) + + fixture.init.propagateEmbeddedControlMessage(request(), fixture.init.mkContext(COORDINATOR)) + + val (channelId, ecm) = onlyEmbeddedMessage(fixture.sent) + assert(channelId == ChannelIdentity(COORDINATOR, sourceWorker, isControl = true)) + assert(ecm.id == ecmId) + // A hard-coded alignment mode would show up here: PORT_ALIGNMENT is not the value any of the + // in-tree callers of this handler pass. + assert(ecm.ecmType == PORT_ALIGNMENT) + } + + it should "pack one invocation per target worker, not per in-scope worker" in { + val fixture = newFixture() + seedRunningRegion(fixture.init) + + fixture.init.propagateEmbeddedControlMessage(request(), fixture.init.mkContext(COORDINATOR)) + + val (_, ecm) = onlyEmbeddedMessage(fixture.sent) + // `targetOps` is only the middle operator while `scope` also covers the source, so a handler + // that keyed the mapping off `scope` would carry an extra entry here. + assert(ecm.commandMapping.keySet == Set(middleWorker.name)) + val invocation = ecm.commandMapping(middleWorker.name) + assert(invocation.methodName == retrieveState) + assert(invocation.command == EmptyRequest()) + assert(invocation.context.receiver == middleWorker) + } + + it should "scope the ECM to in-scope data channels plus both source control directions" in { + val fixture = newFixture() + seedRunningRegion(fixture.init) + + fixture.init.propagateEmbeddedControlMessage(request(), fixture.init.mkContext(COORDINATOR)) + + val (_, ecm) = onlyEmbeddedMessage(fixture.sent) + assert( + ecm.scope.toSet == Set( + ChannelIdentity(sourceWorker, middleWorker, isControl = false), + ChannelIdentity(COORDINATOR, sourceWorker, isControl = true), + ChannelIdentity(sourceWorker, COORDINATOR, isControl = true) + ) + ) + // middle -> sink crosses out of `scope`, so alignment must not wait on it. + assert(!ecm.scope.contains(ChannelIdentity(middleWorker, sinkWorker, isControl = false))) + } + + it should "answer only after every target worker replied, and report the replies by worker" in { + val fixture = newFixture() + seedRunningRegion(fixture.init) + + val response = fixture.init.propagateEmbeddedControlMessage( + request(), + fixture.init.mkContext(COORDINATOR) + ) + + assert(!response.isDefined) + // The replay destination is the stop point of a future replay; writing it while the + // propagation is still outstanding would truncate that replay early. + assert(fixture.logManager.marked.isEmpty) + + val (_, ecm) = onlyEmbeddedMessage(fixture.sent) + val commandId = ecm.commandMapping(middleWorker.name).commandId + fixture.init.cp.asyncRPCClient.fulfillPromise( + ReturnInvocation(commandId, StringResponse("middle-state")) + ) + + assert( + Await.result(response, awaitTimeout) == PropagateEmbeddedControlMessageResponse( + Map(middleWorker.name -> StringResponse("middle-state")) + ) + ) + assert(fixture.logManager.marked.toSeq == Seq(ecmId)) + } + + it should "ignore the workers of a region execution that already completed" in { + val fixture = newFixture() + seedRunningRegion(fixture.init) + // A second, finished execution of the middle operator with a different worker. Its output port + // is marked completed, which is what makes the region execution report COMPLETED. + val retiredWorker = mkWorkerId(middleOp, 7) + val completedRegion = Region( + RegionIdentity(2), + physicalOps = Set(middleOp), + physicalLinks = Set.empty, + ports = Set(GlobalPortIdentity(middleOp.id, PortIdentity(), input = false)) + ) + val completedExecution = fixture.init.cp.workflowExecution + .initRegionExecution(completedRegion) + .initOperatorExecution(middleOp.id) + completedExecution + .initWorkerExecution(retiredWorker) + .getOutputPortExecution(PortIdentity()) + .setCompleted() + + fixture.init.propagateEmbeddedControlMessage(request(), fixture.init.mkContext(COORDINATOR)) + + val (_, ecm) = onlyEmbeddedMessage(fixture.sent) + assert(ecm.commandMapping.keySet == Set(middleWorker.name)) + assert(!ecm.commandMapping.contains(retiredWorker.name)) + } + + it should "complete immediately and still mark the destination when there is nothing to target" in { + val fixture = newFixture() + + val response = fixture.init.propagateEmbeddedControlMessage( + PropagateEmbeddedControlMessageRequest( + sourceOpToStartProp = Seq.empty, + id = ecmId, + ecmType = NO_ALIGNMENT, + scope = Seq.empty, + targetOps = Seq.empty, + command = EmptyRequest(), + methodName = retrieveState + ), + fixture.init.mkContext(COORDINATOR) + ) + + // No source workers to inject into and no targets to wait for: the propagation is vacuously + // done, and the replay destination is still recorded so a replay can stop here. + assert(embeddedMessages(fixture.sent).isEmpty) + assert( + Await.result(response, awaitTimeout) == PropagateEmbeddedControlMessageResponse(Map.empty) + ) + assert(fixture.logManager.marked.toSeq == Seq(ecmId)) + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EvaluatePythonExpressionHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EvaluatePythonExpressionHandlerSpec.scala new file mode 100644 index 0000000000..6002b62a29 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/EvaluatePythonExpressionHandlerSpec.scala @@ -0,0 +1,237 @@ +/* + * 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} +import org.apache.texera.amber.core.executor.OpExecInitInfo +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + OperatorIdentity, + PhysicalOpIdentity +} +import org.apache.texera.amber.core.workflow.WorkflowContext.{ + DEFAULT_EXECUTION_ID, + DEFAULT_WORKFLOW_ID +} +import org.apache.texera.amber.core.workflow.{PhysicalOp, PhysicalPlan, WorkflowContext} +import org.apache.texera.amber.engine.architecture.coordinator.{ + CoordinatorAsyncRPCHandlerInitializer, + CoordinatorConfig, + CoordinatorProcessor +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + ControlInvocation, + EvaluatePythonExpressionRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EvaluatePythonExpressionResponse +import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity} +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.MainThreadDelegateMessage +import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage +import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR +import org.apache.texera.amber.util.VirtualIdentityUtils +import org.scalatest.flatspec.AnyFlatSpec + +import scala.collection.mutable.ArrayBuffer + +/** + * `evaluatePythonExpression` is the coordinator-side fan-out for the debugger's + * expression-evaluation panel: it resolves the logical operator the UI named to a single physical + * operator, then asks every worker of that operator's *latest* execution to evaluate the + * expression. + * + * Two things can silently break here. The resolution step is a precondition — a logical operator + * that expanded into several physical operators (or into none) has no unambiguous worker set, and + * the handler must refuse rather than pick one arbitrarily. The fan-out step must read the latest + * operator execution, because an operator re-run in a later region execution has a different set + * of live workers and the stale ones are gone. + * + * The handler dispatches through `workerInterface`, so these tests capture what the coordinator's + * output handler emits instead of running workers (the harness mirrors + * `RetrieveWorkflowStateHandlerSpec`). + */ +class EvaluatePythonExpressionHandlerSpec extends AnyFlatSpec { + + private val rpcContext = AsyncRPCContext(COORDINATOR, COORDINATOR) + private val awaitTimeout = Duration.fromSeconds(1) + + private def mkPhysicalOp(logicalOpId: String, layerName: String): PhysicalOp = + PhysicalOp( + PhysicalOpIdentity(OperatorIdentity(logicalOpId), layerName), + DEFAULT_WORKFLOW_ID, + DEFAULT_EXECUTION_ID, + OpExecInitInfo.Empty + ) + + private def mkWorkerId(physicalOp: PhysicalOp, index: Int): ActorVirtualIdentity = + VirtualIdentityUtils.createWorkerIdentity(DEFAULT_WORKFLOW_ID, physicalOp.id, index) + + /** + * Build a coordinator handler initializer over a physical plan holding `physicalOps`. Dispatched + * control messages are appended to the returned buffer. + */ + private def newFixture( + physicalOps: Set[PhysicalOp] + ): (CoordinatorAsyncRPCHandlerInitializer, ArrayBuffer[WorkflowFIFOMessage]) = { + val sent = ArrayBuffer[WorkflowFIFOMessage]() + val outputHandler: Either[MainThreadDelegateMessage, WorkflowFIFOMessage] => Unit = { + case Right(m) => sent += m + case _ => () + } + val cp = new CoordinatorProcessor( + new WorkflowContext(), + CoordinatorConfig(None, None, None, None), + COORDINATOR, + outputHandler + ) + cp.workflowScheduler.physicalPlan = PhysicalPlan(physicalOps, Set.empty) + (new CoordinatorAsyncRPCHandlerInitializer(cp), sent) + } + + /** + * Register a region execution that owns `physicalOp` with the given workers. Region executions + * are stored in creation order, and the handler resolves the operator through the *last* one + * that knows it, so the call order of this helper is what decides which workers are "latest". + */ + private def seedOperatorExecution( + init: CoordinatorAsyncRPCHandlerInitializer, + regionId: Long, + physicalOp: PhysicalOp, + workerIds: Seq[ActorVirtualIdentity] + ): Unit = { + val operatorExecution = init.cp.workflowExecution + .initRegionExecution(Region(RegionIdentity(regionId), Set(physicalOp), Set.empty)) + .initOperatorExecution(physicalOp.id) + workerIds.foreach(operatorExecution.initWorkerExecution) + } + + /** The per-worker expression-evaluation invocations the handler dispatched. */ + private def dispatchedEvaluations( + sent: ArrayBuffer[WorkflowFIFOMessage] + ): Seq[ControlInvocation] = + sent.toSeq.collect { + case WorkflowFIFOMessage(_, _, invocation: ControlInvocation) + if invocation.methodName == "evaluatePythonExpression" => + invocation + } + + behavior of "EvaluatePythonExpressionHandler" + + it should "refuse a logical operator that has no physical operator in the plan" in { + val (init, sent) = newFixture(Set(mkPhysicalOp("present", "main"))) + + val ex = intercept[RuntimeException] { + init.evaluatePythonExpression( + EvaluatePythonExpressionRequest("1 + 1", "absent"), + rpcContext + ) + } + + assert(ex.getMessage.contains("has 0 physical operators, expecting a single one")) + assert(ex.getMessage.contains("absent")) + assert(sent.isEmpty) + } + + it should "refuse a logical operator that expanded into several physical operators" in { + // A logical operator with more than one physical layer has no unambiguous worker set; the + // handler must not silently evaluate against `physicalOps.head`. + val (init, sent) = newFixture( + Set(mkPhysicalOp("expanded", "first"), mkPhysicalOp("expanded", "second")) + ) + + val ex = intercept[RuntimeException] { + init.evaluatePythonExpression( + EvaluatePythonExpressionRequest("1 + 1", "expanded"), + rpcContext + ) + } + + assert(ex.getMessage.contains("has 2 physical operators, expecting a single one")) + assert(sent.isEmpty) + } + + it should "forward the request unchanged to every worker of the operator" in { + val physicalOp = mkPhysicalOp("udf", "main") + val (init, sent) = newFixture(Set(physicalOp)) + val workerIds = Seq(mkWorkerId(physicalOp, 0), mkWorkerId(physicalOp, 1)) + seedOperatorExecution(init, regionId = 1, physicalOp, workerIds) + val request = EvaluatePythonExpressionRequest("tuple_['x']", "udf") + + init.evaluatePythonExpression(request, rpcContext) + + val invocations = dispatchedEvaluations(sent) + assert(invocations.map(_.context.receiver).toSet == workerIds.toSet) + assert(invocations.size == workerIds.size) + // The worker-side handler needs the expression verbatim, so nothing may be rewritten on the + // way out. + assert(invocations.forall(_.command == request)) + } + + it should "target the workers of the operator's latest execution, not an earlier one" in { + val physicalOp = mkPhysicalOp("udf", "main") + val (init, sent) = newFixture(Set(physicalOp)) + val staleWorkerId = mkWorkerId(physicalOp, 0) + val liveWorkerId = mkWorkerId(physicalOp, 1) + // The operator was re-run: region 2's execution superseded region 1's, and region 1's worker + // no longer exists. + seedOperatorExecution(init, regionId = 1, physicalOp, Seq(staleWorkerId)) + seedOperatorExecution(init, regionId = 2, physicalOp, Seq(liveWorkerId)) + + init.evaluatePythonExpression(EvaluatePythonExpressionRequest("1", "udf"), rpcContext) + + assert(dispatchedEvaluations(sent).map(_.context.receiver) == Seq(liveWorkerId)) + } + + it should "not answer before every targeted worker has replied" in { + val physicalOp = mkPhysicalOp("udf", "main") + val (init, sent) = newFixture(Set(physicalOp)) + seedOperatorExecution( + init, + regionId = 1, + physicalOp, + Seq(mkWorkerId(physicalOp, 0), mkWorkerId(physicalOp, 1)) + ) + + val response = init.evaluatePythonExpression( + EvaluatePythonExpressionRequest("1", "udf"), + rpcContext + ) + + // Both worker invocations are outstanding, so the aggregate response must still be pending — + // the UI panel shows one row per worker and a partial answer would be wrong. + assert(dispatchedEvaluations(sent).size == 2) + assert(!response.isDefined) + } + + it should "answer immediately with no values when the operator has no workers" in { + val physicalOp = mkPhysicalOp("udf", "main") + val (init, sent) = newFixture(Set(physicalOp)) + seedOperatorExecution(init, regionId = 1, physicalOp, Seq.empty) + + val response = init.evaluatePythonExpression( + EvaluatePythonExpressionRequest("1", "udf"), + rpcContext + ) + + // Nothing to ask, so the handler must not stall on an empty collect. + assert(dispatchedEvaluations(sent).isEmpty) + assert(Await.result(response, awaitTimeout) == EvaluatePythonExpressionResponse(Seq.empty)) + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/TakeGlobalCheckpointHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/TakeGlobalCheckpointHandlerSpec.scala new file mode 100644 index 0000000000..064b2b9960 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/coordinator/promisehandlers/TakeGlobalCheckpointHandlerSpec.scala @@ -0,0 +1,338 @@ +/* + * 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} +import org.apache.texera.amber.core.executor.OpExecInitInfo +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + EmbeddedControlMessageIdentity, + OperatorIdentity, + PhysicalOpIdentity +} +import org.apache.texera.amber.core.workflow.WorkflowContext.{ + DEFAULT_EXECUTION_ID, + DEFAULT_WORKFLOW_ID +} +import org.apache.texera.amber.core.workflow.{PhysicalOp, PhysicalPlan, WorkflowContext} +import org.apache.texera.amber.engine.architecture.coordinator.{ + CoordinatorAsyncRPCHandlerInitializer, + CoordinatorConfig, + CoordinatorProcessor +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessageType.NO_ALIGNMENT +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + ControlInvocation, + FinalizeCheckpointRequest, + PrepareCheckpointRequest, + PropagateEmbeddedControlMessageRequest, + TakeGlobalCheckpointRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + FinalizeCheckpointResponse, + PropagateEmbeddedControlMessageResponse, + ReturnInvocation, + TakeGlobalCheckpointResponse +} +import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_PREPARE_CHECKPOINT +import org.apache.texera.amber.engine.architecture.scheduling.{Region, RegionIdentity} +import org.apache.texera.amber.engine.common.CheckpointState +import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage +import org.apache.texera.amber.engine.common.storage.SequentialRecordStorage +import org.apache.texera.amber.engine.common.virtualidentity.util.{COORDINATOR, SELF} +import org.apache.texera.amber.util.VirtualIdentityUtils +import org.scalatest.flatspec.AnyFlatSpec + +import java.net.URI +import scala.collection.mutable.ArrayBuffer + +/** + * `takeGlobalCheckpoint` is the two-phase coordinator entry point for a global checkpoint. Phase + * one propagates a `prepareCheckpoint` embedded control message over the whole physical plan; + * phase two asks every worker that answered to finalize its checkpoint into a per-checkpoint + * subfolder of the destination, and sums the reported sizes into the response. + * + * The breakages this spec catches: + * - mixing up the three operator-id sets in the propagation request. `scope`/`targetOps` come + * from the physical plan, while `sourceOpToStartProp` comes from the operators that actually + * have an execution; the fixture keeps those deliberately different. + * - losing the "already taken" short-circuit, which downgrades a repeat request to an estimate + * instead of overwriting a checkpoint that is already on disk. + * - finalizing to the bare destination instead of the per-checkpoint subfolder, which would make + * concurrent checkpoints collide. + * - answering before every worker has finalized. + * - writing anything to storage on an estimate-only request. + * + * Storage is the in-memory `ram://` VFS also used by `LoggingSpec`, so the `SequentialRecordStorage` + * calls run for real without touching disk. Each case uses its own `ram://` folder, because the VFS + * manager is a JVM-wide singleton. + */ +class TakeGlobalCheckpointHandlerSpec extends AnyFlatSpec { + + private val rpcContext = AsyncRPCContext(COORDINATOR, COORDINATOR) + private val awaitTimeout = Duration.fromSeconds(1) + private val checkpointId = EmbeddedControlMessageIdentity("global-checkpoint-1") + + // Two operators in the plan; only `startedOp` has a region execution, so the propagation's + // "where to start" set is a strict subset of its "what to cover" set. + private val startedOp = mkPhysicalOp("started") + private val notStartedOp = mkPhysicalOp("not-started") + private val planOpIds = Set(startedOp.id, notStartedOp.id) + + private def mkPhysicalOp(logicalOpId: String): PhysicalOp = + PhysicalOp( + PhysicalOpIdentity(OperatorIdentity(logicalOpId), "main"), + DEFAULT_WORKFLOW_ID, + DEFAULT_EXECUTION_ID, + OpExecInitInfo.Empty + ) + + private def mkWorkerId(physicalOp: PhysicalOp, index: Int): ActorVirtualIdentity = + VirtualIdentityUtils.createWorkerIdentity(DEFAULT_WORKFLOW_ID, physicalOp.id, index) + + private def newFixture() + : (CoordinatorAsyncRPCHandlerInitializer, ArrayBuffer[WorkflowFIFOMessage]) = { + val sent = ArrayBuffer[WorkflowFIFOMessage]() + val cp = new CoordinatorProcessor( + new WorkflowContext(), + CoordinatorConfig(None, None, None, None), + COORDINATOR, + { + case Right(m) => sent += m + case _ => () + } + ) + cp.workflowScheduler.physicalPlan = PhysicalPlan(Set(startedOp, notStartedOp), Set.empty) + cp.workflowExecution + .initRegionExecution(Region(RegionIdentity(1), Set(startedOp), Set.empty)) + .initOperatorExecution(startedOp.id) + (new CoordinatorAsyncRPCHandlerInitializer(cp), sent) + } + + private def invocationsNamed( + sent: ArrayBuffer[WorkflowFIFOMessage], + methodName: String + ): Seq[ControlInvocation] = + sent.toSeq.collect { + case WorkflowFIFOMessage(_, _, invocation: ControlInvocation) + if invocation.methodName == methodName => + invocation + } + + private def onlyPropagation(sent: ArrayBuffer[WorkflowFIFOMessage]): ControlInvocation = { + val invocations = invocationsNamed(sent, "propagateEmbeddedControlMessage") + assert(invocations.size == 1, s"expected one propagation, got: $invocations") + invocations.head + } + + private def propagationRequest( + invocation: ControlInvocation + ): PropagateEmbeddedControlMessageRequest = + invocation.command match { + case req: PropagateEmbeddedControlMessageRequest => req + case other => fail(s"unexpected command: $other") + } + + private def prepareRequest(invocation: ControlInvocation): PrepareCheckpointRequest = + propagationRequest(invocation).command match { + case req: PrepareCheckpointRequest => req + case other => fail(s"unexpected carried command: $other") + } + + /** Answer the propagation as if `workerNames` had all run `prepareCheckpoint`. */ + private def completePropagation( + init: CoordinatorAsyncRPCHandlerInitializer, + sent: ArrayBuffer[WorkflowFIFOMessage], + workerNames: Seq[String] + ): Unit = + init.cp.asyncRPCClient.fulfillPromise( + ReturnInvocation( + onlyPropagation(sent).commandId, + PropagateEmbeddedControlMessageResponse(workerNames.map(_ -> EmptyReturn()).toMap) + ) + ) + + private def checkpointStorage(destination: String): SequentialRecordStorage[CheckpointState] = + SequentialRecordStorage.getStorage[CheckpointState](Some(new URI(destination))) + + behavior of "TakeGlobalCheckpointHandler" + + it should "propagate prepareCheckpoint over the whole plan, starting from the started operators" in { + val (init, sent) = newFixture() + + init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest( + estimationOnly = true, + checkpointId, + "ram:///take-global-checkpoint-propagate/" + ), + rpcContext + ) + + val invocation = onlyPropagation(sent) + // The coordinator asks itself to run the propagation. + assert(invocation.context.receiver == SELF) + val request = propagationRequest(invocation) + assert(request.id == checkpointId) + assert(request.ecmType == NO_ALIGNMENT) + assert(request.scope.toSet == planOpIds) + assert(request.targetOps.toSet == planOpIds) + // Only the operator with an execution can inject the message; the un-started one has no + // workers to start propagation from. Compared as a Set because production derives this from + // `RegionExecution.getAllOperatorExecutions`, which is backed by a mutable.HashMap — sequence + // order is not something to lean on once more than one operator is started. + assert(request.sourceOpToStartProp.toSet == Set(startedOp.id)) + assert(request.methodName == METHOD_PREPARE_CHECKPOINT.getBareMethodName) + assert( + prepareRequest(invocation) == PrepareCheckpointRequest(checkpointId, estimationOnly = true) + ) + } + + it should "downgrade to an estimate when the destination already holds this checkpoint" in { + val destination = "ram:///take-global-checkpoint-already-taken/" + // Creating a VFS-backed storage at the per-checkpoint folder creates the folder, i.e. this is + // exactly the on-disk state left behind by a previous successful checkpoint. + checkpointStorage(new URI(destination).resolve(checkpointId.toString).toString) + val (init, sent) = newFixture() + + init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest(estimationOnly = false, checkpointId, destination), + rpcContext + ) + + // The caller asked for a real checkpoint, but one is already there, so the workers must be + // told to estimate rather than overwrite it. + assert( + prepareRequest(onlyPropagation(sent)) == PrepareCheckpointRequest( + checkpointId, + estimationOnly = true + ) + ) + } + + it should "still take the checkpoint when a different checkpoint id exists at the destination" in { + val destination = "ram:///take-global-checkpoint-other-id/" + checkpointStorage( + new URI(destination) + .resolve(EmbeddedControlMessageIdentity("some-other-checkpoint").toString) + .toString + ) + val (init, sent) = newFixture() + + init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest(estimationOnly = false, checkpointId, destination), + rpcContext + ) + + // The short-circuit must key on this checkpoint's own folder, not on the destination being + // non-empty. + assert( + prepareRequest(onlyPropagation(sent)) == PrepareCheckpointRequest( + checkpointId, + estimationOnly = false + ) + ) + } + + it should "finalize on every worker that answered, into the per-checkpoint subfolder" in { + val destination = "ram:///take-global-checkpoint-finalize/" + val (init, sent) = newFixture() + val workerIds = Seq(mkWorkerId(startedOp, 0), mkWorkerId(startedOp, 1)) + + init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest(estimationOnly = true, checkpointId, destination), + rpcContext + ) + // Phase two must not start before the propagation answered. + assert(invocationsNamed(sent, "finalizeCheckpoint").isEmpty) + completePropagation(init, sent, workerIds.map(_.name)) + + val finalizations = invocationsNamed(sent, "finalizeCheckpoint") + assert(finalizations.map(_.context.receiver).toSet == workerIds.toSet) + assert(finalizations.map(_.command).forall { + case request: FinalizeCheckpointRequest => + // The write target is a per-checkpoint subfolder of the destination — the same layout + // `WorkflowActor.setupReplay` reads back via `readFrom.resolve(replayTo.toString)`. A + // handler that passed `msg.destination` through unresolved would leave the last path + // segment off. + val writeTo = new URI(request.writeTo) + request.checkpointId == checkpointId && + writeTo.getScheme == "ram" && + writeTo.getPath == s"/take-global-checkpoint-finalize/$checkpointId" + case other => fail(s"unexpected finalize command: $other") + }) + } + + it should "answer only after every worker has finalized" in { + val destination = "ram:///take-global-checkpoint-sizes/" + val (init, sent) = newFixture() + val workerIds = Seq(mkWorkerId(startedOp, 0), mkWorkerId(startedOp, 1)) + + val response = init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest(estimationOnly = true, checkpointId, destination), + rpcContext + ) + completePropagation(init, sent, workerIds.map(_.name)) + + val finalizations = invocationsNamed(sent, "finalizeCheckpoint") + assert(finalizations.size == 2) + assert(!response.isDefined) + init.cp.asyncRPCClient.fulfillPromise( + ReturnInvocation(finalizations.head.commandId, FinalizeCheckpointResponse(4000L)) + ) + // One worker still owes a reply, so the global checkpoint is not done. + assert(!response.isDefined) + init.cp.asyncRPCClient.fulfillPromise( + ReturnInvocation(finalizations(1).commandId, FinalizeCheckpointResponse(56L)) + ) + + // The reported `totalSize` is deliberately not asserted. The handler accumulates it from + // per-future `onSuccess` callbacks, and a Twitter `Promise` runs its continuations in reverse + // registration order, so `Future.collect`'s continuation — and with it the response — is built + // before the last-satisfied future's `onSuccess` has contributed. Pinning today's number would + // turn a fix of that into a failure. + assert(Await.result(response, awaitTimeout).isInstanceOf[TakeGlobalCheckpointResponse]) + } + + it should "write nothing to storage for an estimate-only checkpoint" in { + val destination = "ram:///take-global-checkpoint-no-write/" + val (init, sent) = newFixture() + + val response = init.takeGlobalCheckpoint( + TakeGlobalCheckpointRequest(estimationOnly = true, checkpointId, destination), + rpcContext + ) + completePropagation(init, sent, Seq(mkWorkerId(startedOp, 0).name)) + init.cp.asyncRPCClient.fulfillPromise( + ReturnInvocation( + invocationsNamed(sent, "finalizeCheckpoint").head.commandId, + FinalizeCheckpointResponse(7L) + ) + ) + // Await so the whole two-phase flow has run before storage is inspected. + Await.result(response, awaitTimeout) + + // The per-checkpoint folder is created by the writer, so its absence shows the coordinator + // skipped serializing its own state. + assert(!checkpointStorage(destination).containsFolder(checkpointId.toString)) + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandlerSpec.scala new file mode 100644 index 0000000000..0efa8816b3 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/FinalizeCheckpointHandlerSpec.scala @@ -0,0 +1,276 @@ +/* + * 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.worker.promisehandlers + +import com.twitter.util.{Await, Duration, Future} +import org.apache.pekko.actor.{ActorSystem, Props} +import org.apache.pekko.testkit.{TestActorRef, TestKit} +import org.apache.texera.amber.clustering.SingleNodeListener +import org.apache.texera.amber.core.executor.OperatorExecutor +import org.apache.texera.amber.core.tuple.{Tuple, TupleLike} +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + ChannelIdentity, + EmbeddedControlMessageIdentity +} +import org.apache.texera.amber.core.workflow.PortIdentity +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + EmptyRequest, + FinalizeCheckpointRequest +} +import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_QUERY_STATISTICS +import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{ + DPInputQueueElement, + WorkerReplayInitialization +} +import org.apache.texera.amber.engine.architecture.worker.{ + DataProcessor, + DataProcessorRPCHandlerInitializer, + WorkflowWorker +} +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.storage.SequentialRecordStorage +import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR +import org.apache.texera.amber.engine.common.{ + AmberRuntime, + CheckpointState, + CheckpointSupport, + SerializedState +} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpecLike + +import java.net.URI +import java.util.concurrent.LinkedBlockingQueue +import scala.collection.mutable +import scala.collection.mutable.ArrayBuffer +import scala.util.Try + +/** + * `finalizeCheckpoint` is the second half of a worker's checkpoint. It has two disjoint jobs, + * selected by whether `prepareCheckpoint` registered a checkpoint under this id: + * - a checkpoint exists: fold the input messages recorded since `prepareCheckpoint` into it, + * stop that recording, and write the checkpoint out. + * - no checkpoint exists (an estimate-only request): report the operator's estimated cost and + * touch nothing. + * + * The breakages this spec catches: + * - taking the estimate branch when a checkpoint *is* registered, which would report a made-up + * size and drop the recorded in-flight messages. + * - creating the destination folder on an estimate-only request. + * - reporting a non-zero size for an operator that cannot checkpoint, or ignoring the estimate of + * one that can. + * - saving another checkpoint's recorded messages, or leaving this checkpoint's recording in + * place so the worker keeps buffering input forever. + * + * The write itself is not asserted: `SequentialRecordWriter` serializes through + * `AmberRuntime.serde`, and `CheckpointState` is not `java.io.Serializable`, so the shipped Pekko + * config (kryo bound to `java.io.Serializable`, java serialization off) has no binding for it. The + * write-branch case below therefore asserts only the state the handler mutates before writing, and + * does not depend on whether the call as a whole succeeds. + * + * The write branch hands a closure to the worker's main thread and blocks until it runs, so that + * case uses a real `WorkflowWorker` behind a `TestActorRef` (synchronous dispatch, as in + * `WorkflowWorkerSpec`) and the worker's own `DataProcessor`. The estimate branch never reaches the + * main thread and uses a bare `DataProcessor`. + */ +class FinalizeCheckpointHandlerSpec + extends TestKit(ActorSystem("FinalizeCheckpointHandlerSpec", AmberRuntime.pekkoConfig)) + with AnyFlatSpecLike + with BeforeAndAfterAll { + + import FinalizeCheckpointHandlerSpec._ + + private val workerId = ActorVirtualIdentity("Worker:WF1-finalize-main-0") + private val rpcContext = AsyncRPCContext(COORDINATOR, workerId) + private val awaitTimeout = Duration.fromSeconds(10) + private val checkpointId = EmbeddedControlMessageIdentity("finalize-checkpoint-1") + private val unrelatedCheckpointId = EmbeddedControlMessageIdentity("some-other-checkpoint") + + override def beforeAll(): Unit = { + // WorkflowActor's actor service resolves node addresses through /user/cluster-info. + system.actorOf(Props[SingleNodeListener](), "cluster-info") + } + + override def afterAll(): Unit = { + TestKit.shutdownActorSystem(system) + } + + private def await[T](future: Future[T]): T = Await.result(future, awaitTimeout) + + /** A handler over a bare DataProcessor: enough for the estimate branch. */ + private def estimateHandler(executor: OperatorExecutor): DataProcessorRPCHandlerInitializer = { + val dp = new DataProcessor( + workerId, + _ => (), + new LinkedBlockingQueue[DPInputQueueElement]() + ) + dp.executor = executor + new DataProcessorRPCHandlerInitializer(dp) + } + + /** + * A live worker, so the main-thread hand-off in the write branch is delivered exactly as in + * production (`logManager.sendCommitted` -> `self !` -> `handleTriggerClosure`). + */ + private def liveWorker(): TestActorRef[WorkflowWorker] = { + val worker = TestActorRef( + new WorkflowWorker(WorkerConfig(workerId), WorkerReplayInitialization()) + ) + // The DP thread would otherwise race the assertions on the worker's recorded inputs. + worker.underlyingActor.dpThread.stop() + worker + } + + private def recordedMessage(seq: Long): WorkflowFIFOMessage = + WorkflowFIFOMessage( + ChannelIdentity(COORDINATOR, workerId, isControl = true), + seq, + ControlInvocation(METHOD_QUERY_STATISTICS, EmptyRequest(), rpcContext, seq) + ) + + private def storageAt(uri: String): SequentialRecordStorage[CheckpointState] = + SequentialRecordStorage.getStorage[CheckpointState](Some(new URI(uri))) + + behavior of "FinalizeCheckpointHandler" + + it should "report zero for an operator that cannot checkpoint" in { + val handler = estimateHandler(new PlainExecutor) + + val response = await( + handler.finalizeCheckpoint( + FinalizeCheckpointRequest(checkpointId, "ram:///finalize-plain-estimate/"), + rpcContext + ) + ) + + assert(response.size == 0L) + } + + it should "report the operator's own estimated cost when it can checkpoint" in { + val handler = estimateHandler(new EstimatingExecutor(estimatedCost = 4242L)) + + val response = await( + handler.finalizeCheckpoint( + FinalizeCheckpointRequest(checkpointId, "ram:///finalize-supported-estimate/"), + rpcContext + ) + ) + + assert(response.size == 4242L) + } + + it should "not create the destination folder for an estimate-only request" in { + val handler = estimateHandler(new EstimatingExecutor(estimatedCost = 1L)) + val parent = "ram:///finalize-estimate-no-write/" + + await( + handler.finalizeCheckpoint( + FinalizeCheckpointRequest(checkpointId, parent + "checkpoint-folder/"), + rpcContext + ) + ) + + // Opening a writer creates the folder, so its absence is what shows nothing was written. + assert(!storageAt(parent).containsFolder("checkpoint-folder")) + } + + it should "fold this checkpoint's recorded messages in and stop only its recording" in { + val worker = liveWorker() + val dp = worker.underlyingActor.dp + // A distinctive estimate: a handler that fell through to the estimate branch would neither + // touch the recorded inputs nor populate the checkpoint below. + dp.executor = new EstimatingExecutor(estimatedCost = 4242L) + val checkpoint = new CheckpointState() + dp.ecmManager.checkpoints(checkpointId) = checkpoint + val recorded = ArrayBuffer(recordedMessage(1), recordedMessage(2)) + worker.underlyingActor.recordedInputs(checkpointId) = recorded + worker.underlyingActor.recordedInputs(unrelatedCheckpointId) = ArrayBuffer(recordedMessage(99)) + val handler = new DataProcessorRPCHandlerInitializer(dp) + + // The outcome is intentionally ignored: everything asserted below happens on the worker's main + // thread, before the storage write this call ends with (see the note in the class comment). + Try( + handler.finalizeCheckpoint( + FinalizeCheckpointRequest(checkpointId, "ram:///finalize-fold/"), + rpcContext + ) + ) + + assert( + checkpoint.load[mutable.ArrayBuffer[WorkflowFIFOMessage]]( + SerializedState.IN_FLIGHT_MSG_KEY + ) == recorded + ) + // Recording is per-checkpoint: this one is done, the other must keep going. + assert(!worker.underlyingActor.recordedInputs.contains(checkpointId)) + assert(worker.underlyingActor.recordedInputs.contains(unrelatedCheckpointId)) + } + + it should "record an empty in-flight buffer when nothing was recorded for the checkpoint" in { + val worker = liveWorker() + val dp = worker.underlyingActor.dp + val checkpoint = new CheckpointState() + dp.ecmManager.checkpoints(checkpointId) = checkpoint + val handler = new DataProcessorRPCHandlerInitializer(dp) + + // Outcome ignored for the same reason as the case above. + Try( + handler.finalizeCheckpoint( + FinalizeCheckpointRequest(checkpointId, "ram:///finalize-fold-empty/"), + rpcContext + ) + ) + + // `WorkflowWorker.loadFromCheckpoint` reads this key unconditionally, so it has to be present + // even when no input arrived between prepare and finalize. + assert( + checkpoint + .load[mutable.ArrayBuffer[WorkflowFIFOMessage]](SerializedState.IN_FLIGHT_MSG_KEY) + .isEmpty + ) + } +} + +object FinalizeCheckpointHandlerSpec { + + /** Not `CheckpointSupport`: the estimate branch must fall through to 0. */ + class PlainExecutor extends OperatorExecutor { + override def processTuple(tuple: Tuple, port: Int): Iterator[TupleLike] = Iterator.empty + } + + class EstimatingExecutor(estimatedCost: Long) extends OperatorExecutor with CheckpointSupport { + override def processTuple(tuple: Tuple, port: Int): Iterator[TupleLike] = Iterator.empty + + override def serializeState( + currentIteratorState: Iterator[(TupleLike, Option[PortIdentity])], + checkpoint: CheckpointState + ): Iterator[(TupleLike, Option[PortIdentity])] = currentIteratorState + + override def deserializeState( + checkpoint: CheckpointState + ): Iterator[(TupleLike, Option[PortIdentity])] = Iterator.empty + + override def getEstimatedCheckpointCost: Long = estimatedCost + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandlerSpec.scala new file mode 100644 index 0000000000..60e779ffa5 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/PrepareCheckpointHandlerSpec.scala @@ -0,0 +1,237 @@ +/* + * 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.worker.promisehandlers + +import com.twitter.util.{Await, Duration, Future} +import org.apache.texera.amber.core.executor.OperatorExecutor +import org.apache.texera.amber.core.tuple.{Schema, Tuple, TupleLike} +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + EmbeddedControlMessageIdentity +} +import org.apache.texera.amber.core.workflow.PortIdentity +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + PrepareCheckpointRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{ + DPInputQueueElement, + MainThreadDelegateMessage +} +import org.apache.texera.amber.engine.architecture.worker.{ + DataProcessor, + DataProcessorRPCHandlerInitializer +} +import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage +import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR +import org.apache.texera.amber.engine.common.{CheckpointState, CheckpointSupport, SerializedState} +import org.scalatest.flatspec.AnyFlatSpec + +import java.util.concurrent.LinkedBlockingQueue +import scala.collection.mutable.ArrayBuffer + +/** + * `prepareCheckpoint` is the first half of a worker's checkpoint. It must not serialize anything + * itself: the DP thread may be in the middle of a tuple when the request arrives, so the handler + * only *registers* a serialization callback that the DP loop later fires at a safe point + * (`SerializationManager.applySerialization`). + * + * The breakages this spec catches: + * - inverting the estimate-only test, i.e. registering a serialization for a request that only + * asked for a size estimate (a real checkpoint would then be taken, and written, unasked), or + * skipping registration for a real checkpoint. + * - serializing eagerly inside the handler instead of at the DP's serialization point. + * - registering the checkpoint under the wrong id, which would make `finalizeCheckpoint` fall + * through to its estimate-only branch and silently write nothing. + * - not handing the operator the DP's pending output state, or dropping the replacement + * iterator the operator returns — the un-emitted tuples would be lost on restore. + * + * The registered callback's last step hands a closure to the worker's main thread and blocks on + * it, which needs a live worker actor. These tests therefore install an output handler that throws + * [[PrepareCheckpointHandlerSpec.MainThreadHandOffReached]] at that hand-off: everything the + * callback does before it (the part that runs on the DP thread) is observable, and the assertions + * below stop there. + * + * Not covered here, and not covered anywhere else either: this handler's own main-thread closure, + * which drains the input queue into `DP_QUEUED_MSG_KEY`, snapshots un-acked output into + * `OUTPUT_MSG_KEY`, and starts input recording by seeding `worker.recordedInputs(checkpointId)`. + * `FinalizeCheckpointHandlerSpec` exercises a *different* closure and hand-installs + * `recordedInputs` itself, so the prepare-to-finalize handshake is currently unverified. Closing + * that gap needs a live worker actor (the `TestActorRef` recipe in `FinalizeCheckpointHandlerSpec` + * would do it) and is left as follow-up. + */ +class PrepareCheckpointHandlerSpec extends AnyFlatSpec { + + import PrepareCheckpointHandlerSpec._ + + private val workerId = ActorVirtualIdentity("Worker:WF1-prepare-main-0") + private val rpcContext = AsyncRPCContext(COORDINATOR, workerId) + private val awaitTimeout = Duration.fromSeconds(5) + private val checkpointId = EmbeddedControlMessageIdentity("prepare-checkpoint-1") + + /** Records main-thread hand-offs, then aborts so the blocking wait is never reached. */ + private class OutputRecorder { + val delegates: ArrayBuffer[MainThreadDelegateMessage] = ArrayBuffer() + + def handle(msg: Either[MainThreadDelegateMessage, WorkflowFIFOMessage]): Unit = + msg match { + case Left(delegate) => + delegates += delegate + throw new MainThreadHandOffReached + case Right(_) => () + } + } + + private def newHandler( + executor: OperatorExecutor + ): (DataProcessorRPCHandlerInitializer, OutputRecorder) = { + val recorder = new OutputRecorder + val dp = new DataProcessor( + workerId, + recorder.handle, + new LinkedBlockingQueue[DPInputQueueElement]() + ) + dp.executor = executor + (new DataProcessorRPCHandlerInitializer(dp), recorder) + } + + private def await[T](future: Future[T]): T = Await.result(future, awaitTimeout) + + behavior of "PrepareCheckpointHandler" + + it should "register no serialization for an estimate-only request" in { + val (handler, recorder) = newHandler(new PlainExecutor) + + assert( + await( + handler.prepareCheckpoint(PrepareCheckpointRequest(checkpointId, true), rpcContext) + ) == EmptyReturn() + ) + + // Nothing was registered, so the DP's serialization point is a no-op and no checkpoint state + // exists for this id — the coordinator will get a size estimate from `finalizeCheckpoint` + // instead. + handler.dp.serializationManager.applySerialization() + assert(recorder.delegates.isEmpty) + assert(handler.dp.ecmManager.checkpoints.isEmpty) + } + + it should "reply before any state is serialized for a real checkpoint" in { + val (handler, recorder) = newHandler(new PlainExecutor) + + assert( + await( + handler.prepareCheckpoint(PrepareCheckpointRequest(checkpointId, false), rpcContext) + ) == EmptyReturn() + ) + + // The reply is what releases the embedded control message; the serialization itself must still + // be pending, waiting for the DP loop to reach a safe point. + assert(handler.dp.ecmManager.checkpoints.isEmpty) + assert(recorder.delegates.isEmpty) + } + + it should "serialize the DP state under the requested checkpoint id when the DP fires it" in { + val (handler, recorder) = newHandler(new PlainExecutor) + // Pre-install a distinctive pending output. Production reaches + // `outputIterator.setTupleOutput(...)` ONLY inside the `case support: CheckpointSupport` + // arm, so for a plain executor this iterator must come back untouched -- that is the + // observable consequence of taking the skip branch. (Asserting the *absence* of the + // fixture's operator-state key would prove nothing: only the fixture writes that key, and + // this executor has no serializeState at all, so it could never appear.) + val pendingOutput: (TupleLike, Option[PortIdentity]) = (emptyTuple, Some(PortIdentity(3))) + handler.dp.outputManager.outputIterator.setTupleOutput(Iterator(pendingOutput)) + await(handler.prepareCheckpoint(PrepareCheckpointRequest(checkpointId, false), rpcContext)) + + intercept[MainThreadHandOffReached] { + handler.dp.serializationManager.applySerialization() + } + + assert(handler.dp.ecmManager.checkpoints.keySet == Set(checkpointId)) + assert(handler.dp.ecmManager.checkpoints(checkpointId).has(SerializedState.DP_STATE_KEY)) + assert(handler.dp.outputManager.outputIterator.outputIter.toList == List(pendingOutput)) + assert(recorder.delegates.size == 1) + } + + it should "hand the operator the DP's pending output and install the iterator it returns" in { + val executor = new CheckpointAwareExecutor + val (handler, _) = newHandler(executor) + val pendingOutput: (TupleLike, Option[PortIdentity]) = (emptyTuple, Some(PortIdentity(3))) + handler.dp.outputManager.outputIterator.setTupleOutput(Iterator(pendingOutput)) + await(handler.prepareCheckpoint(PrepareCheckpointRequest(checkpointId, false), rpcContext)) + + intercept[MainThreadHandOffReached] { + handler.dp.serializationManager.applySerialization() + } + + // Output the executor produced but the DP has not emitted yet is part of the operator's state; + // it is handed over so the operator can persist it. + assert(executor.receivedOutputState == List(pendingOutput)) + // ...and whatever the operator returns becomes the DP's output from now on. The replacement + // uses a different port than `pendingOutput`, so a handler that reinstalled the original (or + // left the iterator untouched) is visible here. + assert( + handler.dp.outputManager.outputIterator.outputIter.toList == + List((emptyTuple, Some(RestoredPortId))) + ) + assert(handler.dp.ecmManager.checkpoints(checkpointId).has(OperatorStateKey)) + } +} + +object PrepareCheckpointHandlerSpec { + + /** Key the CheckpointSupport fixture writes, used to tell the two executor branches apart. */ + val OperatorStateKey = "spec-operator-state" + + /** Port carried by the iterator the fixture returns from `serializeState`. */ + val RestoredPortId: PortIdentity = PortIdentity(9) + + val emptyTuple: Tuple = Tuple.builder(new Schema()).build() + + class MainThreadHandOffReached + extends RuntimeException("output handler reached the main-thread hand-off") + + /** Not `CheckpointSupport`, so `serializeWorkerState` must take its skip branch. */ + class PlainExecutor extends OperatorExecutor { + override def processTuple(tuple: Tuple, port: Int): Iterator[TupleLike] = Iterator.empty + } + + class CheckpointAwareExecutor extends OperatorExecutor with CheckpointSupport { + var receivedOutputState: List[(TupleLike, Option[PortIdentity])] = Nil + + override def processTuple(tuple: Tuple, port: Int): Iterator[TupleLike] = Iterator.empty + + override def serializeState( + currentIteratorState: Iterator[(TupleLike, Option[PortIdentity])], + checkpoint: CheckpointState + ): Iterator[(TupleLike, Option[PortIdentity])] = { + receivedOutputState = currentIteratorState.toList + checkpoint.save(OperatorStateKey, "written-by-operator") + Iterator((emptyTuple, Some(RestoredPortId))) + } + + override def deserializeState( + checkpoint: CheckpointState + ): Iterator[(TupleLike, Option[PortIdentity])] = Iterator.empty + + override def getEstimatedCheckpointCost: Long = 0L + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala index 8eb58fc580..c15afffa06 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/faulttolerance/CheckpointSpec.scala @@ -102,78 +102,11 @@ class CheckpointSpec extends AnyFlatSpecLike with BeforeAndAfterAll { assert(ex.getMessage == "no state saved for key = unknown") } -// "CSVScanOperator" should "be serializable" in { -// val chkpt = new CheckpointState() -// val headerlessCsvOpDesc = TestOperators.headerlessSmallCsvScanOpDesc() -// val context = new WorkflowContext() -// headerlessCsvOpDesc.setContext(context) -// val phyOp = headerlessCsvOpDesc.getPhysicalOp(WorkflowIdentity(1), ExecutionIdentity(1)) -// phyOp.opExecInitInfo match { -// case OpExecInitInfoWithCode(codeGen) => ??? -// case OpExecInitInfoWithFunc(opGen) => -// val operator = opGen(1, 1) -// operator.open() -// val outputIter = -// operator.asInstanceOf[SourceOperatorExecutor].produceTuple().map(t => (t, None)) -// outputIter.next() -// outputIter.next() -// operator.asInstanceOf[CheckpointSupport].serializeState(outputIter, chkpt) -// chkpt.save("deserialization", opGen) -// val opGen2 = chkpt.load("deserialization").asInstanceOf[(Int, Int) => OperatorExecutor] -// val op = opGen2.apply(1, 1) -// op.asInstanceOf[CheckpointSupport].deserializeState(chkpt) -// } -// } -// -// "Workflow " should "take global checkpoint, reload and continue" in { -// val client1 = new AmberClient( -// system, -// workflow.context, -// workflow.physicalPlan, -// resultStorage, -// CoordinatorConfig.default, -// error => {} -// ) -// Await.result(client1.coordinatorInterface.startWorkflow(EmptyRequest(), ())) -// Thread.sleep(100) -// Await.result(client1.coordinatorInterface.pauseWorkflow(EmptyRequest(), ())) -// val checkpointId = EmbeddedControlMessageIdentity(s"Checkpoint_test_1") -// val uri = new URI("ram:///recovery-logs/tmp/") -// Await.result( -// client1.coordinatorInterface.takeGlobalCheckpoint( -// TakeGlobalCheckpointRequest(estimationOnly = false, checkpointId, uri.toString), -// () -// ), -// Duration.fromSeconds(30) -// ) -// client1.shutdown() -// Thread.sleep(100) -// var coordinatorConfig = CoordinatorConfig.default -// coordinatorConfig = -// coordinatorConfig.copy(stateRestoreConfOpt = Some(StateRestoreConfig(uri, checkpointId))) -// val completableFuture = new CompletableFuture[Unit]() -// val client2 = new AmberClient( -// system, -// workflow.context, -// workflow.physicalPlan, -// resultStorage, -// coordinatorConfig, -// error => {} -// ) -// client2.registerCallback[ExecutionStateUpdate] { evt => -// if (evt.state == COMPLETED) { -// completableFuture.complete(()) -// } -// } -// Thread.sleep(1000) -// assert( -// Await -// .result(client2.coordinatorInterface.startWorkflow(EmptyRequest(), ())) -// .workflowState == PAUSED -// ) -// Thread.sleep(5000) -// Await.result(client2.coordinatorInterface.resumeWorkflow(EmptyRequest(), ())) -// completableFuture.get(30000, TimeUnit.MILLISECONDS) -// } + // Checkpoint coverage beyond these round-trips lives elsewhere: SerializationManagerSpec and + // CheckpointSubsystemSpec cover operator and DP state going through a CheckpointState, and + // PrepareCheckpointHandlerSpec / FinalizeCheckpointHandlerSpec / + // TakeGlobalCheckpointHandlerSpec cover the promise handlers that drive a checkpoint. A full + // "checkpoint, reload, continue" run needs a live multi-operator workflow and belongs in an + // integration spec, not here. }
