You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@pekko.apache.org by "He-Pin (via GitHub)" <gi...@apache.org> on 2026/06/01 05:12:27 UTC
[PR] Optimize lazy stage actor dispatch [pekko]
He-Pin opened a new pull request, #3035:
URL: https://github.com/apache/pekko/pull/3035
### Motivation
The public issue https://github.com/akka/akka-core/issues/26857 describes an opportunity to avoid an extra actor hop for non-eager stage actor refs, which is now Apache licensed.
### Modification
Add separate eager and lazy `StageActorRef` handlers:
- non-eager `getStageActor` refs are attached to the graph interpreter actor and still dispatch through `GraphInterpreter.onAsyncInput`
- eager stage actor refs keep the existing materializer supervisor + async callback path so they remain usable before stream demand
- directional tests cover lazy parent placement, lazy message delivery, eager pre-demand usability, and eager supervisor placement
- add a JMH benchmark for lazy `StageActorRef.tell` throughput
### Result
Non-eager stage actor messages target the graph interpreter actor while preserving async-input semantics and eager compatibility.
### Benchmark
JMH smoke result:
```text
StageActorRefBenchmark.lazy_stage_actor_ref_tell_10k thrpt 5611073.982 ops/s
```
Command:
```shell
sbt "bench-jmh / clean" "bench-jmh/Jmh/run -i 1 -wi 0 -f 1 -r 1s -w 1s .*StageActorRefBenchmark.*"
```
### Tests
- `scalafmt stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/StageActorRefSpec.scala bench-jmh/src/main/scala/org/apache/pekko/stream/StageActorRefBenchmark.scala`
- `sbt "stream-tests / Test / testOnly org.apache.pekko.stream.scaladsl.StageActorRefSpec"`
- `sbt "bench-jmh / Compile / compile" "stream / mimaReportBinaryIssues" "bench-jmh / headerCheck"`
- `sbt "bench-jmh / clean" "bench-jmh/Jmh/run -i 1 -wi 0 -f 1 -r 1s -w 1s .*StageActorRefBenchmark.*"`
- `git diff --check`
### Clean-room note
This implementation was developed against the Pekko codebase and the public issue statement only. No upstream code, tests, comments, or diff content were copied.
### References
Refs https://github.com/akka/akka-core/issues/26857
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3333598145
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -1382,7 +1529,7 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount:
* Override and return a name to be given to the StageActor of this operator.
*
* This method will be only invoked and used once, during the first [[getStageActor]]
- * invocation whichc reates the actor, since subsequent `getStageActors` calls function
+ * invocation which creates the actor, since subsequent `getStageActors` calls function
* like `become`, rather than creating new actors.
Review Comment:
Scaladoc refers to `getStageActors` (plural), but the public API method is `getStageActor`. This makes the documentation misleading and the inline code reference unsearchable.
##########
stream/src/main/java/org/apache/pekko/stream/stage/AbstractMpscDispatch.java:
##########
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * license agreements; and to You under the Apache License, version 2.0:
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * This file is part of the Apache Pekko project, which was derived from Akka.
+ */
+
+package org.apache.pekko.stream.stage;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+import org.apache.pekko.dispatch.AbstractNodeQueue;
+
+/**
+ * INTERNAL API.
+ *
+ * <p>Vyukov MPSC linked queue augmented with a volatile {@code int} election slot used by
+ * single-consumer drain coalescing: producers race to elect a drainer via IDLE -> SCHEDULED CAS
+ * on the state field. The state is driven through a {@link VarHandle} declared here so Scala
+ * callers can use the typed {@link #stateAcquire()} / {@link #stateSetRelease(int)} / {@link
+ * #stateCompareAndSet(int, int)} accessors without hitting Scala's polymorphic-signature inference
+ * quirks for {@code VarHandle.getAcquire} (which differ between Scala 2 and Scala 3). Same layering
+ * as {@code AbstractMailbox} for {@code Mailbox}.
+ */
+public abstract class AbstractMpscDispatch<T> extends AbstractNodeQueue<T> {
+
+ @SuppressWarnings("unused")
+ private volatile int _stateDoNotCallMeDirectly = 0;
+
+ private static final VarHandle stateHandle;
+
+ static {
+ try {
+ stateHandle =
+ MethodHandles.privateLookupIn(AbstractMpscDispatch.class, MethodHandles.lookup())
+ .findVarHandle(AbstractMpscDispatch.class, "_stateDoNotCallMeDirectly", int.class);
+ } catch (Throwable t) {
+ throw new ExceptionInInitializerError(t);
+ }
+ }
+
+ /** Acquire-load of the election state. Use only from the consumer (drain) side. */
+ public final int stateAcquire() {
+ return (int) stateHandle.getAcquire(this);
+ }
Review Comment:
`stateAcquire()` is documented as consumer-only, but it is used from producer threads in `StageActor.LazyDispatch.apply` to implement the double-checked election. The Javadoc should reflect that it is safe/expected to call from any thread (or provide a separate consumer-only accessor).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3345632116
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,123 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
Review Comment:
The PR description claims the lazy dispatch election state is implemented as a `@volatile var Int` driven by a `VarHandle` (to avoid per-instance wrapper allocations), but the implementation here uses `new AtomicInteger(...)`. Please either update the PR description to match the code, or switch the implementation to the described `VarHandle` pattern if that is still the intended design goal.
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,123 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
+
+ // Reused across all drain batches; allocated once at construction.
+ private[this] val drainCallback: Any => Unit = (_: Any) => drain()
+
+ override def apply(pair: (ActorRef, Any)): Unit = {
+ // Producer-side completion guard. `runAsyncInput` short-circuits the drain handler when the stage
+ // is completed, so without this check a CAS-winning producer would leave state=SCHEDULED forever
+ // and subsequent producers would skip the mailbox push, growing the queue unbounded. Dropping
+ // post-completion sends here matches the original per-tell behaviour where `runAsyncInput`
+ // silently ignored them. The read is racy against the interpreter thread that flips
+ // `shutdownCounter`, but eventually visible — enough to bound the leak to messages enqueued
+ // before completion becomes visible to this thread.
+ if (interpreter.isStageCompleted(logic)) return
+ add(pair) // Vyukov producer path: getAndSet + release-store, no CAS spin
+ // Double-checked CAS: uncontended fast path is one volatile read; only the IDLE->SCHEDULED winner
+ // pays a CAS + mailbox push.
+ if (state.get() == SchedStateIdle && state.compareAndSet(SchedStateIdle, SchedStateScheduled)) {
+ // Re-check after winning election: if completion landed between the initial guard and the CAS,
+ // `scheduleDrain` would post an envelope that `runAsyncInput` skips, leaving state=SCHEDULED.
+ // Reset to IDLE and skip the schedule.
+ if (interpreter.isStageCompleted(logic)) state.set(SchedStateIdle)
+ else scheduleDrain()
+ }
+ }
+
+ private def scheduleDrain(): Unit =
+ // 1 AsyncInput + 1 Envelope per drain batch (amortized across up to drainBatchSize tells).
+ interpreter.onAsyncInput(logic, null, NoPromise, drainCallback)
+
+ private def drain(): Unit = {
+ val limit = drainBatchSize // hoisted to a local so JIT treats it as a loop-invariant constant
+ var processed = 0
+ while (processed < limit) {
+ if (interpreter.isStageCompleted(logic)) {
+ // Stage completed mid-drain; drop the remainder (matches the original per-tell behaviour where
+ // runAsyncInput silently skipped completed stages). Don't reschedule — no future drain will run.
+ while (poll() ne null) ()
+ state.set(SchedStateIdle)
+ return
+ }
+ val item = poll()
+ if (item eq null) {
+ state.set(SchedStateIdle)
+ // Recheck race: a producer may have added between `poll == null` and the IDLE publish above.
+ // That producer saw state=SCHEDULED and skipped the mailbox send, so we must re-elect.
+ if (!isEmpty && state.compareAndSet(SchedStateIdle, SchedStateScheduled))
+ scheduleDrain()
+ return
+ }
+ handler(item)
+ processed += 1
+ }
+ // Hit batch cap with items potentially still queued. Re-schedule another envelope so other
+ // BoundaryEvents (pull/push/complete) can interleave via the actor mailbox. `scheduledState` stays
+ // SCHEDULED: concurrent producers correctly observe SCHEDULED and skip; the new envelope will drain.
+ scheduleDrain()
Review Comment:
`LazyDispatch.drain()` always schedules another drain envelope after hitting the batch limit, even if the last processed item completed the stage. In that case `GraphInterpreter.runAsyncInput` will skip the scheduled drain (because the stage is already completed), leaving `state` stuck at SCHEDULED and retaining any already-enqueued items until the FunctionRef is stopped. Add a completion check before re-scheduling and drop any remaining queued items when completed (matching the existing mid-drain completion branch).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447496003
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
I will check this soon.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3343517196
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -206,29 +207,61 @@ object GraphStageLogic {
*
* Not for user instantiation, use [[GraphStageLogic.getStageActor]].
*/
- final class StageActor @InternalApi() private[pekko] (
+ final class StageActor @InternalApi() private (
materializer: Materializer,
- getAsyncCallback: StageActorRef.Receive => AsyncCallback[(ActorRef, Any)],
initialReceive: StageActorRef.Receive,
- name: String) {
+ name: String,
+ cell: ActorCell,
+ buildDispatch: StageActorRef.Receive => ((ActorRef, Any)) => Unit) {
+
+ @InternalApi private[pekko] def this(
+ materializer: Materializer,
+ getAsyncCallback: StageActorRef.Receive => AsyncCallback[(ActorRef, Any)],
+ initialReceive: StageActorRef.Receive,
+ name: String) =
+ this(
+ materializer,
+ initialReceive,
+ name,
+ StageActor.localCell(materializer.supervisor, "Stream supervisor"),
+ receive => getAsyncCallback(receive).invoke)
+
+ @InternalApi private[pekko] def this(
+ materializer: Materializer,
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ initialReceive: StageActorRef.Receive,
+ name: String) =
+ this(
+ materializer,
+ initialReceive,
+ name,
+ StageActor.localCell(materializer.supervisor, "Stream supervisor"),
+ // Coalesce per-tell mailbox traffic: N tells produce 1 AsyncInput envelope (amortized).
+ receive =>
+ new StageActor.LazyDispatch(
+ interpreter,
+ logic,
+ receive.asInstanceOf[Any => Unit],
+ StageActor.drainBatchSize(materializer)))
+
+ // Monomorphic Function1 captured once; JIT can inline the apply at the FunctionRef call site.
+ private val dispatch: ((ActorRef, Any)) => Unit = buildDispatch(internalReceive)
- private val callback = getAsyncCallback(internalReceive)
-
- private def cell = materializer.supervisor match {
- case ref: LocalActorRef => ref.underlying
- case unknown =>
- throw new IllegalStateException(s"Stream supervisor must be a local actor, was [${unknown.getClass.getName}]")
- }
private val functionRef: FunctionRef = {
- val f: (ActorRef, Any) => Unit = {
- case (_, m @ (PoisonPill | Kill)) =>
- materializer.logger.warning(
- "{} message sent to StageActor({}) will be ignored, since it is not a real Actor." +
- "Use a custom message type to communicate with it instead.",
- m,
- functionRef.path)
- case pair => callback.invoke(pair)
- }
+ // Explicit (sender, msg) lambda (not a pattern-match Function2 literal) so the PoisonPill / Kill
+ // branch matches on `msg` directly and does not allocate a Tuple2. The regular branch still
+ // constructs one tuple per tell, as required by the `((ActorRef, Any)) => Unit` public Receive type.
+ val f: (ActorRef, Any) => Unit = (sender, msg) =>
+ msg match {
+ case PoisonPill | Kill =>
+ materializer.logger.warning(
+ "{} message sent to StageActor({}) will be ignored, since it is not a real Actor." +
+ "Use a custom message type to communicate with it instead.",
Review Comment:
Fixed in da487e3c74: added trailing space so the concatenation reads `...not a real Actor. Use a custom message type...`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3340382157
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,110 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
+
+ // Reused across all drain batches; allocated once at construction.
+ private[this] val drainCallback: Any => Unit = (_: Any) => drain()
+
+ override def apply(pair: (ActorRef, Any)): Unit = {
+ add(pair) // Vyukov producer path: getAndSet + release-store, no CAS spin
+ // Double-checked CAS: uncontended fast path is one volatile read; only the IDLE->SCHEDULED winner
+ // pays a CAS + mailbox push.
+ if (state.get() == SchedStateIdle && state.compareAndSet(SchedStateIdle, SchedStateScheduled))
+ scheduleDrain()
+ }
Review Comment:
`LazyDispatch` enqueues and schedules drains even after the stage has completed. However `GraphInterpreter.runAsyncInput` explicitly skips invoking the handler when `isStageCompleted(logic)` is true, so a drain scheduled after completion will never run and queued items will never be polled (risking unbounded queue growth until the stage actor is stopped/GC’d). Guard in the producer path and drop messages once the stage is completed (matching the old semantics where completed stages silently ignored async inputs).
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -206,29 +207,61 @@ object GraphStageLogic {
*
* Not for user instantiation, use [[GraphStageLogic.getStageActor]].
*/
- final class StageActor @InternalApi() private[pekko] (
+ final class StageActor @InternalApi() private (
materializer: Materializer,
- getAsyncCallback: StageActorRef.Receive => AsyncCallback[(ActorRef, Any)],
initialReceive: StageActorRef.Receive,
- name: String) {
+ name: String,
+ cell: ActorCell,
+ buildDispatch: StageActorRef.Receive => ((ActorRef, Any)) => Unit) {
+
+ @InternalApi private[pekko] def this(
+ materializer: Materializer,
+ getAsyncCallback: StageActorRef.Receive => AsyncCallback[(ActorRef, Any)],
+ initialReceive: StageActorRef.Receive,
+ name: String) =
+ this(
+ materializer,
+ initialReceive,
+ name,
+ StageActor.localCell(materializer.supervisor, "Stream supervisor"),
+ receive => getAsyncCallback(receive).invoke)
+
+ @InternalApi private[pekko] def this(
+ materializer: Materializer,
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ initialReceive: StageActorRef.Receive,
+ name: String) =
+ this(
+ materializer,
+ initialReceive,
+ name,
+ StageActor.localCell(materializer.supervisor, "Stream supervisor"),
+ // Coalesce per-tell mailbox traffic: N tells produce 1 AsyncInput envelope (amortized).
+ receive =>
+ new StageActor.LazyDispatch(
+ interpreter,
+ logic,
+ receive.asInstanceOf[Any => Unit],
+ StageActor.drainBatchSize(materializer)))
+
+ // Monomorphic Function1 captured once; JIT can inline the apply at the FunctionRef call site.
+ private val dispatch: ((ActorRef, Any)) => Unit = buildDispatch(internalReceive)
- private val callback = getAsyncCallback(internalReceive)
-
- private def cell = materializer.supervisor match {
- case ref: LocalActorRef => ref.underlying
- case unknown =>
- throw new IllegalStateException(s"Stream supervisor must be a local actor, was [${unknown.getClass.getName}]")
- }
private val functionRef: FunctionRef = {
- val f: (ActorRef, Any) => Unit = {
- case (_, m @ (PoisonPill | Kill)) =>
- materializer.logger.warning(
- "{} message sent to StageActor({}) will be ignored, since it is not a real Actor." +
- "Use a custom message type to communicate with it instead.",
- m,
- functionRef.path)
- case pair => callback.invoke(pair)
- }
+ // Explicit (sender, msg) lambda (not a pattern-match Function2 literal) so the PoisonPill / Kill
+ // branch matches on `msg` directly and does not allocate a Tuple2. The regular branch still
+ // constructs one tuple per tell, as required by the `((ActorRef, Any)) => Unit` public Receive type.
+ val f: (ActorRef, Any) => Unit = (sender, msg) =>
+ msg match {
+ case PoisonPill | Kill =>
+ materializer.logger.warning(
+ "{} message sent to StageActor({}) will be ignored, since it is not a real Actor." +
+ "Use a custom message type to communicate with it instead.",
Review Comment:
The warning message concatenation is missing a space between sentences, producing "Actor.Use ..." in logs. Add a trailing/leading space on one of the concatenated string literals.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447500043
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
Maybe this is a problem and can cause performance regression, let me check it.
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
@pjfanning https://github.com/apache/pekko/pull/3116
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
I will check this soon.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447514593
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
Rechecked this. `AtomicIntegerFieldUpdater` would avoid the field-retention concern, but it fails at runtime for this Scala-generated private field with `IllegalAccessException`. I opened #3116 to use a plain `AtomicInteger` instead, keeping the same IDLE/SCHEDULED CAS protocol while avoiding the VarHandle/private-field concern.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3343514000
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,110 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
+
+ // Reused across all drain batches; allocated once at construction.
+ private[this] val drainCallback: Any => Unit = (_: Any) => drain()
+
+ override def apply(pair: (ActorRef, Any)): Unit = {
+ add(pair) // Vyukov producer path: getAndSet + release-store, no CAS spin
+ // Double-checked CAS: uncontended fast path is one volatile read; only the IDLE->SCHEDULED winner
+ // pays a CAS + mailbox push.
+ if (state.get() == SchedStateIdle && state.compareAndSet(SchedStateIdle, SchedStateScheduled))
+ scheduleDrain()
+ }
Review Comment:
Fixed in da487e3c74: producer-side guard added in `LazyDispatch.apply`. It now pre-checks `interpreter.isStageCompleted(logic)` and drops the message (matching the original per-tell behaviour where `runAsyncInput` silently ignored post-completion sends), and re-checks after winning the IDLE→SCHEDULED CAS to reset state if completion landed in the race window. Bounds any leak to messages enqueued before completion becomes visible to the producer thread.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "pjfanning (via GitHub)" <gi...@apache.org>.
pjfanning commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3347850380
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,152 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
+
+ override def apply(msg: Any): Unit = {
+ // Drain path: onAsyncInput passes null as the message when scheduling a drain envelope.
+ if (msg.asInstanceOf[AnyRef] eq null) { drain(); return }
Review Comment:
why do you need to cast before checking `eq null`? can you also avoid `;` and just use new lines? can you avoid return here by using an `else` on the `if`.
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,152 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
+
+ override def apply(msg: Any): Unit = {
+ // Drain path: onAsyncInput passes null as the message when scheduling a drain envelope.
+ if (msg.asInstanceOf[AnyRef] eq null) { drain(); return }
+ // Producer path: msg is an (ActorRef, Any) tuple from the FunctionRef.
+ val pair = msg.asInstanceOf[(ActorRef, Any)]
+ // Producer-side completion guard. `runAsyncInput` short-circuits the drain handler when the stage
+ // is completed, so without this check a CAS-winning producer would leave state=SCHEDULED forever
+ // and subsequent producers would skip the mailbox push, growing the queue unbounded. Dropping
+ // post-completion sends here matches the original per-tell behaviour where `runAsyncInput`
+ // silently ignored them. The read is racy against the interpreter thread that flips
+ // `shutdownCounter`, but eventually visible — enough to bound the leak to messages enqueued
+ // before completion becomes visible to this thread.
+ if (interpreter.isStageCompleted(logic)) return
+ add(pair) // Vyukov producer path: getAndSet + release-store, no CAS spin
+ // Double-checked CAS: uncontended fast path is one volatile read; only the IDLE->SCHEDULED winner
+ // pays a CAS + mailbox push.
+ val u = LazyDispatch.stateHandle
+ // Scala 3's strict inference cannot pick the Int-returning signature-polymorphic overload of
+ // `VarHandle.get` without explicit return-type context — a typed local witnesses it. On Scala 2
+ // this is a no-op. Keep the double-checked plain read fast path: under producer contention
+ // only the IDLE→SCHEDULED winner pays a CAS, the rest just read.
+ val cur: Int = u.get(this)
+ if (cur == SchedStateIdle && u.compareAndSet(this, SchedStateIdle, SchedStateScheduled)) {
+ // Re-check after winning election: if completion landed between the initial guard and the CAS,
+ // `scheduleDrain` would post an envelope that `runAsyncInput` skips, leaving state=SCHEDULED.
+ // Reset to IDLE and skip the schedule.
+ if (interpreter.isStageCompleted(logic)) u.set(this, SchedStateIdle)
+ else scheduleDrain()
+ }
+ }
+
+ private def scheduleDrain(): Unit =
+ // 1 AsyncInput + 1 Envelope per drain batch (amortized across up to drainBatchSize tells).
+ // `this` serves as the drain callback (Any => Unit); onAsyncInput calls apply(null).
+ interpreter.onAsyncInput(logic, null, NoPromise, this)
+
+ private def drain(): Unit = {
+ val u = LazyDispatch.stateHandle
+ val limit = drainBatchSize // hoisted to a local so JIT treats it as a loop-invariant constant
+ var processed = 0
+ while (processed < limit) {
+ if (interpreter.isStageCompleted(logic)) {
+ // Stage completed mid-drain; drop the remainder (matches the original per-tell behaviour where
+ // runAsyncInput silently skipped completed stages). Don't reschedule — no future drain will run.
+ while (poll() ne null) ()
+ u.set(this, SchedStateIdle)
+ return
+ }
+ val item = poll()
+ if (item eq null) {
+ u.set(this, SchedStateIdle)
+ // Recheck race: a producer may have added between `poll == null` and the IDLE publish above.
+ // That producer saw state=SCHEDULED and skipped the mailbox send, so we must re-elect.
+ if (!isEmpty && u.compareAndSet(this, SchedStateIdle, SchedStateScheduled))
+ scheduleDrain()
+ return
+ }
+ handler(item)
+ processed += 1
+ }
+ // Hit batch cap with items potentially still queued. The last `handler(item)` call may have
+ // completed the stage (e.g. user code called `completeStage()`); a fresh `scheduleDrain` would
+ // post an envelope that `runAsyncInput` skips, leaving state=SCHEDULED forever. Mirror the
+ // mid-loop branch: drain remainder, publish IDLE, do not reschedule.
+ if (interpreter.isStageCompleted(logic)) {
+ while (poll() ne null) ()
+ u.set(this, SchedStateIdle)
+ return
+ }
+ // Re-schedule another envelope so other BoundaryEvents (pull/push/complete) can interleave via
+ // the actor mailbox. `state` stays SCHEDULED: concurrent producers observe SCHEDULED and skip;
+ // the new envelope will drain.
+ scheduleDrain()
+ }
+ }
+
+ object LazyDispatch {
+ private val stateHandle: VarHandle = {
+ val lookup = MethodHandles.privateLookupIn(classOf[LazyDispatch], MethodHandles.lookup())
Review Comment:
this is our own code - why do we need to use var handles? Can't we just change LazyDispatch to expose methods that work with the 'state' field?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "pjfanning (via GitHub)" <gi...@apache.org>.
pjfanning commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3410499697
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
https://github.com/apache/pekko/blob/main/actor/src/main/scala/org/apache/pekko/actor/RepointableActorRef.scala#L60-L65
In existing code, where we use VarHandles to do the sets and gets and don't access the var directly, we expose some unused private code that references the underlying field due to fears that the field might be erased by the JVM as an optimisation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447514593
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
Rechecked this and updated #3116. Instead of switching to `AtomicInteger`, the PR now keeps the VarHandle-backed `@volatile var state`, uses `getVolatile` / `setVolatile` / `compareAndSet`, and retains a direct private reference to the field (`_preventPrivateUnusedErasure`) like the existing Pekko VarHandle fields. This avoids per-instance `AtomicInteger` allocation while addressing the private-field retention concern.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447514593
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
Rechecked this and updated #3116. The PR now keeps the VarHandle-backed `@volatile var state` and uses explicit `getVolatile` / `setVolatile` / `compareAndSet` instead of plain VarHandle access. This avoids per-instance `AtomicInteger` allocation while making the state access mode explicit and conservative.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447516733
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
@pjfanning https://github.com/apache/pekko/pull/3116
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin closed pull request #3035: Optimize lazy stage actor dispatch
URL: https://github.com/apache/pekko/pull/3035
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "pjfanning (via GitHub)" <gi...@apache.org>.
pjfanning commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3410499697
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
https://github.com/apache/pekko/blob/main/actor/src/main/scala/org/apache/pekko/actor/RepointableActorRef.scala#L60-L65
In existing code (like the example in RepointableActorRef), where we use VarHandles to do the sets and gets and don't access the var directly, we expose some unused private code that references the underlying field due to fears that the field might be erased by the JVM as an optimisation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin merged PR #3035:
URL: https://github.com/apache/pekko/pull/3035
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447500043
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
Maybe this is a problem and can cause performance regression, let me check it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3345691257
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,123 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
+
+ // Reused across all drain batches; allocated once at construction.
+ private[this] val drainCallback: Any => Unit = (_: Any) => drain()
+
+ override def apply(pair: (ActorRef, Any)): Unit = {
+ // Producer-side completion guard. `runAsyncInput` short-circuits the drain handler when the stage
+ // is completed, so without this check a CAS-winning producer would leave state=SCHEDULED forever
+ // and subsequent producers would skip the mailbox push, growing the queue unbounded. Dropping
+ // post-completion sends here matches the original per-tell behaviour where `runAsyncInput`
+ // silently ignored them. The read is racy against the interpreter thread that flips
+ // `shutdownCounter`, but eventually visible — enough to bound the leak to messages enqueued
+ // before completion becomes visible to this thread.
+ if (interpreter.isStageCompleted(logic)) return
+ add(pair) // Vyukov producer path: getAndSet + release-store, no CAS spin
+ // Double-checked CAS: uncontended fast path is one volatile read; only the IDLE->SCHEDULED winner
+ // pays a CAS + mailbox push.
+ if (state.get() == SchedStateIdle && state.compareAndSet(SchedStateIdle, SchedStateScheduled)) {
+ // Re-check after winning election: if completion landed between the initial guard and the CAS,
+ // `scheduleDrain` would post an envelope that `runAsyncInput` skips, leaving state=SCHEDULED.
+ // Reset to IDLE and skip the schedule.
+ if (interpreter.isStageCompleted(logic)) state.set(SchedStateIdle)
+ else scheduleDrain()
+ }
+ }
+
+ private def scheduleDrain(): Unit =
+ // 1 AsyncInput + 1 Envelope per drain batch (amortized across up to drainBatchSize tells).
+ interpreter.onAsyncInput(logic, null, NoPromise, drainCallback)
+
+ private def drain(): Unit = {
+ val limit = drainBatchSize // hoisted to a local so JIT treats it as a loop-invariant constant
+ var processed = 0
+ while (processed < limit) {
+ if (interpreter.isStageCompleted(logic)) {
+ // Stage completed mid-drain; drop the remainder (matches the original per-tell behaviour where
+ // runAsyncInput silently skipped completed stages). Don't reschedule — no future drain will run.
+ while (poll() ne null) ()
+ state.set(SchedStateIdle)
+ return
+ }
+ val item = poll()
+ if (item eq null) {
+ state.set(SchedStateIdle)
+ // Recheck race: a producer may have added between `poll == null` and the IDLE publish above.
+ // That producer saw state=SCHEDULED and skipped the mailbox send, so we must re-elect.
+ if (!isEmpty && state.compareAndSet(SchedStateIdle, SchedStateScheduled))
+ scheduleDrain()
+ return
+ }
+ handler(item)
+ processed += 1
+ }
+ // Hit batch cap with items potentially still queued. Re-schedule another envelope so other
+ // BoundaryEvents (pull/push/complete) can interleave via the actor mailbox. `scheduledState` stays
+ // SCHEDULED: concurrent producers correctly observe SCHEDULED and skip; the new envelope will drain.
+ scheduleDrain()
Review Comment:
Fixed in ad3a9d77f0: added the symmetric `isStageCompleted` check at the post-batch tail before re-scheduling. If the last `handler(item)` call completed the stage, drain the remainder, publish IDLE and return — same as the existing mid-loop completion branch. Since `drain()` runs on the interpreter thread and completion within drain() can only happen via our own `handler` call, the check is race-free.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3343523873
##########
stream/src/main/java/org/apache/pekko/stream/stage/AbstractMpscDispatch.java:
##########
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * license agreements; and to You under the Apache License, version 2.0:
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * This file is part of the Apache Pekko project, which was derived from Akka.
+ */
+
+package org.apache.pekko.stream.stage;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+import org.apache.pekko.dispatch.AbstractNodeQueue;
+
+/**
+ * INTERNAL API.
+ *
+ * <p>Vyukov MPSC linked queue augmented with a volatile {@code int} election slot used by
+ * single-consumer drain coalescing: producers race to elect a drainer via IDLE -> SCHEDULED CAS
+ * on the state field. The state is driven through a {@link VarHandle} declared here so Scala
+ * callers can use the typed {@link #stateAcquire()} / {@link #stateSetRelease(int)} / {@link
+ * #stateCompareAndSet(int, int)} accessors without hitting Scala's polymorphic-signature inference
+ * quirks for {@code VarHandle.getAcquire} (which differ between Scala 2 and Scala 3). Same layering
+ * as {@code AbstractMailbox} for {@code Mailbox}.
+ */
+public abstract class AbstractMpscDispatch<T> extends AbstractNodeQueue<T> {
+
+ @SuppressWarnings("unused")
+ private volatile int _stateDoNotCallMeDirectly = 0;
+
+ private static final VarHandle stateHandle;
+
+ static {
+ try {
+ stateHandle =
+ MethodHandles.privateLookupIn(AbstractMpscDispatch.class, MethodHandles.lookup())
+ .findVarHandle(AbstractMpscDispatch.class, "_stateDoNotCallMeDirectly", int.class);
+ } catch (Throwable t) {
+ throw new ExceptionInInitializerError(t);
+ }
+ }
+
+ /** Acquire-load of the election state. Use only from the consumer (drain) side. */
+ public final int stateAcquire() {
+ return (int) stateHandle.getAcquire(this);
+ }
Review Comment:
Stale — the `AbstractMpscDispatch.java` helper was removed in commit 1b533f8 ("fix: keep lazy stage actor refs under supervisor") when the dispatch was refactored to a pure Scala `LazyDispatch` that extends `AbstractNodeQueue` directly and uses `AtomicInteger` for the election state, dropping the Java helper entirely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on PR #3035:
URL: https://github.com/apache/pekko/pull/3035#issuecomment-4702417420
I think this PR is ready.
refs https://github.com/akkadotnet/akka.net/pull/8263
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "pjfanning (via GitHub)" <gi...@apache.org>.
pjfanning commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3447487752
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +309,135 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - Implements `Any => Unit` directly — serves as both the producer callback and the drain callback,
+ * eliminating the separate `drainCallback` lambda allocation.
+ * - `state` is a plain `@volatile var Int` driven by a static `VarHandle` in the companion object
+ * (via `MethodHandles.privateLookupIn`), same pattern as `AbstractNodeQueue` itself;
+ * avoids per-instance `AtomicInteger`.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (Any => Unit) {
+
+ // IDLE/SCHEDULED election state. VarHandle avoids per-instance AtomicInteger;
+ // the handle lives in the companion object (true JVM static), same pattern as
+ // AbstractNodeQueue._tailDoNotCallMeDirectly.
+ @volatile var state: Int = SchedStateIdle
Review Comment:
@He-Pin is there a risk here? That we could have the hotspot compiler remove this field because it might appear not to be used.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on PR #3035:
URL: https://github.com/apache/pekko/pull/3035#issuecomment-4759966196
@pjfanning This is not related akka.net, I read the issue from akka-core.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3345692957
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -275,6 +308,123 @@ object GraphStageLogic {
type Receive = ((ActorRef, Any)) => Unit
}
+ private object StageActor {
+ def localCell(ref: ActorRef, description: String): ActorCell =
+ ref match {
+ case ref: LocalActorRef => ref.underlying
+ case ref: RepointableActorRef =>
+ ref.underlying match {
+ case cell: ActorCell => cell
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+ case unknown =>
+ throw new IllegalStateException(s"$description must be a local actor, was [${unknown.getClass.getName}]")
+ }
+
+ /**
+ * Reads `pekko.stream.materializer.stage-actor-drain-batch` from the materializer's ActorSystem config.
+ * Called once per lazy StageActor construction (never on the hot path). Bounded to `>= 1`.
+ */
+ def drainBatchSize(materializer: Materializer): Int =
+ Math.max(1, materializer.system.settings.config.getInt("pekko.stream.materializer.stage-actor-drain-batch"))
+
+ private final val SchedStateIdle: Int = 0
+ private final val SchedStateScheduled: Int = 1
+
+ /**
+ * Lazy-path dispatch: producers enqueue into a Vyukov MPSC queue and elect a single drain via
+ * IDLE -> SCHEDULED CAS; only the elected producer pays a mailbox enqueue. The drain runs on the
+ * interpreter thread, polls in a tight loop bounded by `drainBatchSize`, then either publishes IDLE
+ * (with a recheck for the publish-window race) or re-schedules another envelope to yield to other
+ * BoundaryEvents.
+ *
+ * JIT/GC notes:
+ * - `final class` + monomorphic per-StageActor instance → JIT devirtualizes the apply at the
+ * FunctionRef call site.
+ * - Extends `AbstractNodeQueue` directly so the queue head atomic and the dispatch function share one
+ * object (one allocation per StageActor, one fewer field deref on the producer hot path).
+ * - All hot-path state is `private[this]` → direct field access, no accessor methods.
+ * - `drainBatchSize` is read once into a stack-local at the top of `drain` so the JIT can treat the loop
+ * bound as a constant.
+ * - Per-tell allocation = 1 Node (`AbstractNodeQueue.Node`, ~24 bytes) + 1 Tuple2 (~24 bytes). The
+ * Tuple2 is forced by the public `StageActorRef.Receive` type. No AsyncInput / Envelope per tell —
+ * those are amortized across the batch.
+ */
+ // Not marked `private` so that `class StageActor`'s aux constructor (compiled outside of the companion
+ // object on Scala 3) can reference it; the enclosing `object StageActor` is itself private.
+ final class LazyDispatch(
+ interpreter: GraphInterpreter,
+ logic: GraphStageLogic,
+ handler: Any => Unit,
+ drainBatchSize: Int)
+ extends AbstractNodeQueue[(ActorRef, Any)]
+ with (((ActorRef, Any)) => Unit) {
+
+ // IDLE/SCHEDULED election state. AtomicInteger gives us volatile read + CAS without the cross-Scala
+ // VarHandle / field-updater access fuss; the wrapper costs one extra reference per StageActor, which
+ // is negligible against the per-tell mailbox traffic we are saving.
+ private[this] val state = new AtomicInteger(SchedStateIdle)
Review Comment:
Updated the PR description to match the implementation. The `@volatile var Int` + `VarHandle` design lived in `AbstractMpscDispatch.java` in an earlier revision, but the helper hit cross-Scala polymorphic-signature inference quirks for `VarHandle.getAcquire` between Scala 2 and Scala 3, so commit 1b533f8 dropped it in favour of a single `AtomicInteger` field. The extra reference per `StageActor` is negligible against the per-tell mailbox traffic this PR already saves. Inline comment at GraphStage.scala:364 also explains the rationale.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org
Re: [PR] Optimize lazy stage actor dispatch [pekko]
Posted by "He-Pin (via GitHub)" <gi...@apache.org>.
He-Pin commented on code in PR #3035:
URL: https://github.com/apache/pekko/pull/3035#discussion_r3343520738
##########
stream/src/main/scala/org/apache/pekko/stream/stage/GraphStage.scala:
##########
@@ -1382,7 +1529,7 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount:
* Override and return a name to be given to the StageActor of this operator.
*
* This method will be only invoked and used once, during the first [[getStageActor]]
- * invocation whichc reates the actor, since subsequent `getStageActors` calls function
+ * invocation which creates the actor, since subsequent `getStageActors` calls function
* like `become`, rather than creating new actors.
Review Comment:
Stale — current scaladoc at GraphStage.scala:1534-1535 already uses the singular `getStageActor`, no `getStageActors` (plural) remains in the file. Likely fixed during an earlier rebase.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@pekko.apache.org
For additional commands, e-mail: notifications-help@pekko.apache.org