diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferService.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferService.scala index cba9b0b2ee4..065e8777a64 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferService.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferService.scala @@ -172,7 +172,7 @@ class PekkoMessageTransferService( if (msgsNeedResend.nonEmpty) { logger.debug(s"output for $channel: ${cc.getStatusReport}") } - if (refService.hasActorRef(channel.fromWorkerId)) { + if (refService.hasActorRef(channel.toWorkerId)) { msgsNeedResend.foreach { msg => refService.forwardToActor(msg) } diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferServiceSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferServiceSpec.scala new file mode 100644 index 00000000000..0653cd63488 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoMessageTransferServiceSpec.scala @@ -0,0 +1,280 @@ +/* + * 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.common + +import org.apache.pekko.actor.{Actor, ActorContext, ActorSystem, Cancellable, Props} +import org.apache.pekko.testkit.{TestActorRef, TestKit} +import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity} +import org.apache.texera.amber.engine.architecture.common.WorkflowActor.NetworkMessage +import org.apache.texera.amber.engine.architecture.messaginglayer.{CongestionControl, FlowControl} +import org.apache.texera.amber.engine.common.ambermessage.{ + WorkflowFIFOMessage, + WorkflowFIFOMessagePayload, + WorkflowMessage +} +import org.apache.texera.common.config.ApplicationConfig +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpecLike + +import scala.collection.mutable +import scala.concurrent.duration.FiniteDuration + +/** + * Unit tests for [[PekkoMessageTransferService]]. + * + * There is no Mockito in the amber test tree, so instead of mocking `ActorContext` + * we obtain a real one from Pekko TestKit: a minimal [[CtxHolder]] actor is spawned + * via `TestActorRef`, and `underlyingActor.context` supplies the live context that + * `PekkoActorService` / `PekkoActorRefMappingService` eagerly dereference + * (self / dispatcher / parent). + * + * `initialize()` is deliberately NOT called in the getAllUnAckedMessages and + * backpressure recipes — it would schedule real repeating timers on the system + * scheduler. The scheduler-driven `checkResend` path is covered separately by + * overriding `scheduleWithFixedDelay` to capture (rather than schedule) the + * callback (see [[CapturingActorService]]). + */ +class PekkoMessageTransferServiceSpec + extends TestKit(ActorSystem("PekkoMessageTransferServiceSpec")) + with AnyFlatSpecLike + with BeforeAndAfterAll { + + override def afterAll(): Unit = { + TestKit.shutdownActorSystem(system) + } + + private val actorId: ActorVirtualIdentity = ActorVirtualIdentity("transfer-test-worker") + + // A non-DataFrame payload so `WorkflowMessage.getInMemSize` falls through to the + // 200L default branch. (DataFrame(Array.empty) is 0 bytes and would never be able + // to overload flow control regardless of the configured credit.) + private case class FixedSizePayload() extends WorkflowFIFOMessagePayload + + private def dataChannel(): ChannelIdentity = + ChannelIdentity( + fromWorkerId = ActorVirtualIdentity("from"), + toWorkerId = ActorVirtualIdentity("to"), + isControl = false + ) + + private def fifo(chan: ChannelIdentity, seq: Long): WorkflowFIFOMessage = + WorkflowFIFOMessage(chan, seq, FixedSizePayload()) + + private def net(id: Long, chan: ChannelIdentity, seq: Long): NetworkMessage = + NetworkMessage(id, fifo(chan, seq)) + + private def backdateSentTime(cc: CongestionControl, id: Long): Unit = { + val field = classOf[CongestionControl].getDeclaredField("sentTime") + field.setAccessible(true) + val sentTime = field.get(cc).asInstanceOf[mutable.LongMap[Long]] + sentTime(id) = System.currentTimeMillis() - cc.resendTimeLimit - 1 + } + + // Pin the assumed payload size so this test fails loudly if the size accounting + // changes in a way that would invalidate the credit math below. + assert(WorkflowMessage.getInMemSize(fifo(dataChannel(), 0L)) == 200L) + + private val maxBytes: Long = ApplicationConfig.maxCreditAllowedInBytesPerChannel + + // Each service instance needs its own live ActorContext; spawn a fresh holder + // actor (unique name) and hand back its context. + private val ctxCounter = new java.util.concurrent.atomic.AtomicInteger(0) + private def freshContext(): ActorContext = { + val holder = + TestActorRef[CtxHolder](Props(new CtxHolder), s"ctx-holder-${ctxCounter.incrementAndGet()}") + holder.underlyingActor.context + } + + private def newActorService(): PekkoActorService = + new PekkoActorService(actorId, freshContext()) + + // --------------------------------------------------------------------------- + // getAllUnAckedMessages (PekkoMessageTransferService lines 137-145) + // --------------------------------------------------------------------------- + + "getAllUnAckedMessages" should + "return both flow-control-stashed and congestion-control in-transit messages" in { + val actorService = newActorService() + val refService = new PekkoActorRefMappingService(actorService) + val service = new PekkoMessageTransferService(actorService, refService, _ => ()) + + val chan = dataChannel() + + // FlowControl carrying a stashed message that is drainable at read time: + // stash it while credit is 0, then restore credit so getAllUnAckedMessages' + // internal `fc.getMessagesToSend` actually returns it. + val fc = new FlowControl() + val fcNet = net(1L, chan, 1L) + fc.updateQueuedCredit(maxBytes) // getCredit -> 0 + fc.getMessagesToSend(fcNet) // stash fcNet (returns empty, becomes overloaded) + fc.updateQueuedCredit(0L) // restore credit; fcNet remains in the stash + service.channelToFC(chan) = fc + + // CongestionControl holding one in-transit message. + val cc = new CongestionControl() + val ccNet = net(2L, chan, 2L) + cc.markMessageInTransit(ccNet) + service.channelToCC(chan) = cc + + val result = service.getAllUnAckedMessages.toList + assert(result.contains(fcNet.internalMessage)) + assert(result.contains(ccNet.internalMessage)) + assert(result.size == 2) + } + + // --------------------------------------------------------------------------- + // send / updateChannelCreditFromReceiver -> checkForBackPressure + // (PekkoMessageTransferService lines 79-104, 147-165) + // --------------------------------------------------------------------------- + + "send / updateChannelCreditFromReceiver" should + "raise backpressure when a data channel overloads and lower it when credit returns" in { + var backpressured = false + val actorService = newActorService() + val refService = new PekkoActorRefMappingService(actorService) + val service = + new PekkoMessageTransferService(actorService, refService, b => backpressured = b) + + val dataChan = dataChannel() + + // Exhaust the receiver-side credit so getCredit drops to 0. The stash is still + // empty here, so no channel is overloaded yet and backpressure stays off. + service.updateChannelCreditFromReceiver(dataChan, maxBytes) + assert(!backpressured) + + // A 200-byte data message cannot fit into 0 credit, so FlowControl stashes it and + // the channel becomes overloaded -> checkForBackPressure flips false -> true. + service.send(fifo(dataChan, 0L)) + assert(backpressured) + + // Restore credit: the stash drains, the channel is no longer overloaded, and + // checkForBackPressure flips true -> false. + service.updateChannelCreditFromReceiver(dataChan, 0L) + assert(!backpressured) + } + + // --------------------------------------------------------------------------- + // initialize + checkResend (PekkoMessageTransferService lines 58-63, 167-181) + // --------------------------------------------------------------------------- + + "initialize" should + "schedule the resend/credit-polling callbacks and let checkResend run without error" in { + val actorService = new CapturingActorService(actorId, freshContext()) + val refService = new PekkoActorRefMappingService(actorService) + val service = new PekkoMessageTransferService(actorService, refService, _ => ()) + + // At least one CongestionControl entry so the checkResend foreach body executes. + val chan = dataChannel() + service.channelToCC(chan) = new CongestionControl() + + service.initialize() + // initialize() schedules exactly two fixed-delay callbacks (resend + credit poll). + assert(actorService.capturedCallables.size == 2) + + // Invoke the captured resend callback directly. With no timed-out in-transit + // messages, the body simply iterates the CC map and returns without throwing. + val resendCallback = actorService.capturedCallables.head + resendCallback() + + service.stop() // cancels the (already-cancelled) captured handles; must not throw + assert(service.channelToCC.contains(chan)) + } + + "checkResend" should "resend a timed-out message when its destination is registered" in { + val actorService = new CapturingActorService(actorId, freshContext()) + val sender = ActorVirtualIdentity("known-sender") + val destination = ActorVirtualIdentity("known-destination") + val refService = new RecordingRefService(actorService, Set(destination)) + val service = new PekkoMessageTransferService(actorService, refService, _ => ()) + val channel = ChannelIdentity(sender, destination, isControl = false) + val message = net(10L, channel, 0L) + val cc = new CongestionControl() + cc.markMessageInTransit(message) + backdateSentTime(cc, message.messageId) + service.channelToCC(channel) = cc + + service.initialize() + actorService.capturedCallables.head() + + assert(refService.queriedIds == Seq(destination)) + assert(refService.forwardedMessages == Seq(message)) + } + + it should "not resend a timed-out message when only its sender is registered" in { + val actorService = new CapturingActorService(actorId, freshContext()) + val sender = ActorVirtualIdentity("known-sender") + val destination = ActorVirtualIdentity("removed-destination") + val refService = new RecordingRefService(actorService, Set(sender)) + val service = new PekkoMessageTransferService(actorService, refService, _ => ()) + val channel = ChannelIdentity(sender, destination, isControl = false) + val message = net(11L, channel, 0L) + val cc = new CongestionControl() + cc.markMessageInTransit(message) + backdateSentTime(cc, message.messageId) + service.channelToCC(channel) = cc + + service.initialize() + actorService.capturedCallables.head() + + assert(refService.queriedIds == Seq(destination)) + assert(refService.forwardedMessages.isEmpty) + } +} + +/** Minimal actor used only to obtain a real `ActorContext` from Pekko TestKit. */ +class CtxHolder extends Actor { + override def receive: Receive = { case _ => () } +} + +/** + * A [[PekkoActorService]] that captures scheduled callbacks instead of registering + * real repeating timers, so `initialize()` can run and the resend callback can be + * invoked deterministically from the test thread. + */ +class CapturingActorService(vid: ActorVirtualIdentity, ac: ActorContext) + extends PekkoActorService(vid, ac) { + + val capturedCallables: mutable.ArrayBuffer[() => Unit] = mutable.ArrayBuffer() + + override def scheduleWithFixedDelay( + initialDelay: FiniteDuration, + delay: FiniteDuration, + callable: () => Unit + ): Cancellable = { + capturedCallables += callable + Cancellable.alreadyCancelled + } +} + +class RecordingRefService( + actorService: PekkoActorService, + knownIds: Set[ActorVirtualIdentity] +) extends PekkoActorRefMappingService(actorService) { + + val queriedIds: mutable.ArrayBuffer[ActorVirtualIdentity] = mutable.ArrayBuffer() + val forwardedMessages: mutable.ArrayBuffer[NetworkMessage] = mutable.ArrayBuffer() + + override def hasActorRef(id: ActorVirtualIdentity): Boolean = { + queriedIds += id + knownIds.contains(id) + } + + override def forwardToActor(msg: NetworkMessage): Unit = forwardedMessages += msg +}