This is an automated email from the ASF dual-hosted git repository.
He-Pin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/pekko.git
The following commit(s) were added to refs/heads/main by this push:
new 3ac83fbd31 fix(stream): prevent double materialization of SourceRef
and SinkRef with clear error (#3114)
3ac83fbd31 is described below
commit 3ac83fbd31dd802c2aa5fb114ab44f5e44486517
Author: He-Pin(kerr) <[email protected]>
AuthorDate: Tue Jun 23 17:18:25 2026 +0800
fix(stream): prevent double materialization of SourceRef and SinkRef with
clear error (#3114)
Motivation:
SourceRefImpl.source and SinkRefImpl.sink() create a new stage instance on
every access. If a user accidentally materializes the returned Source/Sink
more than once, multiple stage instances compete for the same partner
handshake, causing the second materialization to fail with a confusing
error message.
Modification:
Add an AtomicBoolean guard per SourceRef/SinkRef instance. When the
backing stage is materialized, the guard is atomically claimed via
compareAndSet. A second materialization attempt throws a clear
IllegalStateException explaining that SourceRef/SinkRef is single-use,
instead of the previous opaque handshake failure.
Result:
Double materialization of SourceRef or SinkRef now fails fast with a
descriptive error: "This SourceRef/SinkRef has already been materialized.
SourceRef/SinkRef is single-use: calling .source/.sink() and materializing
it more than once is not supported." Concurrent materialization attempts
are also correctly rejected (exactly one succeeds, all others fail).
Tests:
- stream-tests/testOnly org.apache.pekko.stream.scaladsl.StreamRefsSpec:
27/27 passed, including sequential and concurrent (10-thread latch)
double-materialization tests for both SourceRef and SinkRef
References:
Fixes #3106
---
.../pekko/stream/scaladsl/StreamRefsSpec.scala | 63 +++++++++++++++++++---
.../pekko/stream/impl/streamref/SinkRefImpl.scala | 18 ++++++-
.../stream/impl/streamref/SourceRefImpl.scala | 18 ++++++-
3 files changed, 87 insertions(+), 12 deletions(-)
diff --git
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/StreamRefsSpec.scala
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/StreamRefsSpec.scala
index 70a006a749..d877cac479 100644
---
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/StreamRefsSpec.scala
+++
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/StreamRefsSpec.scala
@@ -17,8 +17,11 @@ import scala.collection.immutable
import scala.concurrent.{ Await, Future }
import scala.concurrent.Promise
import scala.concurrent.duration._
+import scala.util.Try
import scala.util.control.NoStackTrace
+import java.util.concurrent.CountDownLatch
+
import org.apache.pekko
import pekko.{ Done, NotUsed }
import pekko.actor.{ Actor, ActorIdentity, ActorLogging, ActorRef,
ActorSystem, ActorSystemImpl, Identify, Props }
@@ -26,7 +29,6 @@ import pekko.actor.Status.Failure
import pekko.pattern._
import pekko.stream._
import pekko.stream.impl.streamref.{ SinkRefImpl, SourceRefImpl }
-import pekko.stream.testkit.TestPublisher
import pekko.stream.testkit.Utils.TE
import pekko.stream.testkit.scaladsl._
import pekko.testkit.{ PekkoSpec, TestKit, TestProbe }
@@ -480,6 +482,36 @@ class StreamRefsSpec extends
PekkoSpec(StreamRefsSpec.config()) {
expectTerminated(sinkRefStageActorRef)
}
+ "not allow materializing multiple times" in {
+ val remoteProbe = TestProbe()(remoteSystem)
+ remoteActor.tell("give", remoteProbe.ref)
+ val sourceRef = remoteProbe.expectMsgType[SourceRef[String]]
+
+ sourceRef.source.to(Sink.ignore).run()
+
+ val ex = the[IllegalStateException] thrownBy
sourceRef.source.to(Sink.ignore).run()
+ ex.getMessage should include("already been materialized")
+ }
+
+ "not allow concurrent materialization" in {
+ val remoteProbe = TestProbe()(remoteSystem)
+ remoteActor.tell("give", remoteProbe.ref)
+ val sourceRef = remoteProbe.expectMsgType[SourceRef[String]]
+
+ implicit val ec: scala.concurrent.ExecutionContext = system.dispatcher
+ val latch = new CountDownLatch(1)
+ val results = (1 to 10).map { _ =>
+ Future {
+ latch.await()
+ Try(sourceRef.source.to(Sink.ignore).run())
+ }
+ }
+ latch.countDown()
+ val outcomes = Await.result(Future.sequence(results), 10.seconds)
+ outcomes.count(_.isSuccess) shouldBe 1
+ outcomes.count(_.isFailure) shouldBe 9
+ }
+
}
"A SinkRef" must {
@@ -613,15 +645,30 @@ class StreamRefsSpec extends
PekkoSpec(StreamRefsSpec.config()) {
remoteActor.tell(Command("receive", elementProbe.ref), remoteProbe.ref)
val sinkRef = remoteProbe.expectMsgType[SinkRef[String]]
- val p1: TestPublisher.Probe[String] =
TestSource[String]().to(sinkRef).run()
- val p2: TestPublisher.Probe[String] =
TestSource[String]().to(sinkRef).run()
+ TestSource[String]().to(sinkRef).run()
+
+ val ex = the[IllegalStateException] thrownBy
TestSource[String]().to(sinkRef).run()
+ ex.getMessage should include("already been materialized")
+ }
- p1.ensureSubscription()
- p1.expectRequest()
+ "not allow concurrent materialization" in {
+ val remoteProbe = TestProbe()(remoteSystem)
+ val elementProbe = TestProbe()(remoteSystem)
+ remoteActor.tell(Command("receive", elementProbe.ref), remoteProbe.ref)
+ val sinkRef = remoteProbe.expectMsgType[SinkRef[String]]
- // will be cancelled immediately, since it's 2nd:
- p2.ensureSubscription()
- p2.expectCancellation()
+ implicit val ec: scala.concurrent.ExecutionContext = system.dispatcher
+ val latch = new CountDownLatch(1)
+ val results = (1 to 10).map { _ =>
+ Future {
+ latch.await()
+ Try(TestSource[String]().to(sinkRef).run())
+ }
+ }
+ latch.countDown()
+ val outcomes = Await.result(Future.sequence(results), 10.seconds)
+ outcomes.count(_.isSuccess) shouldBe 1
+ outcomes.count(_.isFailure) shouldBe 9
}
}
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SinkRefImpl.scala
b/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SinkRefImpl.scala
index 53b8a89c8d..614a1f8f64 100644
---
a/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SinkRefImpl.scala
+++
b/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SinkRefImpl.scala
@@ -27,11 +27,19 @@ import pekko.stream.scaladsl.Sink
import pekko.stream.stage._
import pekko.util.{ OptionVal, PrettyDuration }
+import java.util.concurrent.atomic.AtomicBoolean
+
/** INTERNAL API: Implementation class, not intended to be touched directly by
end-users */
@InternalApi
private[stream] final case class SinkRefImpl[In](initialPartnerRef: ActorRef)
extends SinkRef[In] {
+ private val materialized = new AtomicBoolean(false)
+
+ private[stream] def tryClaimMaterialization(): Boolean =
+ materialized.compareAndSet(false, true)
+
override def sink(): Sink[In, NotUsed] =
- Sink.fromGraph(new
SinkRefStageImpl[In](OptionVal.Some(initialPartnerRef))).mapMaterializedValue(_
=> NotUsed)
+ Sink.fromGraph(new SinkRefStageImpl[In](OptionVal.Some(initialPartnerRef),
this))
+ .mapMaterializedValue(_ => NotUsed)
}
/**
@@ -48,7 +56,9 @@ private[stream] final case class
SinkRefImpl[In](initialPartnerRef: ActorRef) ex
* the ref.
*/
@InternalApi
-private[stream] final class SinkRefStageImpl[In] private[pekko] (val
initialPartnerRef: OptionVal[ActorRef])
+private[stream] final class SinkRefStageImpl[In] private[pekko] (
+ val initialPartnerRef: OptionVal[ActorRef],
+ private val owner: SinkRefImpl[?] = null)
extends GraphStageWithMaterializedValue[SinkShape[In], SourceRef[In]] {
import SinkRefStageImpl.ActorRefStage
@@ -67,6 +77,10 @@ private[stream] final class SinkRefStageImpl[In]
private[pekko] (val initialPart
private[pekko] override def createLogicAndMaterializedValue(
inheritedAttributes: Attributes,
eagerMaterializer: Materializer): (GraphStageLogic, SourceRef[In]) = {
+ if (owner != null && !owner.tryClaimMaterialization())
+ throw new IllegalStateException(
+ "This SinkRef has already been materialized. " +
+ "SinkRef is single-use: calling .sink() and materializing it more than
once is not supported.")
val logic = new TimerGraphStageLogic(shape) with StageLogging with
ActorRefStage with InHandler {
override protected def logSource: Class[?] = classOf[SinkRefStageImpl[?]]
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SourceRefImpl.scala
b/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SourceRefImpl.scala
index 4e392cc35b..59d383c5ad 100644
---
a/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SourceRefImpl.scala
+++
b/stream/src/main/scala/org/apache/pekko/stream/impl/streamref/SourceRefImpl.scala
@@ -26,11 +26,19 @@ import pekko.stream.scaladsl.Source
import pekko.stream.stage._
import pekko.util.{ OptionVal, PrettyDuration }
+import java.util.concurrent.atomic.AtomicBoolean
+
/** INTERNAL API: Implementation class, not intended to be touched directly by
end-users */
@InternalApi
private[stream] final case class SourceRefImpl[T](initialPartnerRef: ActorRef)
extends SourceRef[T] {
+ private val materialized = new AtomicBoolean(false)
+
+ private[stream] def tryClaimMaterialization(): Boolean =
+ materialized.compareAndSet(false, true)
+
def source: Source[T, NotUsed] =
- Source.fromGraph(new
SourceRefStageImpl(OptionVal.Some(initialPartnerRef))).mapMaterializedValue(_
=> NotUsed)
+ Source.fromGraph(new SourceRefStageImpl(OptionVal.Some(initialPartnerRef),
this))
+ .mapMaterializedValue(_ => NotUsed)
}
/**
@@ -107,7 +115,9 @@ private[stream] final case class
SourceRefImpl[T](initialPartnerRef: ActorRef) e
* If it is none, then we are the side creating the ref.
*/
@InternalApi
-private[stream] final class SourceRefStageImpl[Out](val initialPartnerRef:
OptionVal[ActorRef])
+private[stream] final class SourceRefStageImpl[Out](
+ val initialPartnerRef: OptionVal[ActorRef],
+ private val owner: SourceRefImpl[?] = null)
extends GraphStageWithMaterializedValue[SourceShape[Out], SinkRef[Out]] {
stage =>
import SourceRefStageImpl._
@@ -126,6 +136,10 @@ private[stream] final class SourceRefStageImpl[Out](val
initialPartnerRef: Optio
private[pekko] override def createLogicAndMaterializedValue(
inheritedAttributes: Attributes,
eagerMaterializer: Materializer): (GraphStageLogic, SinkRef[Out]) = {
+ if (owner != null && !owner.tryClaimMaterialization())
+ throw new IllegalStateException(
+ "This SourceRef has already been materialized. " +
+ "SourceRef is single-use: calling .source and materializing it more
than once is not supported.")
val logic = new TimerGraphStageLogic(shape) with StageLogging with
ActorRefStage with OutHandler {
override protected def logSource: Class[?] =
classOf[SourceRefStageImpl[?]]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
