You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pekko.apache.org by he...@apache.org on 2026/06/19 11:21:48 UTC
(pekko) 06/06: refactor(stream): address code review findings for ActorRefSource optimization
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: 虎鸣 <he...@alibaba-inc.com>
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: commits-unsubscribe@pekko.apache.org
For additional commands, e-mail: commits-help@pekko.apache.org