This is an automated email from the ASF dual-hosted git repository. He-Pin pushed a commit to branch perf/actorref-source-optimization in repository https://gitbox.apache.org/repos/asf/pekko.git
commit 29b17cdb6bac7f188937efa388487773781dacdb Author: 虎鸣 <[email protected]> AuthorDate: Fri Jun 19 18:07:39 2026 +0800 perf(stream): bypass StageActor in ActorRefSource for message-only async callback Motivation: ActorRefSource used getEagerStageActor which creates a (ActorRef, Any) Tuple2 (~24 bytes) per message via the StageActor FunctionRef lambda. The sender is never used by ActorRefSource — it is always Actor.noSender for external tells — so this allocation is pure waste on the hot path. Modification: Replace getEagerStageActor with a direct FunctionRef created on the stream supervisor's ActorCell. The FunctionRef lambda intercepts PoisonPill/Kill (ignoring them) and delegates to a message-only AsyncCallback[Any] — no tuple boxing. The supervisor cell is resolved inline with UnstartedCell race handling. postStop removes the FunctionRef to prevent leaks. Result: Each element avoids ~24 bytes of Tuple2 allocation. At 1M elem/s this saves ~24 MB/s of GC pressure. PoisonPill/Kill are explicitly ignored at the FunctionRef level, preventing them from leaking as data when T = Any. Tests: - sbt "stream-tests / Test / testOnly ...ActorRefSourceSpec" — 17/17 pass - New regression tests: PoisonPill and Kill ignored for T = Any References: Refs akkadotnet/akka.net#8263 --- .../pekko/stream/scaladsl/ActorRefSourceSpec.scala | 36 +++- .../apache/pekko/stream/impl/ActorRefSource.scala | 188 +++++++++++++-------- 2 files changed, 154 insertions(+), 70 deletions(-) diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/ActorRefSourceSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/ActorRefSourceSpec.scala index ac350d9c56..e7315eb402 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/ActorRefSourceSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/ActorRefSourceSpec.scala @@ -19,7 +19,7 @@ import scala.concurrent.duration._ import org.apache.pekko import pekko.Done -import pekko.actor.{ ActorRef, ActorSystem, Status } +import pekko.actor.{ ActorRef, ActorSystem, Kill, PoisonPill, Status } import pekko.stream.{ OverflowStrategy, _ } import pekko.stream.testkit._ import pekko.stream.testkit.Utils._ @@ -265,5 +265,39 @@ class ActorRefSourceSpec extends StreamSpec { ref ! "ok" s.expectComplete() } + + "ignore PoisonPill and not emit it as a data element" in { + val s = TestSubscriber.manualProbe[Any]() + val ref = Source + .actorRef[Any]({ case "ok" => CompletionStrategy.draining }, PartialFunction.empty, 10, OverflowStrategy.fail) + .to(Sink.fromSubscriber(s)) + .run() + val sub = s.expectSubscription() + sub.request(10) + ref ! PoisonPill + // PoisonPill must not be emitted and must not complete the stream + s.expectNoMessage(300.millis) + ref ! "real-element" + s.expectNext("real-element") + ref ! "ok" + s.expectComplete() + } + + "ignore Kill and not emit it as a data element" in { + val s = TestSubscriber.manualProbe[Any]() + val ref = Source + .actorRef[Any]({ case "ok" => CompletionStrategy.draining }, PartialFunction.empty, 10, OverflowStrategy.fail) + .to(Sink.fromSubscriber(s)) + .run() + val sub = s.expectSubscription() + sub.request(10) + ref ! Kill + // Kill must not be emitted and must not complete the stream + s.expectNoMessage(300.millis) + ref ! "real-element" + s.expectNext("real-element") + ref ! "ok" + s.expectComplete() + } } } diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/ActorRefSource.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/ActorRefSource.scala index 1b36f4fd20..5cea69f3a8 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/impl/ActorRefSource.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/impl/ActorRefSource.scala @@ -13,14 +13,17 @@ package org.apache.pekko.stream.impl +import scala.concurrent.Await + import org.apache.pekko -import pekko.actor.ActorRef +import pekko.actor.{ ActorCell, ActorRef, FunctionRef, Kill, LocalActorRef, PoisonPill, RepointableActorRef, UnstartedCell } import pekko.annotation.InternalApi +import pekko.pattern.ask import pekko.stream._ import pekko.stream.OverflowStrategies._ import pekko.stream.impl.Stages.DefaultAttributes import pekko.stream.stage._ -import pekko.util.OptionVal +import pekko.util.{ OptionVal, Timeout } private object ActorRefSource { private sealed trait ActorRefStage { def ref: ActorRef } @@ -62,87 +65,134 @@ private object ActorRefSource { } private var isCompleting: Boolean = false - override protected def stageActorName: String = inheritedAttributes.nameForActorRef(super.stageActorName) - private val name = inheritedAttributes.nameOrDefault(getClass.toString) - override val ref: ActorRef = getEagerStageActor(eagerMaterializer) { - case (_, m) if failureMatcher.isDefinedAt(m) => - failStage(failureMatcher(m)) - case (_, m) if completionMatcher.isDefinedAt(m) => - completionMatcher(m) match { + + private val cell: ActorCell = { + val supervisor = eagerMaterializer.supervisor + supervisor match { + case ref: LocalActorRef => ref.underlying + case r: RepointableActorRef => + r.underlying match { + case c: ActorCell => c + case _: UnstartedCell => + implicit val timeout: Timeout = r.system.settings.CreationTimeout + Await.result(r ? pekko.actor.Identify(null), timeout.duration) + r.underlying match { + case c: ActorCell => c + case unknown => + throw new IllegalStateException( + s"Stream supervisor must be a local actor, was [${unknown.getClass.getName}]") + } + case unknown => + throw new IllegalStateException( + s"Stream supervisor must be a local actor, was [${unknown.getClass.getName}]") + } + case unknown => + throw new IllegalStateException( + s"Stream supervisor must be a local actor, was [${unknown.getClass.getName}]") + } + } + + private val callback = getAsyncCallback[Any](onMessage) + + private val functionRef: FunctionRef = { + val actorName = inheritedAttributes.nameForActorRef(super.stageActorName) + cell.addFunctionRef( + (_, msg) => + msg match { + case PoisonPill | Kill => // ignored — stage actor lifecycle messages are not honored + case _ => callback.invoke(msg) + }, + actorName) + } + + override val ref: ActorRef = functionRef + + override def postStop(): Unit = cell.removeFunctionRef(functionRef) + + private def onMessage(msg: Any): Unit = { + if (failureMatcher.isDefinedAt(msg)) { + failStage(failureMatcher(msg)) + } else if (completionMatcher.isDefinedAt(msg)) { + completionMatcher(msg) match { case CompletionStrategy.Draining => isCompleting = true tryPush() case CompletionStrategy.Immediately => completeStage() } - case (_, m: T @unchecked) => - buffer match { - case OptionVal.Some(buf) => - if (isCompleting) { - log.warning( - "Dropping element because Status.Success received already, only draining already buffered elements: [{}] (pending: [{}] in stream [{}])", - m, - buf.used, - name) - } else if (buf.isEmpty && isAvailable(out)) { - push(out, m) - } else if (!buf.isFull) { - buf.enqueue(m) - tryPush() - } else - overflowStrategy match { - case s: DropHead => - log.log( - s.logLevel, - "Dropping the head element because buffer is full and overflowStrategy is: [DropHead] in stream [{}]", + } else { + msg match { + case m: T @unchecked => + buffer match { + case OptionVal.Some(buf) => + if (isCompleting) { + log.warning( + "Dropping element because Status.Success received already, only draining already buffered elements: [{}] (pending: [{}] in stream [{}])", + m, + buf.used, name) - buf.dropHead() + } else if (buf.isEmpty && isAvailable(out)) { + push(out, m) + } else if (!buf.isFull) { buf.enqueue(m) tryPush() - case s: DropTail => - log.log( - s.logLevel, - "Dropping the tail element because buffer is full and overflowStrategy is: [DropTail] in stream [{}]", + } else + overflowStrategy match { + case s: DropHead => + log.log( + s.logLevel, + "Dropping the head element because buffer is full and overflowStrategy is: [DropHead] in stream [{}]", + name) + buf.dropHead() + buf.enqueue(m) + tryPush() + case s: DropTail => + log.log( + s.logLevel, + "Dropping the tail element because buffer is full and overflowStrategy is: [DropTail] in stream [{}]", + name) + buf.dropTail() + buf.enqueue(m) + tryPush() + case s: DropBuffer => + log.log( + s.logLevel, + "Dropping all the buffered elements because buffer is full and overflowStrategy is: [DropBuffer] in stream [{}]", + name) + buf.clear() + buf.enqueue(m) + tryPush() + case s: Fail => + log.log( + s.logLevel, + "Failing because buffer is full and overflowStrategy is: [Fail] in stream [{}]", + name) + val bufferOverflowException = + BufferOverflowException(s"Buffer overflow (max capacity was: $maxBuffer)!") + failStage(bufferOverflowException) + case _: Backpressure => + failStage(new IllegalStateException("Backpressure is not supported")) + } + case _ => + if (isCompleting) { + log.warning( + "Dropping element because Status.Success received already: [{}] in stream [{}]", + m, name) - buf.dropTail() - buf.enqueue(m) - tryPush() - case s: DropBuffer => - log.log( - s.logLevel, - "Dropping all the buffered elements because buffer is full and overflowStrategy is: [DropBuffer] in stream [{}]", - name) - buf.clear() - buf.enqueue(m) - tryPush() - case s: Fail => - log.log( - s.logLevel, - "Failing because buffer is full and overflowStrategy is: [Fail] in stream [{}]", + } else if (isAvailable(out)) { + push(out, m) + } else { + log.debug( + "Dropping element because there is no downstream demand and no buffer: [{}] in stream [{}]", + m, name) - val bufferOverflowException = - BufferOverflowException(s"Buffer overflow (max capacity was: $maxBuffer)!") - failStage(bufferOverflowException) - case _: Backpressure => - // there is a precondition check in Source.actorRefSource factory method to not allow backpressure as strategy - failStage(new IllegalStateException("Backpressure is not supported")) - } - case _ => - if (isCompleting) { - log.warning("Dropping element because Status.Success received already: [{}] in stream [{}]", m, name) - } else if (isAvailable(out)) { - push(out, m) - } else { - log.debug( - "Dropping element because there is no downstream demand and no buffer: [{}] in stream [{}]", - m, - name) + } } + case _ => // unmatched message, ignore } - - case _ => throw new IllegalArgumentException() // won't happen, compiler exhaustiveness check pleaser - }.ref + } + } private def tryPush(): Unit = { if (isAvailable(out) && buffer.isDefined && buffer.get.nonEmpty) { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
