From 76eef843217d99b5f1469e964d0bf50a7a31b837 Mon Sep 17 00:00:00 2001
From: "He-Pin(kerr)" <hepin1989@gmail.com>
Date: Sat, 11 Jul 2026 04:20:30 +0800
Subject: [PATCH] Fix byte string transformer cleanup on stage stop (#1154)

Motivation:
A byte string transformer marked itself finished before a pending final emission was delivered, so downstream cancellation in that window skipped cleanup.

Modification:
Treat cleanup as the stage finalizer by invoking it unconditionally from postStop, and add directional tests for normal completion and cancellation while the final emission is pending.

Result:
All stage termination paths release transformer resources without relying on a fragile finished flag. Existing compressor cleanup remains safe because it is idempotent.

Tests:
- sbt "http-core / Test / testOnly org.apache.pekko.http.impl.util.StreamUtilsSpec": 7 passed
- sbt "http-tests / Test / testOnly org.apache.pekko.http.scaladsl.coding.DeflateSpec org.apache.pekko.http.scaladsl.coding.GzipSpec": 54 passed
- sbt checkCodeStyle: passed
- sbt headerCreateAll and +headerCheckAll: passed
- scalafmt --list --mode diff-ref=origin/main: passed
- sbt sortImports: environment failure, Scalafix/scala.meta NoSuchMethodError
- sbt validatePullRequest: not run locally per maintainer request; delegated to CI

References:
Refs #1133; follow-up to #1142
---
 .../pekko/http/impl/util/StreamUtils.scala    |  8 +---
 .../http/impl/util/StreamUtilsSpec.scala      | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

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 cd29412b14..1a27751c0f 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
@@ -44,8 +44,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,
@@ -54,8 +53,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)
@@ -66,12 +63,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 49c3a4f74f..13a18860c3 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 {