diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/util/StreamUtils.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/util/StreamUtils.scala index 5f43eca17..0613a0a9b 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/util/StreamUtils.scala +++ b/http-core/src/main/scala/org/apache/pekko/http/impl/util/StreamUtils.scala @@ -43,8 +43,7 @@ private[http] object StreamUtils { * Creates a transformer that will call `f` for each incoming ByteString and output its result. After the complete * input has been read it will call `finish` once to determine the final ByteString to post to the output. * Empty ByteStrings are discarded. - * If the stage is stopped before the input is fully consumed (e.g. on downstream cancellation or upstream failure), - * `cleanup` is called to release any resources held by the transformer. + * When the stage stops, `cleanup` is called to release any resources held by the transformer. */ def byteStringTransformer( f: ByteString => ByteString, @@ -53,8 +52,6 @@ private[http] object StreamUtils { new SimpleLinearGraphStage[ByteString] { override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with InHandler with OutHandler { - private var finished = false - override def onPush(): Unit = { val data = f(grab(in)) if (data.nonEmpty) push(out, data) @@ -65,12 +62,11 @@ private[http] object StreamUtils { override def onUpstreamFinish(): Unit = { val data = finish() - finished = true if (data.nonEmpty) emit(out, data) completeStage() } - override def postStop(): Unit = if (!finished) cleanup() + override def postStop(): Unit = cleanup() setHandlers(in, out, this) } diff --git a/http-core/src/test/scala/org/apache/pekko/http/impl/util/StreamUtilsSpec.scala b/http-core/src/test/scala/org/apache/pekko/http/impl/util/StreamUtilsSpec.scala index 49c3a4f74..13a18860c 100644 --- a/http-core/src/test/scala/org/apache/pekko/http/impl/util/StreamUtilsSpec.scala +++ b/http-core/src/test/scala/org/apache/pekko/http/impl/util/StreamUtilsSpec.scala @@ -16,6 +16,7 @@ package org.apache.pekko.http.impl.util import org.apache.pekko import pekko.stream.Attributes import pekko.stream.scaladsl.{ Sink, Source } +import pekko.stream.testkit.scaladsl.TestSink import pekko.util.ByteString import pekko.testkit._ import org.scalatest.concurrent.ScalaFutures @@ -25,6 +26,44 @@ import scala.util.Failure class StreamUtilsSpec extends PekkoSpec with ScalaFutures { + "byteStringTransformer" should { + "clean up after normal completion" in { + val events = TestProbe() + val transformed = ByteString("transformed") + val trailer = ByteString("trailer") + val result = + Source + .single(ByteString("input")) + .via(StreamUtils.byteStringTransformer(_ => transformed, () => trailer, () => events.ref ! "cleanup")) + .runWith(Sink.seq) + + Await.result(result, 3.seconds.dilated) shouldBe Seq(transformed, trailer) + events.expectMsg("cleanup") + } + + "clean up if downstream cancels before the final emission" in { + val events = TestProbe() + val transformed = ByteString("transformed") + val trailer = ByteString("trailer") + val downstream = + Source + .single(ByteString("input")) + .via(StreamUtils.byteStringTransformer( + _ => transformed, + () => { + events.ref ! "finish" + trailer + }, + () => events.ref ! "cleanup")) + .runWith(TestSink[ByteString]()) + + downstream.request(1).expectNext(transformed) + events.expectMsg("finish") + downstream.cancel() + events.expectMsg("cleanup") + } + } + "captureTermination" should { "signal completion" when { "upstream terminates" in {