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 07cda18251 feat: add supervision strategy support for doOnFirst (#3179)
07cda18251 is described below
commit 07cda182517b1eb0d6d53ac274488e5327af56b8
Author: He-Pin(kerr) <[email protected]>
AuthorDate: Sun Jun 28 02:44:48 2026 +0800
feat: add supervision strategy support for doOnFirst (#3179)
* feat: add supervision strategy support for doOnFirst
Motivation:
The doOnFirst operator runs a user-provided side-effect function on the
first
element, but it was called without a try-catch. Any exception failed the
stream unconditionally, ignoring the configured SupervisionStrategy.
Modification:
Wrap f(elem) in DoOnFirst with a try-catch that consults the
SupervisionStrategy decider. Stop fails the stage (unchanged), Resume drops
the failing first element and switches to pass-through without retrying the
function, and Restart re-arms first-element detection so the next element is
treated as the first. The decider is a lazy val for zero overhead on the
happy path. The first-element handler was extracted into a named val so it
can be re-armed on Restart. Also fix a fully-qualified type reference in the
createLogic signature to use the project's two-line import style, and update
the Scala/Java DSL scaladoc and the operator reference page.
Result:
doOnFirst now adheres to the SupervisionStrategy attribute with clear,
distinguishable Resume and Restart semantics.
Tests:
- sbt "stream-tests/Test/testOnly
org.apache.pekko.stream.scaladsl.FlowDoOnFirstSpec" -- 6/6 passed
References:
Refs #3110
* test: add directional test for doOnFirst restart that succeeds on retry
Motivation:
The existing restart test only covered the path where f throws for
every element. The more common restart scenario is a transient failure
on the first element followed by success, after which the operator
must switch to pass-through.
Modification:
Add "restart and continue with pass-through once f succeeds on retry"
which asserts that the second element is treated as the first (f runs
again), and that once f succeeds the operator transitions to the
pass-through handler so later elements do not re-invoke f.
Result:
The doOnFirst restart semantics are now covered for both the
all-failing and the succeed-on-retry paths.
Tests:
- sbt "stream-tests/Test/testOnly
org.apache.pekko.stream.scaladsl.FlowDoOnFirstSpec" (7/7 passed)
References:
Refs #3110
---
.../stream/operators/Source-or-Flow/doOnFirst.md | 2 +
.../pekko/stream/scaladsl/FlowDoOnFirstSpec.scala | 83 ++++++++++++++++++++--
.../pekko/stream/impl/fusing/DoOnFirst.scala | 42 ++++++++---
.../org/apache/pekko/stream/javadsl/Flow.scala | 4 ++
.../org/apache/pekko/stream/javadsl/Source.scala | 4 ++
.../org/apache/pekko/stream/javadsl/SubFlow.scala | 4 ++
.../apache/pekko/stream/javadsl/SubSource.scala | 4 ++
.../org/apache/pekko/stream/scaladsl/Flow.scala | 4 ++
8 files changed, 134 insertions(+), 13 deletions(-)
diff --git a/docs/src/main/paradox/stream/operators/Source-or-Flow/doOnFirst.md
b/docs/src/main/paradox/stream/operators/Source-or-Flow/doOnFirst.md
index 13a57301b3..42ac2e5355 100644
--- a/docs/src/main/paradox/stream/operators/Source-or-Flow/doOnFirst.md
+++ b/docs/src/main/paradox/stream/operators/Source-or-Flow/doOnFirst.md
@@ -13,6 +13,8 @@ Run the given function when the first element is received.
Run the given function when the first element is received.
+The `doOnFirst` operator adheres to the ActorAttributes.SupervisionStrategy
attribute. On `Supervision.Resume` the failing first element is dropped and the
function is not retried; on `Supervision.Restart` the next element is treated
as the first.
+
## Examples
Scala
diff --git
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowDoOnFirstSpec.scala
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowDoOnFirstSpec.scala
index 9c832d8003..6ced5fa70d 100644
---
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowDoOnFirstSpec.scala
+++
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowDoOnFirstSpec.scala
@@ -17,8 +17,12 @@
package org.apache.pekko.stream.scaladsl
-import org.apache.pekko.Done
-import org.apache.pekko.stream.testkit._
+import java.util.concurrent.atomic.AtomicInteger
+
+import org.apache.pekko
+import pekko.Done
+import pekko.stream.{ ActorAttributes, Supervision }
+import pekko.stream.testkit._
class FlowDoOnFirstSpec extends StreamSpec("""
pekko.stream.materializer.initial-input-buffer-size = 2
@@ -27,7 +31,7 @@ class FlowDoOnFirstSpec extends StreamSpec("""
"A DoOnFirst" must {
"can only invoke on first" in {
- val invoked = new java.util.concurrent.atomic.AtomicInteger(0)
+ val invoked = new AtomicInteger(0)
Source(1 to 10)
.via(Flow[Int].doOnFirst(invoked.addAndGet))
.runWith(Sink.ignore)
@@ -37,7 +41,7 @@ class FlowDoOnFirstSpec extends StreamSpec("""
}
"will not invoke on empty stream" in {
- val invoked = new java.util.concurrent.atomic.AtomicInteger(0)
+ val invoked = new AtomicInteger(0)
Source.empty
.via(Flow[Int].doOnFirst(invoked.addAndGet))
.runWith(Sink.ignore)
@@ -46,6 +50,77 @@ class FlowDoOnFirstSpec extends StreamSpec("""
invoked.get() shouldBe 0
}
+ "fail the stream when f throws and supervision is Stop" in {
+ val ex = new RuntimeException("doOnFirst-stop")
+ val result = Source(1 to 5)
+ .via(Flow[Int].doOnFirst(_ => throw ex))
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.stoppingDecider))
+ .runWith(Sink.ignore)
+
+ result.failed.futureValue shouldBe ex
+ }
+
+ "fail the stream when f throws and no supervision attribute is set
(default Stop)" in {
+ val ex = new RuntimeException("doOnFirst-default")
+ val result = Source(1 to 5)
+ .via(Flow[Int].doOnFirst(_ => throw ex))
+ .runWith(Sink.ignore)
+
+ result.failed.futureValue shouldBe ex
+ }
+
+ "resume drops the first element, does not retry f, and passes the rest
through" in {
+ val invocationCount = new AtomicInteger(0)
+ val ex = new RuntimeException("doOnFirst-resume")
+
+ val result = Source(1 to 5)
+ .via(Flow[Int].doOnFirst { _ =>
+ if (invocationCount.incrementAndGet() == 1) throw ex
+ })
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .runWith(Sink.seq)
+ .futureValue
+
+ result shouldBe Seq(2, 3, 4, 5)
+ invocationCount.get() shouldBe 1
+ }
+
+ "restart re-arms so f is retried on the next element" in {
+ val invocationCount = new AtomicInteger(0)
+ val ex = new RuntimeException("doOnFirst-restart")
+
+ val result = Source(1 to 3)
+ .via(Flow[Int].doOnFirst { _ =>
+ invocationCount.incrementAndGet()
+ throw ex
+ })
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.restartingDecider))
+ .runWith(Sink.seq)
+ .futureValue
+
+ result shouldBe Seq.empty
+ invocationCount.get() shouldBe 3
+ }
+
+ "restart and continue with pass-through once f succeeds on retry" in {
+ // Element 1 triggers f which throws -> Restart drops element 1 and
re-arms.
+ // Element 2 triggers f which succeeds -> handler switches to
pass-through.
+ // Elements 3 and 4 pass through without invoking f.
+ val invocationCount = new AtomicInteger(0)
+ val ex = new RuntimeException("doOnFirst-restart-then-ok")
+
+ val result = Source(1 to 4)
+ .via(Flow[Int].doOnFirst { _ =>
+ if (invocationCount.incrementAndGet() == 1) throw ex
+ })
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.restartingDecider))
+ .runWith(Sink.seq)
+ .futureValue
+
+ result shouldBe Seq(2, 3, 4)
+ invocationCount.get() shouldBe 2
+ }
+
}
}
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/DoOnFirst.scala
b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/DoOnFirst.scala
index 009d15b5dc..5eebe57655 100644
--- a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/DoOnFirst.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/DoOnFirst.scala
@@ -17,10 +17,14 @@
package org.apache.pekko.stream.impl.fusing
+import scala.util.control.NonFatal
+
import org.apache.pekko
import pekko.annotation.InternalApi
+import pekko.stream.ActorAttributes.SupervisionStrategy
import pekko.stream.{ Attributes, FlowShape, Inlet, Outlet }
import pekko.stream.Attributes.SourceLocation
+import pekko.stream.Supervision
import pekko.stream.impl.Stages.DefaultAttributes
import pekko.stream.stage.{ GraphStage, GraphStageLogic, InHandler, OutHandler
}
@@ -34,20 +38,40 @@ import pekko.stream.stage.{ GraphStage, GraphStageLogic,
InHandler, OutHandler }
override def initialAttributes: Attributes = DefaultAttributes.doOnFirst and
SourceLocation.forLambda(f)
- override def createLogic(inheritedAttributes:
org.apache.pekko.stream.Attributes) =
+ override def createLogic(inheritedAttributes: Attributes) =
new GraphStageLogic(shape) with InHandler with OutHandler {
self =>
+ private lazy val decider =
inheritedAttributes.mandatoryAttribute[SupervisionStrategy].decider
+
final override def onPush(): Unit = push(out, grab(in))
final override def onPull(): Unit = pull(in)
setHandler(out, this)
- setHandler(in,
- new InHandler {
- override def onPush(): Unit = {
- setHandler(in, self)
- val elem = grab(in)
- f(elem)
- push(out, elem)
+
+ // Resume consumes the one-shot and switches to pass-through; Restart
keeps this handler armed for retry.
+ val firstHandler: InHandler = new InHandler {
+ override def onPush(): Unit = {
+ val elem = grab(in)
+ try f(elem)
+ catch {
+ case NonFatal(ex) =>
+ decider(ex) match {
+ case Supervision.Stop =>
+ failStage(ex)
+ return
+ case Supervision.Resume =>
+ setHandler(in, self)
+ pull(in)
+ return
+ case Supervision.Restart =>
+ pull(in)
+ return
+ }
}
- })
+ setHandler(in, self)
+ push(out, elem)
+ }
+ }
+
+ setHandler(in, firstHandler)
}
}
diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Flow.scala
b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Flow.scala
index cf3119f85b..2896116c72 100755
--- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Flow.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Flow.scala
@@ -1150,6 +1150,10 @@ final class Flow[In, Out, Mat](delegate:
scaladsl.Flow[In, Out, Mat]) extends Gr
/**
* Run the given function when the first element is received.
*
+ * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute. On
`Supervision.Resume` the
+ * first element is dropped and the function is not retried; on
`Supervision.Restart` the next
+ * element is treated as the first.
+ *
* '''Emits when''' upstream emits an element
*
* '''Backpressures when''' downstream backpressures
diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Source.scala
b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Source.scala
index 7dd27fbe52..0d7bd784a7 100755
--- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Source.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Source.scala
@@ -3046,6 +3046,10 @@ final class Source[Out, Mat](delegate:
scaladsl.Source[Out, Mat]) extends Graph[
/**
* Run the given function when the first element is received.
*
+ * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute. On
`Supervision.Resume` the
+ * first element is dropped and the function is not retried; on
`Supervision.Restart` the next
+ * element is treated as the first.
+ *
* '''Emits when''' upstream emits an element
*
* '''Backpressures when''' downstream backpressures
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubFlow.scala
b/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubFlow.scala
index ce766b5628..8bb0109a2d 100755
--- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubFlow.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubFlow.scala
@@ -552,6 +552,10 @@ final class SubFlow[In, Out, Mat](
/**
* Run the given function when the first element is received.
*
+ * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute. On
`Supervision.Resume` the
+ * first element is dropped and the function is not retried; on
`Supervision.Restart` the next
+ * element is treated as the first.
+ *
* '''Emits when''' upstream emits an element
*
* '''Backpressures when''' downstream backpressures
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubSource.scala
b/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubSource.scala
index 67d339a1c0..822fc76113 100755
--- a/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubSource.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/javadsl/SubSource.scala
@@ -543,6 +543,10 @@ final class SubSource[Out, Mat](
/**
* Run the given function when the first element is received.
*
+ * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute. On
`Supervision.Resume` the
+ * first element is dropped and the function is not retried; on
`Supervision.Restart` the next
+ * element is treated as the first.
+ *
* '''Emits when''' upstream emits an element
*
* '''Backpressures when''' downstream backpressures
diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Flow.scala
b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Flow.scala
index cbcb86432f..7fb3938e81 100755
--- a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Flow.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Flow.scala
@@ -1594,6 +1594,10 @@ trait FlowOps[+Out, +Mat] {
/**
* Run the given function when the first element is received.
*
+ * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute. On
`Supervision.Resume` the
+ * first element is dropped and the function is not retried; on
`Supervision.Restart` the next
+ * element is treated as the first.
+ *
* '''Emits when''' upstream emits an element
*
* '''Backpressures when''' downstream backpressures
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]