Ensure deflater/inflater end() is called in DeflateCompressor and GzipCompressor

http-core/src/main/scala/org/apache/pekko/http/impl/util/StreamUtils.scala

@@ -43,12 +43,18 @@ 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.
    */
   def byteStringTransformer(
-      f: ByteString => ByteString, finish: () => ByteString): GraphStage[FlowShape[ByteString, ByteString]] =
+      f: ByteString => ByteString,
+      finish: () => ByteString,
+      cleanup: () => Unit = () => ()): GraphStage[FlowShape[ByteString, ByteString]] =
     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)
@@ -59,10 +65,13 @@ 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()
+
           setHandlers(in, out, this)
         }
     }

http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/DeflateSpec.scala

@@ -92,11 +92,41 @@ class DeflateSpec extends CoderSpec {
       decodeWith(inflater, streamEncode(smallTextBytes).dropRight(5))
       inflater.endCalls.get() shouldEqual 1
     }
+    "release the deflater when encoding completes" in {
+      val tracking = new TrackingDeflater
+      Source.single(smallTextBytes)
+        .via(encoderWith(tracking).encoderFlow)
+        .runWith(Sink.ignore)
+        .awaitResult(3.seconds.dilated)
+      tracking.awaitEnd(3.seconds.dilated)
+      tracking.endCalls.get() shouldEqual 1
+    }
+    "release the deflater when encoding is cancelled early" in {
+      val tracking = new TrackingDeflater
+      Source.single(largeTextBytes)
+        .via(encoderWith(tracking).encoderFlow)
+        .take(1)
+        .runWith(Sink.ignore)
+        .awaitResult(3.seconds.dilated)
+      // postStop() (which calls end()) is dispatched to the stage actor after the
+      // Sink.ignore future completes, so we must wait for end() itself rather than
+      // for the stream future to avoid a race.
+      tracking.awaitEnd(3.seconds.dilated)
+      tracking.endCalls.get() shouldEqual 1
+    }
   }
 
   private def decodeWith(inflater: TrackingInflater, bytes: ByteString): ByteString =
     decoderWith(inflater).decode(bytes)(SystemMaterializer(system).materializer).awaitResult(3.seconds.dilated)
 
+  @nowarn("msg=deprecated")
+  private def encoderWith(tracking: TrackingDeflater): Deflate =
+    new Deflate(Encoder.DefaultFilter) {
+      override private[http] def newCompressor: DeflateCompressor = new DeflateCompressor() {
+        override protected lazy val deflater: java.util.zip.Deflater = tracking
+      }
+    }
+
   @nowarn("msg=deprecated")
   private def decoderWith(inflater: TrackingInflater): StreamDecoder =
     new StreamDecoder {
@@ -123,6 +153,20 @@ class DeflateSpec extends CoderSpec {
     }
   }
 
