ensure end is called on inflaters and deflaters (1.4 branch)

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

@@ -44,12 +44,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)
@@ -60,10 +66,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 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

@@ -14,10 +14,20 @@
 package org.apache.pekko.http.scaladsl.coding
 
 import java.io.{ InputStream, OutputStream }
-import java.util.zip.{ GZIPInputStream, GZIPOutputStream, ZipException }
+import java.nio.charset.StandardCharsets
+import java.util.concurrent.{ CountDownLatch, TimeUnit }
+import java.util.concurrent.atomic.AtomicInteger
+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.stream.SystemMaterializer
+import pekko.stream.scaladsl.{ Sink, Source }
+import pekko.testkit._
 import pekko.util.ByteString
 import scala.annotation.nowarn
 
@@ -45,13 +55,112 @@ class GzipSpec extends CoderSpec {
       ex.ultimateCause.getMessage should equal("Truncated GZIP stream")
     }
     "throw an error if compressed data is just missing the trailer at the end" in {
-      def brokenCompress(payload: String) = Coders.Gzip.newCompressor.compress(ByteString(payload, "UTF-8"))
+      def brokenCompress(payload: String) =
+        Coders.Gzip.newCompressor.compress(ByteString(payload, StandardCharsets.UTF_8))
       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 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 tracking = new TrackingInflater
+
+      Source.single(streamEncode(largeTextBytes))
+        .via(decoderWith(tracking).withMaxBytesPerChunk(1).decoderFlow)
+        .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
+    }
+    "release the inflater when decoding fails on truncation" in {
+      val inflater = new TrackingInflater
+
+      val ex = the[RuntimeException] thrownBy decodeWith(inflater, streamEncode(smallTextBytes).dropRight(5))
+      ex.ultimateCause.getMessage should equal("Truncated GZIP stream")
+      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(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 = tracking
+          }
+    }
+
+  @nowarn("msg=deprecated")
+  private def encoderWith(tracking: TrackingDeflater): Gzip =
+    new Gzip(Encoder.DefaultFilter) {
+      override 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
-  def newCompressor = new DeflateCompressor(compressionLevel)
+  def newCompressor: DeflateCompressor = new DeflateCompressor(compressionLevel)
   def newDecompressorStage(maxBytesPerChunk: Int) = () => new DeflateDecompressor(maxBytesPerChunk)
 
   @deprecated("Use Coders.Deflate(compressionLevel = ...) instead", since = "Akka HTTP 10.2.0")

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

@@ -34,6 +34,7 @@ class DeflateCompressor private[coding] (compressionLevel: Int) extends Compress
   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)
@@ -61,10 +62,18 @@ class DeflateCompressor private[coding] (compressionLevel: Int) extends Compress
   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

@@ -70,7 +70,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())
   }
 }
 
@@ -114,4 +114,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
-  def newCompressor = new GzipCompressor(compressionLevel)
+  def newCompressor: GzipCompressor = new GzipCompressor(compressionLevel)
   def newDecompressorStage(maxBytesPerChunk: Int) = () => new GzipDecompressor(maxBytesPerChunk)
 
   @deprecated("Use Coders.Gzip(compressionLevel = ...) instead", since = "Akka HTTP 10.2.0")

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

@@ -22,8 +22,11 @@ import pekko.stream.impl.io.ByteStringParser
 import pekko.stream.impl.io.ByteStringParser.{ ParseResult, ParseStep }
 import pekko.util.ByteString
 
+import scala.annotation.nowarn
+
 /** Internal API */
 @InternalApi
+@nowarn("msg=deprecated .* is internal API")
 private[coding] class GzipCompressor(compressionLevel: Int) extends DeflateCompressor(compressionLevel) {
   override protected lazy val deflater = new Deflater(compressionLevel, true)
   private val checkSum = new CRC32 // CRC32 of uncompressed data
@@ -70,12 +73,27 @@ private[coding] object GzipCompressor {
 @InternalApi
 private[coding] class GzipDecompressor(
     maxBytesPerChunk: Int = Decoder.MaxBytesPerChunkDefault) extends DeflateDecompressorBase(maxBytesPerChunk) {
+  protected[coding] def createInflater(): Inflater = new Inflater(true)
+
   override def createLogic(attr: Attributes) = new ParsingLogic {
-    private[this] val inflater = new Inflater(true)
+    private[this] val inflater = createInflater()
     private[this] val crc32: CRC32 = new CRC32
+    private[this] var inflaterEnded = false
+
+    private def cleanupInflater(): Unit = {
+      if (!inflaterEnded) {
+        inflaterEnded = true
+        inflater.end()
+      }
+    }
+
+    override def postStop(): Unit = cleanupInflater()
 
     trait Step extends ParseStep[ByteString] {
-      override def onTruncation(): Unit = failStage(new ZipException("Truncated GZIP stream"))
+      override def onTruncation(): Unit = {
+        cleanupInflater()
+        failStage(new ZipException("Truncated GZIP stream"))
+      }
     }
     startWith(ReadHeaders)