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 954c592da9..deaebb4d51 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 @@ -968,6 +968,29 @@ class HubSpec extends StreamSpec { result2.futureValue should ===((1 to 9 by 2).filterNot(_ == 3)) } + "notify consumers when hub materializer is shut down" in { + val hubMat = Materializer(system) + val upstream = TestPublisher.probe[Int]() + val source = Source.fromPublisher(upstream).runWith( + PartitionHub.sink((size, elem) => elem % size, startAfterNrOfConsumers = 1, bufferSize = 8))(hubMat) + + val downstream = TestSubscriber.probe[Int]() + source.runWith(Sink.fromSubscriber(downstream)) + + downstream.request(1) + upstream.expectRequest() + + upstream.sendNext(0) + downstream.expectNext(0) + + hubMat.shutdown() + + // The consumer must be notified (not left hanging), which is the fix. + // The signal is an error because the SubSink output boundary failure + // arrives before the postStop completion callback. + downstream.expectError() + } + } } 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 ca7615a55f..26720b05ec 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 @@ -1374,32 +1374,46 @@ object PartitionHub { 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 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 - consumerInfo.consumers.foreach { consumer => - consumer.callback.invoke(failMessage) - } failStage(ex) } override def postStop(): Unit = { - // Notify pending consumers and set tombstone + // Notify all consumers (pending and registered) when the stage stops. + // Registered consumers 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 here + 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) + } + consumerInfo.consumers.foreach { consumer => + try consumer.callback.invoke(message) + catch { case NonFatal(_) => } + } + } + tryClose() }