+  private class TrackingDeflater extends java.util.zip.Deflater(Deflater.DEFAULT_COMPRESSION, false) {
+    val endCalls = new AtomicInteger
+    private val endLatch = new CountDownLatch(1)
+
+    def awaitEnd(atMost: FiniteDuration): Unit =
+      endLatch.await(atMost.toMillis, TimeUnit.MILLISECONDS)
+
+    override def end(): Unit = {
+      endCalls.incrementAndGet()
+      endLatch.countDown()
+      super.end()
+    }
+  }
+
   private def encodeMessage(request: HttpRequest, compressionLevel: Int, noWrap: Boolean): HttpRequest = {
     @nowarn("msg=deprecated .* is internal API")
     val deflaterWithoutWrapping = new Deflate(Encoder.DefaultFilter) {

http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/GzipSpec.scala

@@ -15,17 +15,19 @@ package org.apache.pekko.http.scaladsl.coding
 
 import java.io.{ InputStream, OutputStream }
 import java.nio.charset.StandardCharsets
+import java.util.concurrent.{ CountDownLatch, TimeUnit }
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.zip.{ GZIPInputStream, GZIPOutputStream, ZipException }
+import java.util.zip.{ GZIPInputStream, GZIPOutputStream, Inflater, ZipException }
 
+import scala.annotation.nowarn
+import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.duration._
 
 import org.apache.pekko
 import pekko.http.impl.util._
-import pekko.http.scaladsl.model.headers.HttpEncodings
 import pekko.stream.SystemMaterializer
 import pekko.stream.scaladsl.{ Sink, Source }
-import pekko.testkit.TestDuration
+import pekko.testkit._
 import pekko.util.ByteString
 
 class GzipSpec extends CoderSpec {
@@ -56,23 +58,29 @@ class GzipSpec extends CoderSpec {
       val ex = the[RuntimeException] thrownBy ourDecode(brokenCompress("abcdefghijkl"))
       ex.ultimateCause.getMessage should equal("Truncated GZIP stream")
     }
+    "throw early if header is corrupt" in {
+      val cause = (the[RuntimeException] thrownBy ourDecode(ByteString(0, 1, 2, 3, 4))).ultimateCause
+      cause should ((be(a[ZipException]) and have).message("Not in GZIP format"))
+    }
     "release the inflater when decoding completes" in {
-      val inflater = new TrackingInflater
-
-      decodeWith(inflater, streamEncode(smallTextBytes)) should readAs(smallText)
-      inflater.endCalls.get() shouldEqual 1
+      val tracking = new TrackingInflater
+      decodeWith(tracking, streamEncode(smallTextBytes)) should readAs(smallText)
+      tracking.endCalls.get() shouldEqual 1
     }
     "release the inflater when decoding is cancelled early" in {
-      val inflater = new TrackingInflater
-      val compressed = streamEncode(largeTextBytes)
+      val tracking = new TrackingInflater
 
-      Source.single(compressed)
-        .via(decoderWith(inflater).withMaxBytesPerChunk(1).decoderFlow)
+      Source.single(streamEncode(largeTextBytes))
+        .via(decoderWith(tracking).withMaxBytesPerChunk(1).decoderFlow)
         .take(1)
         .runWith(Sink.ignore)
         .awaitResult(3.seconds.dilated)
 
-      inflater.endCalls.get() shouldEqual 1
+      // postStop() (which calls end()) is dispatched to the stage actor after the
+      // Sink.ignore future completes, so we must wait for end() itself rather than
+      // for the stream future to avoid a race.
+      tracking.awaitEnd(3.seconds.dilated)
+      tracking.endCalls.get() shouldEqual 1
     }
     "release the inflater when decoding fails on truncation" in {
       val inflater = new TrackingInflater
@@ -81,31 +89,75 @@ class GzipSpec extends CoderSpec {
       ex.ultimateCause.getMessage should equal("Truncated GZIP stream")
       inflater.endCalls.get() shouldEqual 1
     }
-    "throw early if header is corrupt" in {
-      val cause = (the[RuntimeException] thrownBy ourDecode(ByteString(0, 1, 2, 3, 4))).ultimateCause
-      cause should ((be(a[ZipException]) and have).message("Not in GZIP format"))
+    "release the deflater when encoding completes" in {
+      val tracking = new TrackingDeflater
+      Source.single(smallTextBytes)
+        .via(encoderWith(tracking).encoderFlow)
+        .runWith(Sink.ignore)
+        .awaitResult(3.seconds.dilated)
+      tracking.awaitEnd(3.seconds.dilated)
+      tracking.endCalls.get() shouldEqual 1
+    }
+    "release the deflater when encoding is cancelled early" in {
+      val tracking = new TrackingDeflater
+      Source.single(largeTextBytes)
+        .via(encoderWith(tracking).encoderFlow)
+        .take(1)
+        .runWith(Sink.ignore)
+        .awaitResult(3.seconds.dilated)
+      // postStop() (which calls end()) is dispatched to the stage actor after the
+      // Sink.ignore future completes, so we must wait for end() itself rather than
+      // for the stream future to avoid a race.
+      tracking.awaitEnd(3.seconds.dilated)
+      tracking.endCalls.get() shouldEqual 1
     }
   }
 
-  private def decodeWith(inflater: TrackingInflater, bytes: ByteString): ByteString =
-    decoderWith(inflater).decode(bytes)(SystemMaterializer(system).materializer).awaitResult(3.seconds.dilated)
-
-  private def decoderWith(inflater: TrackingInflater): StreamDecoder =
-    new StreamDecoder {
-      override val encoding = HttpEncodings.gzip
+  private def decodeWith(tracking: TrackingInflater, bytes: ByteString): ByteString =
+    decoderWith(tracking).decode(bytes)(SystemMaterializer(system).materializer).awaitResult(3.seconds.dilated)
 
+  @nowarn("msg=deprecated")
+  private def decoderWith(tracking: TrackingInflater): Gzip =
+    new Gzip(Encoder.DefaultFilter) {
       override def newDecompressorStage(maxBytesPerChunk: Int) =
         () =>
           new GzipDecompressor(maxBytesPerChunk) {
-            override protected[coding] def createInflater() = inflater
+            override protected[coding] def createInflater(): Inflater = tracking
           }
     }
 
+  @nowarn("msg=deprecated")
+  private def encoderWith(tracking: TrackingDeflater): Gzip =
+    new Gzip(Encoder.DefaultFilter) {
+      override private[http] def newCompressor: GzipCompressor = new GzipCompressor() {
+        override protected lazy val deflater: java.util.zip.Deflater = tracking
+      }
+    }
+
   private class TrackingInflater extends java.util.zip.Inflater(true) {
     val endCalls = new AtomicInteger
+    private val endLatch = new CountDownLatch(1)
+
+    def awaitEnd(atMost: FiniteDuration): Unit =
+      endLatch.await(atMost.toMillis, TimeUnit.MILLISECONDS)
+
+    override def end(): Unit = {
+      endCalls.incrementAndGet()
+      endLatch.countDown()
+      super.end()
+    }
+  }
+
+  private class TrackingDeflater extends java.util.zip.Deflater(java.util.zip.Deflater.DEFAULT_COMPRESSION, true) {
+    val endCalls = new AtomicInteger
+    private val endLatch = new CountDownLatch(1)
+
+    def awaitEnd(atMost: FiniteDuration): Unit =
+      endLatch.await(atMost.toMillis, TimeUnit.MILLISECONDS)
 
     override def end(): Unit = {
       endCalls.incrementAndGet()
+      endLatch.countDown()
       super.end()
     }
   }

http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Deflate.scala

@@ -27,7 +27,7 @@ class Deflate private[http] (compressionLevel: Int, val messageFilter: HttpMessa
   }
 
   val encoding = HttpEncodings.deflate
-  private[http] def newCompressor = new DeflateCompressor(compressionLevel)
+  private[http] def newCompressor: DeflateCompressor = new DeflateCompressor(compressionLevel)
   def newDecompressorStage(maxBytesPerChunk: Int) = () => new DeflateDecompressor(maxBytesPerChunk)
 
   @InternalApi // used by javadsl.coding.Coder

http/src/main/scala/org/apache/pekko/http/scaladsl/coding/DeflateCompressor.scala

@@ -33,6 +33,7 @@ private[coding] class DeflateCompressor private[coding] (compressionLevel: Int)
   def this() = this(DeflateCompressor.DefaultCompressionLevel)
 
   protected lazy val deflater = new Deflater(compressionLevel, false)
+  private var deflaterEnded = false
 
   override final def compressAndFlush(input: ByteString): ByteString = {
     val buffer = newTempBuffer(input.size)
@@ -60,10 +61,18 @@ private[coding] class DeflateCompressor private[coding] (compressionLevel: Int)
   protected def finishWithBuffer(buffer: Array[Byte]): ByteString = {
     deflater.finish()
     val res = drainDeflater(deflater, buffer)
-    deflater.end()
+    endDeflater()
     res
   }
 
+  private[coding] def endDeflater(): Unit =
+    if (!deflaterEnded) {
+      deflaterEnded = true
+      deflater.end()
+    }
+
+  private[coding] override def cleanup(): Unit = endDeflater()
+
   private def newTempBuffer(size: Int = 65536): Array[Byte] = {
     // The default size is somewhat arbitrary, we'd like to guess a better value but Deflater/zlib
     // is buffering in an unpredictable manner.

http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Encoder.scala

@@ -56,7 +56,7 @@ trait Encoder {
     def encodeChunk(bytes: ByteString): ByteString = compressor.compressAndFlush(bytes)
     def finish(): ByteString = compressor.finish()
 
-    StreamUtils.byteStringTransformer(encodeChunk, () => finish())
+    StreamUtils.byteStringTransformer(encodeChunk, () => finish(), () => compressor.cleanup())
   }
 }
 
@@ -100,4 +100,7 @@ abstract class Compressor {
 
   /** Combines `compress` + `finish` */
   def compressAndFinish(input: ByteString): ByteString
+
+  /** Release any native resources held by this compressor. Idempotent. */
+  private[coding] def cleanup(): Unit = ()
 }

http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Gzip.scala

@@ -27,7 +27,7 @@ class Gzip private[http] (compressionLevel: Int, val messageFilter: HttpMessage
   }
 
   val encoding = HttpEncodings.gzip
-  private[http] def newCompressor = new GzipCompressor(compressionLevel)
+  private[http] def newCompressor: GzipCompressor = new GzipCompressor(compressionLevel)
   def newDecompressorStage(maxBytesPerChunk: Int) = () => new GzipDecompressor(maxBytesPerChunk)
 
   @InternalApi // used by javadsl.coding.Coder