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 c85de8e5b345c5c94e96f3b7001e4dad89d7618b Author: 虎鸣 <[email protected]> AuthorDate: Fri Jun 19 19:21:15 2026 +0800 refactor(stream): address code review findings for ActorRefSource optimization - Extract shared resolveSupervisorCell on GraphStageLogic to eliminate duplicated localCell logic between StageActor and ActorRefSource - Add warning log for PoisonPill/Kill (matching StageActor behavior), using systemLog to avoid interpreter initialization dependency - Wrap removeFunctionRef in postStop with NonFatal catch for robustness during ActorSystem shutdown - Add super.postStop() call for forward compatibility - Log unmatched messages at debug level instead of silent drop - Add comment explaining direct-push fast path completion safety - Reduce startup race test from 20 to 5 iterations, verify stream completion via Done future - Add post-completion dead letter test --- .../pekko/stream/scaladsl/ActorRefSourceSpec.scala | 20 ++++++-- .../apache/pekko/stream/impl/ActorRefSource.scala | 54 +++++++++------------- .../org/apache/pekko/stream/stage/GraphStage.scala | 11 +++++ 3 files changed, 49 insertions(+), 36 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 e7315eb402..8372c9f46f 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 @@ -233,16 +233,17 @@ class ActorRefSourceSpec extends StreamSpec { // Regression: the first Source.actorRef materialization on a fresh ActorSystem // can race the async supervisor startup, hitting an UnstartedCell in StageActor.localCell. // The fix force-starts the supervisor via Ask(Identify) before returning the cell. - (1 to 20).foreach { i => + (1 to 5).foreach { i => val sys = ActorSystem(s"ActorRefSource-startup-race-$i") try { val mat = Materializer(sys) - val ref = Source + val (ref, done) = Source .actorRef[Int]({ case "ok" => CompletionStrategy.draining }, PartialFunction.empty, 8, OverflowStrategy.fail) - .to(Sink.ignore) + .toMat(Sink.ignore)(Keep.both) .run()(mat) ref should not be null ref ! "ok" + Await.result(done, 5.seconds) should be(Done) } finally { Await.result(sys.terminate(), 10.seconds) } @@ -299,5 +300,18 @@ class ActorRefSourceSpec extends StreamSpec { ref ! "ok" s.expectComplete() } + + "not fail when messages are sent after stream completion" in { + val (ref, done) = Source + .actorRef[Int]({ case "ok" => CompletionStrategy.draining }, PartialFunction.empty, 10, OverflowStrategy.fail) + .toMat(Sink.ignore)(Keep.both) + .run() + ref ! "ok" + Await.result(done, 5.seconds) should be(Done) + // After completion, FunctionRef should be removed; messages go to dead letters + ref ! 42 + ref ! 43 + // No exception should be thrown + } } } 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 5cea69f3a8..0d73f71231 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,17 +13,16 @@ package org.apache.pekko.stream.impl -import scala.concurrent.Await +import scala.util.control.NonFatal import org.apache.pekko -import pekko.actor.{ ActorCell, ActorRef, FunctionRef, Kill, LocalActorRef, PoisonPill, RepointableActorRef, UnstartedCell } +import pekko.actor.{ ActorRef, FunctionRef, Kill, PoisonPill } 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, Timeout } +import pekko.util.OptionVal private object ActorRefSource { private sealed trait ActorRefStage { def ref: ActorRef } @@ -67,48 +66,34 @@ private object ActorRefSource { private val name = inheritedAttributes.nameOrDefault(getClass.toString) - 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 cell = resolveSupervisorCell(eagerMaterializer) private val callback = getAsyncCallback[Any](onMessage) private val functionRef: FunctionRef = { val actorName = inheritedAttributes.nameForActorRef(super.stageActorName) + val systemLog = eagerMaterializer.system.log cell.addFunctionRef( (_, msg) => msg match { - case PoisonPill | Kill => // ignored — stage actor lifecycle messages are not honored - case _ => callback.invoke(msg) + case PoisonPill | Kill => + systemLog.warning( + "{} message sent to ActorRefSource({}) will be ignored, since it is not a real Actor. " + + "Use a custom message type to communicate with it instead.", + msg, + functionRef.path) + case _ => callback.invoke(msg) }, actorName) } override val ref: ActorRef = functionRef - override def postStop(): Unit = cell.removeFunctionRef(functionRef) + override def postStop(): Unit = { + try cell.removeFunctionRef(functionRef) + catch { case NonFatal(_) => () } + super.postStop() + } private def onMessage(msg: Any): Unit = { if (failureMatcher.isDefinedAt(msg)) { @@ -133,6 +118,8 @@ private object ActorRefSource { buf.used, name) } else if (buf.isEmpty && isAvailable(out)) { + // Direct-push fast path: safe to skip tryPush() because isCompleting + // is checked above, and completion is driven by the completionMatcher branch. push(out, m) } else if (!buf.isFull) { buf.enqueue(m) @@ -189,7 +176,8 @@ private object ActorRefSource { name) } } - case _ => // unmatched message, ignore + case unexpected => + log.debug("Dropping unexpected non-data message [{}] in stream [{}]", unexpected, name) } } } diff --git a/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala b/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala index 8325403b25..7e49ed39f2 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala @@ -1566,6 +1566,17 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount: existing } + /** + * INTERNAL API + * + * Resolves the stream supervisor's [[ActorCell]] so a caller can host a FunctionRef on it. + * If the supervisor is a [[RepointableActorRef]] that has not yet been started (its cell is still + * an [[UnstartedCell]]), it is force-started via Ask(Identify) before the cell is returned. + */ + @InternalApi + private[pekko] def resolveSupervisorCell(materializer: Materializer): ActorCell = + StageActor.localCell(materializer.supervisor, "Stream supervisor") + /** * Override and return a name to be given to the StageActor of this operator. * --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
