perf: avoid copying WebSocket deflate/inflate output

http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/PerMessageDeflate.scala

@@ -17,7 +17,6 @@
 
 package org.apache.pekko.http.impl.engine.ws
 
-import java.io.ByteArrayOutputStream
 import java.util.Random
 import java.util.zip.Deflater
 import java.util.zip.Inflater
@@ -27,6 +26,7 @@ import org.apache.pekko
 import pekko.NotUsed
 import pekko.annotation.InternalApi
 import pekko.http.impl.settings.WebSocketCompressionSettingsImpl
+import pekko.http.impl.util.ByteStringOutputStream
 import pekko.http.scaladsl.model.headers.WebSocketExtension
 import pekko.stream.scaladsl.BidiFlow
 import pekko.stream.scaladsl.Flow
@@ -246,7 +246,7 @@ private[http] object PerMessageDeflate {
       try {
         val input = if (appendTail) data ++ EmptyStoredBlock else data
         inflater.setInput(input.toArrayUnsafe())
-        val output = new ByteArrayOutputStream(1024)
+        val output = new ByteStringOutputStream(1024)
         var count = inflater.inflate(buffer)
         while (count > 0) {
           decompressedMessageBytes += count
@@ -255,7 +255,7 @@ private[http] object PerMessageDeflate {
           output.write(buffer, 0, count)
           count = inflater.inflate(buffer)
         }
-        ByteString.fromArrayUnsafe(output.toByteArray)
+        output.toByteStringUnsafe
       } catch {
         case ex: DataFormatException =>
           throw new ProtocolException(s"Invalid WebSocket compressed message: ${ex.getMessage}")
@@ -345,13 +345,13 @@ private[http] object PerMessageDeflate {
 
     private def deflate(data: ByteString, removeTail: Boolean): ByteString = {
       deflater.setInput(data.toArrayUnsafe())
-      val output = new ByteArrayOutputStream(1024)
+      val output = new ByteStringOutputStream(1024)
       var count = deflater.deflate(buffer, 0, buffer.length, Deflater.SYNC_FLUSH)
       while (count > 0) {
         output.write(buffer, 0, count)
         count = deflater.deflate(buffer, 0, buffer.length, Deflater.SYNC_FLUSH)
       }
-      val bytes = ByteString.fromArrayUnsafe(output.toByteArray)
+      val bytes = output.toByteStringUnsafe
       if (removeTail && bytes.endsWith(EmptyStoredBlock)) bytes.dropRight(EmptyStoredBlock.length) else bytes
     }
 

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

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.http.impl.util
+
+import java.io.ByteArrayOutputStream
+
+import org.apache.pekko
+import pekko.annotation.InternalApi
+import pekko.util.ByteString
+
+/**
+ * INTERNAL API
+ *
+ * An [[java.io.OutputStream]] that buffers into a byte array like [[java.io.ByteArrayOutputStream]] but
+ * that can hand the buffered data over as a [[pekko.util.ByteString]] without copying it, unlike
+ * `ByteArrayOutputStream.toByteArray` which always creates a copy.
+ *
+ * Derived from the `ByteStringOutputStream` in Apache Pekko gRPC
+ * (https://github.com/apache/pekko-grpc/pull/862).
+ */
+@InternalApi
+private[http] final class ByteStringOutputStream(capacity: Int) extends ByteArrayOutputStream(capacity) {
+
+  /**
+   * Wraps the bytes written so far in a `ByteString`. The buffer may be shared with the returned
+   * `ByteString`, so this stream must not be written to, reset or reused afterwards.
+   */
+  def toByteStringUnsafe: ByteString =
+    if (count < 1) ByteString.empty
+    else if (count > (buf.length >> 1))
+      // Most of the buffer is used — wrap it to avoid a copy
+      ByteString.fromArrayUnsafe(buf, 0, count)
+    else
+      // Small amount of data in a large buffer — copy to right-size so the rest can be GC'd
+      ByteString.fromArray(buf, 0, count)
+}

http-core/src/test/scala/org/apache/pekko/http/impl/util/ByteStringOutputStreamSpec.scala

@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.http.impl.util
+
+import org.apache.pekko.util.ByteString
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+
+class ByteStringOutputStreamSpec extends AnyWordSpec with Matchers {
+
+  "ByteStringOutputStream" must {
+
+    "return an empty ByteString when nothing was written" in {
+      new ByteStringOutputStream(16).toByteStringUnsafe should ===(ByteString.empty)
+    }
+
+    "return the bytes written when the buffer is exactly filled" in {
+      val out = new ByteStringOutputStream(4)
+      out.write(Array[Byte](1, 2, 3, 4))
+      out.toByteStringUnsafe should ===(ByteString(1, 2, 3, 4))
+    }
+
+    "return the bytes written when the buffer was grown" in {
+      val out = new ByteStringOutputStream(2)
+      val data = Array.tabulate[Byte](1000)(i => i.toByte)
+      out.write(data)
+      out.toByteStringUnsafe should ===(ByteString(data))
+    }
+
+    "return the bytes written when only a small part of the buffer is used" in {
+      val out = new ByteStringOutputStream(1024)
+      out.write(Array[Byte](1, 2, 3))
+      out.write(4)
+      out.toByteStringUnsafe should ===(ByteString(1, 2, 3, 4))
+    }
+
+    "not retain the buffer when only a small part of it is used" in {
+      val out = new ByteStringOutputStream(1024)
+      out.write(Array[Byte](1, 2, 3))
+      // the ByteString is a copy, so it is not affected by later writes to the stream
+      val result = out.toByteStringUnsafe
+      out.write(Array[Byte](9, 9, 9))
+      result should ===(ByteString(1, 2, 3))
+    }
+
+    "write single bytes and byte ranges" in {
+      val out = new ByteStringOutputStream(8)
+      out.write(1)
+      out.write(Array[Byte](0, 2, 3, 0), 1, 2)
+      out.write(Array[Byte](4, 5, 6, 7, 8))
+      out.toByteStringUnsafe should ===(ByteString(1, 2, 3, 4, 5, 6, 7, 8))
+    }
+  }
+}