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 612851ec4e fix: add supervision strategy support for batch costFn
(#3184)
612851ec4e is described below
commit 612851ec4eaaa988f2465523640d2f563dcb883c
Author: He-Pin(kerr) <[email protected]>
AuthorDate: Sat Jun 27 09:51:05 2026 +0800
fix: add supervision strategy support for batch costFn (#3184)
* fix: add supervision strategy support for batch costFn
Motivation:
The batch operator's costFn was called outside the supervised boundary in
onPush,
so costFn exceptions failed the stream unconditionally. While adding this
support,
we also found coupled exceptional-path hazards around pending-element
handling
(flush/EOS), where state could leak across batches or a closed inlet could
be
pulled after stage failure.
Modification:
Wrap costFn(elem) in Batch.onPush() with supervision handling: Stop fails,
Resume drops the element, Restart resets batch state and continues.
Harden pending-element exceptional paths in Batch:
- clear agg immediately after pushing a batch before processing pending
- evaluate costFn(pending) before seed(pending) in flush()
- avoid pull(in) when the inlet is already closed
- at EOS, complete the stage when pending is dropped by Resume/Restart after
a seed failure
Expand directional regression coverage in FlowBatchSpec:
- explicit Stop + default Stop for costFn throws
- Resume vs Restart semantic distinction
- pending/flush second costFn-evaluation failures under Resume/Restart/Stop
- EOS pending-seed failure under Restart
Result:
batch/batchWeighted now consistently applies supervision for costFn and
preserves
correct stage state/pull behavior across exceptional pending paths.
Tests:
- sbt "stream-tests/Test/testOnly
org.apache.pekko.stream.scaladsl.FlowBatchSpec" -- 14/14 passed
- sbt "stream/mimaReportBinaryIssues" -- clean
References:
Refs #3110
* fix: harden Batch supervision Stop paths, document semantics, and add
directional tests
Motivation:
Review of the batch costFn supervision support identified three latent
correctness risks around the Stop supervision path, missing user-facing
documentation, and gaps in directional test coverage for seed/aggregate
failures in onPush.
Modification:
- Add `return` after `failStage(ex)` in three Batch catch blocks
(flush pending-costFn Stop, onPush seed Stop, onPush aggregate Stop)
for defensive parity with the costFn Stop handler. Without the early
return, a user-configured CancellationStrategy.AfterDelay could leave
the inlet open after failStage, causing flush()/pull(in) to run on a
stage mid-teardown.
- Document supervision semantics for costFn/seed/aggregate in batch.md
and batchWeighted.md.
- Add two directional tests:
* resume when seed throws in onPush (exercises agg==null path under
Resume)
* resume when aggregate throws in onPush with downstream backpressure
(exercises the aggregate branch under Resume and verifies the
accumulated batch is preserved when the offending element is dropped)
Result:
Batch's supervision handling is defensively consistent across all Stop
paths, its docs describe the three-strategy behavior, and the seed and
aggregate failure paths in onPush have directional test coverage.
Tests:
- sbt "stream-tests/Test/testOnly
org.apache.pekko.stream.scaladsl.FlowBatchSpec
org.apache.pekko.stream.scaladsl.FlowBatchWeightedSpec" (17/17 passed)
References:
Refs #3110
---
.../stream/operators/Source-or-Flow/batch.md | 5 +
.../operators/Source-or-Flow/batchWeighted.md | 5 +
.../pekko/stream/scaladsl/FlowBatchSpec.scala | 275 +++++++++++++++++++++
.../org/apache/pekko/stream/impl/fusing/Ops.scala | 45 +++-
4 files changed, 321 insertions(+), 9 deletions(-)
diff --git a/docs/src/main/paradox/stream/operators/Source-or-Flow/batch.md
b/docs/src/main/paradox/stream/operators/Source-or-Flow/batch.md
index 79c862a8ec..08646f7c40 100644
--- a/docs/src/main/paradox/stream/operators/Source-or-Flow/batch.md
+++ b/docs/src/main/paradox/stream/operators/Source-or-Flow/batch.md
@@ -23,6 +23,11 @@ to the summary type.
Will eagerly pull elements, this behavior may result in a single pending (i.e.
buffered) element which cannot be
aggregated to the batched value.
+This operator adheres to the `ActorAttributes.SupervisionStrategy` attribute
for exceptions thrown by the `costFn`,
+`seed`, or `aggregate` functions. On `Supervision.Stop` (the default) the
stream fails. On `Supervision.Resume` the
+offending element is dropped and the accumulated batch (if any) is preserved.
On `Supervision.Restart` the offending
+element is dropped, the accumulated batch is discarded, and the operator
starts fresh.
+
## Reactive Streams semantics
@@@div { .callout }
diff --git
a/docs/src/main/paradox/stream/operators/Source-or-Flow/batchWeighted.md
b/docs/src/main/paradox/stream/operators/Source-or-Flow/batchWeighted.md
index 6db2f37f9f..4ff24c1341 100644
--- a/docs/src/main/paradox/stream/operators/Source-or-Flow/batchWeighted.md
+++ b/docs/src/main/paradox/stream/operators/Source-or-Flow/batchWeighted.md
@@ -21,6 +21,11 @@ backpressure.
Will eagerly pull elements, this behavior may result in a single pending (i.e.
buffered) element which cannot be
aggregated to the batched value.
+This operator adheres to the `ActorAttributes.SupervisionStrategy` attribute
for exceptions thrown by the `costFn`,
+`seed`, or `aggregate` functions. On `Supervision.Stop` (the default) the
stream fails. On `Supervision.Resume` the
+offending element is dropped and the accumulated batch (if any) is preserved.
On `Supervision.Restart` the offending
+element is dropped, the accumulated batch is discarded, and the operator
starts fresh.
+
## Reactive Streams semantics
@@@div { .callout }
diff --git
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowBatchSpec.scala
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowBatchSpec.scala
index feb3829861..7360c15c3c 100644
---
a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowBatchSpec.scala
+++
b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowBatchSpec.scala
@@ -19,7 +19,9 @@ import scala.concurrent.Await
import scala.concurrent.duration._
import org.apache.pekko
+import pekko.stream.ActorAttributes
import pekko.stream.OverflowStrategy
+import pekko.stream.Supervision
import pekko.stream.testkit._
class FlowBatchSpec extends StreamSpec("""
@@ -120,5 +122,278 @@ class FlowBatchSpec extends StreamSpec("""
Await.result(future, 3.seconds) should be((1 to 50).sum)
}
+ "fail stream when costFn throws and supervision is Stop" in {
+ val ex = new RuntimeException("boom")
+ val result = Source(1 to 5)
+ .batchWeighted(
+ max = 10,
+ costFn = (i: Int) => if (i == 3) throw ex else i.toLong,
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.stoppingDecider))
+ .runWith(Sink.ignore)
+ result.failed.futureValue shouldBe ex
+ }
+
+ "fail stream when costFn throws and supervision defaults to Stop" in {
+ // No supervision attribute is set, so the mandatory SupervisionStrategy
attribute
+ // resolves to the default stopping decider. This guards the default
code path.
+ val ex = new RuntimeException("boom")
+ val result = Source(1 to 5)
+ .batchWeighted(
+ max = 10,
+ costFn = (i: Int) => if (i == 3) throw ex else i.toLong,
+ seed = (i: Int) => i)(aggregate = _ + _)
+ .runWith(Sink.ignore)
+ result.failed.futureValue shouldBe ex
+ }
+
+ "resume when costFn throws and supervision is Resume" in {
+ val future = Source(1 to 5)
+ .batchWeighted(
+ max = Long.MaxValue,
+ costFn = (i: Int) => if (i == 3) throw new RuntimeException("boom")
else i.toLong,
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .runFold(0)(_ + _)
+ Await.result(future, 3.seconds) should be(12) // 1 + 2 + 4 + 5, 3 is
skipped
+ }
+
+ // The next two tests are a directional pair: with downstream silent the
whole stream
+ // aggregates into a single batch, so the emitted value reveals whether
the partial
+ // batch accumulated before the failing element was kept (Resume) or
discarded (Restart).
+ "keep the accumulated batch when costFn throws and supervision is Resume"
in {
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = Long.MaxValue,
+ costFn = (i: Int) => if (i == 3) throw new RuntimeException("boom")
else 1L,
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ for (i <- 1 to 5) publisher.sendNext(i)
+ publisher.sendComplete()
+ subscriber.expectNoMessage(300.millis)
+ sub.request(1)
+ // 1 + 2 accumulated, element 3 skipped but the partial batch is
preserved, then 4 + 5 added
+ subscriber.expectNext(1 + 2 + 4 + 5)
+ subscriber.expectComplete()
+ }
+
+ "reset the accumulated batch when costFn throws and supervision is
Restart" in {
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = Long.MaxValue,
+ costFn = (i: Int) => if (i == 3) throw new RuntimeException("boom")
else 1L,
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.restartingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ for (i <- 1 to 5) publisher.sendNext(i)
+ publisher.sendComplete()
+ subscriber.expectNoMessage(300.millis)
+ sub.request(1)
+ // 1 + 2 accumulated, then discarded on Restart, element 3 skipped, only
4 + 5 remain
+ subscriber.expectNext(4 + 5)
+ subscriber.expectComplete()
+ }
+
+ "restart and continue when pending element costFn throws during flush" in {
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+ var callsForThree = 0
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = 3,
+ costFn = (i: Int) => {
+ if (i == 3) {
+ callsForThree += 1
+ if (callsForThree == 2) throw new RuntimeException("boom")
+ }
+ i.toLong
+ },
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.restartingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ // 3 is evaluated once in onPush (becomes pending because left < cost)
and again in flush.
+ publisher.sendNext(1)
+ publisher.sendNext(2)
+ publisher.sendNext(3)
+ subscriber.expectNoMessage(300.millis)
+
+ sub.request(1)
+ subscriber.expectNext(3)
+
+ // After restart in flush(), the stage must still pull and continue
processing new elements.
+ publisher.sendNext(4)
+ publisher.sendComplete()
+ sub.request(1)
+ subscriber.expectNext(4)
+ subscriber.expectComplete()
+
+ callsForThree shouldBe 2
+ }
+
+ "resume and continue when pending element costFn throws during flush" in {
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+ var callsForThree = 0
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = 2,
+ costFn = (i: Int) => {
+ if (i == 3) {
+ callsForThree += 1
+ if (callsForThree == 2) throw new RuntimeException("boom")
+ }
+ 1L
+ },
+ seed = (i: Int) => i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ // 3 is evaluated once in onPush (becomes pending) and again in flush.
+ publisher.sendNext(1)
+ publisher.sendNext(2)
+ publisher.sendNext(3)
+ subscriber.expectNoMessage(300.millis)
+
+ sub.request(1)
+ subscriber.expectNext(3)
+
+ // On Resume, pending should be dropped without leaking the previously
emitted batch into new state.
+ publisher.sendNext(4)
+ publisher.sendComplete()
+ sub.request(1)
+ subscriber.expectNext(4)
+ subscriber.expectComplete()
+
+ callsForThree shouldBe 2
+ }
+
+ "stop with the original exception when pending element costFn throws
during flush" in {
+ val ex = new RuntimeException("boom")
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+ var callsForThree = 0
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = 2,
+ costFn = (i: Int) => {
+ if (i == 3) {
+ callsForThree += 1
+ if (callsForThree == 2) throw ex
+ }
+ 1L
+ },
+ seed = (i: Int) => i)(aggregate = _ + _)
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ publisher.sendNext(1)
+ publisher.sendNext(2)
+ publisher.sendNext(3)
+ subscriber.expectNoMessage(300.millis)
+
+ sub.request(1)
+ subscriber.expectNext(3)
+ subscriber.expectError(ex)
+
+ callsForThree shouldBe 2
+ }
+
+ "complete when pending seed throws at end of stream and supervision is
Restart" in {
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = 1,
+ costFn = (_: Int) => 1L,
+ seed = (i: Int) => if (i == 2) throw new RuntimeException("boom")
else i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.restartingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ publisher.sendNext(1)
+ publisher.sendNext(2)
+ publisher.sendComplete()
+ sub.request(1)
+ subscriber.expectNext(1)
+ subscriber.expectComplete()
+ }
+
+ "resume when seed throws in onPush and supervision is Resume" in {
+ // seed is only called in onPush when agg == null (i.e. no active batch).
+ // Under Resume the failing element must be dropped and the next element
+ // should start a fresh batch.
+ val future = Source(1 to 5)
+ .batchWeighted(
+ max = Long.MaxValue,
+ costFn = (i: Int) => i.toLong,
+ seed = (i: Int) => if (i == 1) throw new RuntimeException("boom")
else i)(aggregate = _ + _)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .runFold(0)(_ + _)
+ Await.result(future, 3.seconds) should be(2 + 3 + 4 + 5)
+ }
+
+ "resume when aggregate throws in onPush and supervision is Resume" in {
+ // aggregate is only called in onPush when downstream backpressures and
a batch
+ // is accumulating. Use a manual subscriber to force batching. Under
Resume
+ // the failing element must be dropped while the accumulated batch is
preserved.
+ val publisher = TestPublisher.probe[Int]()
+ val subscriber = TestSubscriber.manualProbe[Int]()
+
+ Source
+ .fromPublisher(publisher)
+ .batchWeighted(
+ max = 10,
+ costFn = (_: Int) => 1L,
+ seed = (i: Int) => i)(aggregate =
+ (acc, i) => if (i == 3) throw new RuntimeException("boom") else acc
+ i)
+
.withAttributes(ActorAttributes.supervisionStrategy(Supervision.resumingDecider))
+ .to(Sink.fromSubscriber(subscriber))
+ .run()
+ val sub = subscriber.expectSubscription()
+
+ // Send elements without requesting; they accumulate in a single batch.
+ publisher.sendNext(1) // seed -> agg=1
+ publisher.sendNext(2) // aggregate(1,2) -> agg=3
+ publisher.sendNext(3) // aggregate(3,3) throws -> Resume, agg stays 3,
element 3 dropped
+ publisher.sendNext(4) // aggregate(3,4) -> agg=7
+ publisher.sendNext(5) // aggregate(7,5) -> agg=12
+ publisher.sendComplete()
+
+ subscriber.expectNoMessage(300.millis)
+ sub.request(1)
+ subscriber.expectNext(1 + 2 + 4 + 5)
+ subscriber.expectComplete()
+ }
+
}
}
diff --git
a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/Ops.scala
b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/Ops.scala
index 010260e03d..882ec381a8 100644
--- a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/Ops.scala
+++ b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/Ops.scala
@@ -1081,16 +1081,20 @@ private[stream] object Collect {
if (agg != null) {
push(out, agg)
left = max
+ agg = null.asInstanceOf[Out]
}
if (pending != null) {
try {
+ val cost = costFn(pending)
agg = seed(pending)
- left -= costFn(pending)
+ left -= cost
pending = null.asInstanceOf[In]
} catch {
case NonFatal(ex) =>
decider(ex) match {
- case Supervision.Stop => failStage(ex)
+ case Supervision.Stop =>
+ failStage(ex)
+ return
case Supervision.Restart => restartState()
case Supervision.Resume =>
pending = null.asInstanceOf[In]
@@ -1105,7 +1109,23 @@ private[stream] object Collect {
def onPush(): Unit = {
val elem = grab(in)
- val cost = costFn(elem)
+ val cost =
+ try costFn(elem)
+ catch {
+ case NonFatal(ex) =>
+ decider(ex) match {
+ case Supervision.Stop =>
+ failStage(ex)
+ return
+ case Supervision.Resume =>
+ pull(in)
+ return
+ case Supervision.Restart =>
+ restartState()
+ pull(in)
+ return
+ }
+ }
contextPropagation.suspendContext()
if (agg == null) {
@@ -1115,7 +1135,9 @@ private[stream] object Collect {
} catch {
case NonFatal(ex) =>
decider(ex) match {
- case Supervision.Stop => failStage(ex)
+ case Supervision.Stop =>
+ failStage(ex)
+ return
case Supervision.Restart =>
restartState()
case Supervision.Resume =>
@@ -1130,7 +1152,9 @@ private[stream] object Collect {
} catch {
case NonFatal(ex) =>
decider(ex) match {
- case Supervision.Stop => failStage(ex)
+ case Supervision.Stop =>
+ failStage(ex)
+ return
case Supervision.Restart =>
restartState()
case Supervision.Resume =>
@@ -1139,7 +1163,7 @@ private[stream] object Collect {
}
if (isAvailable(out)) flush()
- if (pending == null) pull(in)
+ if (pending == null && !isClosed(in)) pull(in)
}
override def onUpstreamFinish(): Unit = {
@@ -1153,6 +1177,7 @@ private[stream] object Collect {
} else if (isClosed(in)) {
contextPropagation.resumeContext()
push(out, agg)
+ agg = null.asInstanceOf[Out]
if (pending == null) completeStage()
else {
try {
@@ -1160,19 +1185,21 @@ private[stream] object Collect {
} catch {
case NonFatal(ex) =>
decider(ex) match {
- case Supervision.Stop => failStage(ex)
+ case Supervision.Stop =>
+ failStage(ex)
+ return
case Supervision.Resume =>
case Supervision.Restart =>
restartState()
- if (!hasBeenPulled(in)) pull(in)
}
}
pending = null.asInstanceOf[In]
+ if (agg == null) completeStage()
}
} else {
contextPropagation.resumeContext()
flush()
- if (!hasBeenPulled(in)) pull(in)
+ if (!isClosed(in) && !hasBeenPulled(in)) pull(in)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
