diff --git a/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala b/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala index e7d15c7541d..70d3e3e06ed 100644 --- a/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala +++ b/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala @@ -2409,6 +2409,62 @@ class ByteStringSpec extends AnyWordSpec with Matchers with Checkers { (ByteString1(Array[Byte](1, 2, 3), 0, Int.MinValue) should be).theSameInstanceAs(ByteString1.empty) } + "ByteStringBuilder.result hands over a mostly used temp buffer without copying" in { + val builder = ByteString.newBuilder + builder.sizeHint(64) + val data = Array.tabulate[Byte](40)(i => i.toByte) + builder.putBytes(data) + val result = builder.result() + result should ===(ByteString(data)) + // the 64 byte buffer is wrapped rather than copied, so the result covers only part of its array + result.isCompact should ===(false) + } + + "ByteStringBuilder.result hands over an exactly filled temp buffer without copying" in { + val builder = ByteString.newBuilder + builder.sizeHint(64) + val data = Array.tabulate[Byte](64)(i => i.toByte) + builder.putBytes(data) + val result = builder.result() + result should ===(ByteString(data)) + result.isCompact should ===(true) + } + + "ByteStringBuilder.result copies when only a small part of the temp buffer is used" in { + val builder = ByteString.newBuilder + builder.sizeHint(64) + val data = Array.tabulate[Byte](10)(i => i.toByte) + builder.putBytes(data) + val result = builder.result() + result should ===(ByteString(data)) + // a small result must not retain the much larger buffer + result.isCompact should ===(true) + result.toArrayUnsafe().length should ===(10) + } + + "ByteStringBuilder.result is not affected by later writes to the builder" in { + val builder = ByteString.newBuilder + builder.sizeHint(64) + val data = Array.fill[Byte](40)(1) + builder.putBytes(data) + val first = builder.result() + builder.clear() + builder.putBytes(Array.fill[Byte](40)(2)) + val second = builder.result() + first should ===(ByteString(data)) + second should ===(ByteString(Array.fill[Byte](40)(2))) + } + + "ByteStringBuilder.result returns the written bytes for any fill level of the temp buffer" in { + for (n <- 0 to 200) { + val builder = ByteString.newBuilder + builder.sizeHint(128) + val data = Array.tabulate[Byte](n)(i => (i % 128).toByte) + builder.putBytes(data) + builder.result() should ===(ByteString(data)) + } + } + "ByteStringBuilder.sizeHint does not shrink existing capacity" in { val builder = ByteString.newBuilder builder.sizeHint(100) diff --git a/actor/src/main/scala/org/apache/pekko/util/ByteString.scala b/actor/src/main/scala/org/apache/pekko/util/ByteString.scala index af10c991a99..73ffdc6d53d 100644 --- a/actor/src/main/scala/org/apache/pekko/util/ByteString.scala +++ b/actor/src/main/scala/org/apache/pekko/util/ByteString.scala @@ -2853,6 +2853,27 @@ final class ByteStringBuilder extends Builder[Byte, ByteString] { } } + /** + * Like `clearTemp` but for the final flush done by `result()`, where the temporary buffer does not + * have to be kept around for further writes. When most of the buffer holds data, it is handed over + * to the resulting `ByteString` instead of being copied; the reference to it is dropped here so that + * a later write to this builder allocates a fresh buffer rather than mutating the handed over bytes. + * When only a small part of the buffer is used it is still copied, so that a small result does not + * retain a much larger array. + */ + private def clearTempForResult(): Unit = { + if (_tempLength > 0) { + if (_tempLength > (_tempCapacity >> 1)) { + _builder += ByteString1(_temp, 0, _tempLength) + _temp = null + _tempCapacity = 0 + } else { + _builder += ByteString1(java.util.Arrays.copyOf(_temp, _tempLength)) + } + _tempLength = 0 + } + } + private def resizeTemp(size: Int): Unit = { val newtemp = if (_temp eq null) new Array[Byte](size) else java.util.Arrays.copyOf(_temp, size) _temp = newtemp @@ -3090,7 +3111,7 @@ final class ByteStringBuilder extends Builder[Byte, ByteString] { def result(): ByteString = if (_length == 0) ByteString.empty else { - clearTemp() + clearTempForResult() val bytestrings = _builder.result() if (bytestrings.size == 1) bytestrings.head