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 013b08e656..cd29412b14 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,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) } } diff --git a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/DeflateSpec.scala b/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/DeflateSpec.scala index 8d86bd40fa..03f4eabcd7 100644 --- a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/DeflateSpec.scala +++ b/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) { diff --git a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/GzipSpec.scala b/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/GzipSpec.scala index 79f9ea604e..6d2f7802de 100644 --- a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/coding/GzipSpec.scala +++ b/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,7 +55,8 @@ 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") } @@ -53,5 +64,103 @@ class GzipSpec extends CoderSpec { 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() + } } } diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Deflate.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Deflate.scala index 8dbb2a4d00..c4c10929c5 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Deflate.scala +++ b/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") diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/DeflateCompressor.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/DeflateCompressor.scala index a6f596744b..2f66b5c404 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/DeflateCompressor.scala +++ b/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. diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Encoder.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Encoder.scala index b81deec687..dcf7685f6c 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Encoder.scala +++ b/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 = () } diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Gzip.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Gzip.scala index 7046d26469..27b74798c1 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/Gzip.scala +++ b/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") diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/GzipCompressor.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/GzipCompressor.scala index a101565562..3f21f9917f 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/coding/GzipCompressor.scala +++ b/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)