diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/HubSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/HubSpec.scala index b32c8746b1..954c592da9 100644 --- a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/HubSpec.scala +++ b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/HubSpec.scala @@ -19,7 +19,7 @@ import scala.concurrent.duration._ import org.apache.pekko import pekko.Done -import pekko.stream.KillSwitches +import pekko.stream.{ KillSwitches, Materializer } import pekko.stream.ThrottleMode import pekko.stream.impl.ActorPublisher import pekko.stream.testkit.StreamSpec @@ -559,6 +559,28 @@ class HubSpec extends StreamSpec { } } + "notify consumers when hub materializer is shut down" in { + val hubMat = Materializer(system) + val upstream = TestPublisher.probe[Int]() + val source = Source.fromPublisher(upstream).runWith(BroadcastHub.sink(8))(hubMat) + + val downstream = TestSubscriber.probe[Int]() + source.runWith(Sink.fromSubscriber(downstream)) + + downstream.request(1) + upstream.expectRequest() + + upstream.sendNext(1) + downstream.expectNext(1) + + hubMat.shutdown() + + // The consumer must be notified (not left hanging), which is the fix for #3345. + // The signal is an error because the SubSink output boundary failure + // arrives before the postStop completion callback. + downstream.expectError() + } + "handle cancelled Sink" in { val in = TestPublisher.probe[Int]() val hubSource = Source.fromPublisher(in).runWith(BroadcastHub.sink(4)) diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Hub.scala b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Hub.scala index 32e647841f..ca7615a55f 100644 --- a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Hub.scala +++ b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Hub.scala @@ -25,6 +25,7 @@ import scala.collection.immutable.Queue import scala.collection.mutable.LongMap import scala.concurrent.{ Future, Promise } import scala.util.{ Failure, Success, Try } +import scala.util.control.NonFatal import org.apache.pekko import pekko.NotUsed @@ -652,21 +653,13 @@ private[pekko] class BroadcastHub[T](startAfterNrOfConsumers: Int, bufferSize: I override def onUpstreamFailure(ex: Throwable): Unit = { val failMessage = HubCompleted(Some(ex)) - // Notify pending consumers and set tombstone + // Notify pending consumers and set tombstone. + // Registered consumers in the consumerWheel are notified by postStop. state.getAndSet(Closed(Some(ex))).asInstanceOf[Open].registrations.foreach { consumer => - consumer.callback.invoke(failMessage) + try consumer.callback.invoke(failMessage) + catch { case NonFatal(_) => } } - // Notify registered consumers — skip null (empty) slots - var idx = 0 - while (idx < consumerWheel.length) { - val bucket = consumerWheel(idx) - if (bucket ne null) { - val itr = bucket.valuesIterator - while (itr.hasNext) itr.next().callback.invoke(failMessage) - } - idx += 1 - } failStage(ex) } @@ -759,19 +752,43 @@ private[pekko] class BroadcastHub[T](startAfterNrOfConsumers: Int, bufferSize: I } override def postStop(): Unit = { - // Notify pending consumers and set tombstone + // Notify all consumers (pending and registered) when the stage stops. + // Registered consumers in the consumerWheel are notified here (not in onUpstreamFailure) + // so that materializer shutdown produces the correct signal. @tailrec def tryClose(): Unit = state.get() match { - case Closed(_) => // Already closed, ignore + case Closed(_) => // Already closed by onUpstreamFailure — notify registered consumers directly + notifyRegisteredConsumers() case open: Open => if (state.compareAndSet(open, Closed(None))) { val completedMessage = HubCompleted(None) open.registrations.foreach { consumer => - consumer.callback.invoke(completedMessage) + try consumer.callback.invoke(completedMessage) + catch { case NonFatal(_) => } } + notifyRegisteredConsumers() } else tryClose() } + def notifyRegisteredConsumers(): Unit = { + val message = state.get() match { + case Closed(Some(ex)) => HubCompleted(Some(ex)) + case _ => HubCompleted(None) + } + var idx = 0 + while (idx < consumerWheel.length) { + val bucket = consumerWheel(idx) + if (bucket ne null) { + val itr = bucket.valuesIterator + while (itr.hasNext) { + try itr.next().callback.invoke(message) + catch { case NonFatal(_) => } + } + } + idx += 1 + } + } + tryClose() }