`ByteString.readFrom(InputStream)` previously allocated a new scratch `byte[]` inside `readChunk` on every iteration in addition to the copy created by `ByteString.copyFrom(byte[], int, int)`, and continued attempting an extra read (with another buffer allocation) even after a partial chunk had already reached end-of-stream. Similarly, `ByteString.copyFrom(Iterable)` always allocated an `Iterator` even for `RandomAccess` lists, and `ByteString.Output` allocated backing list storage when `toByteString()` was called on a single un-flushed buffer or an empty stream. Reuse the scratch buffer in `ByteString.readFrom` once the chunk size reaches `maxChunkSize` and stop reading once a chunk returns fewer bytes than requested (indicating EOF was reached). Add a `RandomAccess` `List` fast path in `ByteString.copyFrom(Iterable)` and `balancedConcat` to index elements directly without allocating an iterator, and avoid unnecessary buffer and list allocations in `ByteString.Output` for empty or single-buffer outputs. These changes reduce per-operation allocations by 40% to 57% and improve throughput by 1.5x to 1.9x in `readFrom`, while reducing single-buffer `ByteString.Output.toByteString()` allocations by 16% and `copyFrom(List)` iterator overhead. <details> <summary>Expand for JMH benchmark setup and results</summary> Command: ```bash blaze run //javatests/com/google/protobuf:ByteStringComparisonJmhBenchmark_bin_deploy.jar -- \ ".*(readFrom|output|copyFrom).*" -f 1 -wi 3 -i 5 -w 1s -r 1s -prof gc ``` | Benchmark | Before (ns/op) | After (ns/op) | Speedup | Before Alloc (B/op) | After Alloc (B/op) | Alloc Reduction | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | `readFrom100B` | 92.54 | 48.56 | 1.91x (-47.5%) | 1032 | 448 | -584 B/op (-56.6%) | | `readFrom4KB` | 1587.21 | 1038.36 | 1.53x (-34.6%) | 20824 | 12584 | -8240 B/op (-39.6%) | | `readFrom32KB` | 6867.63 | 4500.33 | 1.53x (-34.5%) | 82648 | 49784 | -32864 B/op (-39.8%) | | `outputSingleBuffer` | 33.84 | 29.98 | 1.13x (-11.4%) | 360 | 304 | -56 B/op (-15.6%) | | `copyFrom8ElementList` | 99.94 | 93.67 | 1.07x (-6.3%) | 312 | 280 | -32 B/op (-10.3%) | | `copyFromSingletonList` | 1.00 | 1.04 | ~1.00x | 0 | 0 | 0 B/op | </details> PiperOrigin-RevId: 1000631317
diff --git a/java/core/src/main/java/com/google/protobuf/ByteString.java b/java/core/src/main/java/com/google/protobuf/ByteString.java index 1c44438..1a872ea 100644 --- a/java/core/src/main/java/com/google/protobuf/ByteString.java +++ b/java/core/src/main/java/com/google/protobuf/ByteString.java
@@ -36,6 +36,7 @@ import java.util.List; import java.util.Locale; import java.util.NoSuchElementException; +import java.util.RandomAccess; /** * Immutable sequence of bytes. Provides conversions to and from {@code byte[]}, {@link @@ -663,12 +664,19 @@ // copy the inbound bytes into a list of chunks; the chunk size // grows exponentially to support both short and long streams. int chunkSize = minChunkSize; + byte[] buf = null; while (true) { - ByteString chunk = readChunk(streamToDrain, chunkSize); + if (buf == null || buf.length < chunkSize) { + buf = new byte[chunkSize]; + } + ByteString chunk = readChunk(streamToDrain, buf, chunkSize); if (chunk == null) { break; } results.add(chunk); + if (chunk.size() < chunkSize) { + break; + } chunkSize = Math.min(chunkSize * 2, maxChunkSize); } @@ -683,8 +691,8 @@ * @return A chunk of the desired size, or else a chunk as large as was available when end of * stream was reached. Returns null if the given stream had no more data in it. */ - private static ByteString readChunk(InputStream in, final int chunkSize) throws IOException { - final byte[] buf = new byte[chunkSize]; + private static ByteString readChunk(InputStream in, final byte[] buf, final int chunkSize) + throws IOException { int bytesRead = 0; while (bytesRead < chunkSize) { final int count = in.read(buf, bytesRead, chunkSize - bytesRead); @@ -740,6 +748,17 @@ * Integer.MAX_VALUE */ public static ByteString copyFrom(Iterable<ByteString> byteStrings) { + if (byteStrings instanceof List<ByteString> list && byteStrings instanceof RandomAccess) { + int size = list.size(); + if (size == 0) { + return EMPTY; + } + if (size == 1) { + return list.get(0); + } + return balancedConcat(list, 0, size); + } + // Determine the size; final int size; if (!(byteStrings instanceof Collection)) { @@ -759,6 +778,16 @@ return balancedConcat(byteStrings.iterator(), size); } + private static ByteString balancedConcat(List<ByteString> list, int offset, int length) { + if (length == 1) { + return list.get(offset); + } + int halfLength = length >>> 1; + ByteString left = balancedConcat(list, offset, halfLength); + ByteString right = balancedConcat(list, offset + halfLength, length - halfLength); + return left.concat(right); + } + // Internal function used by copyFrom(Iterable<ByteString>). // Create a balanced concatenation of the next "length" elements from the // iterable. @@ -1190,7 +1219,7 @@ } this.initialCapacity = initialCapacity; this.flushedBuffers = new ArrayList<ByteString>(); - this.buffer = new byte[initialCapacity]; + this.buffer = initialCapacity == 0 ? EMPTY_BYTE_ARRAY : new byte[initialCapacity]; } @Override @@ -1230,6 +1259,14 @@ * @return the current contents of this output stream, as a byte string. */ public synchronized ByteString toByteString() { + if (flushedBuffers.isEmpty()) { + if (bufferPos == 0) { + return EMPTY; + } + if (bufferPos < buffer.length) { + return new LiteralByteString(Arrays.copyOf(buffer, bufferPos)); + } + } flushLastBuffer(); return ByteString.copyFrom(flushedBuffers); } @@ -1256,7 +1293,9 @@ byteString.writeTo(out); } - out.write(Arrays.copyOf(cachedBuffer, cachedBufferPos)); + if (cachedBufferPos > 0) { + out.write(Arrays.copyOf(cachedBuffer, cachedBufferPos)); + } } /** @@ -1292,8 +1331,10 @@ * buffer whose size is at least the specified minimum size. */ private void flushFullBuffer(int minSize) { - flushedBuffers.add(new LiteralByteString(buffer)); - flushedBuffersTotalBytes += buffer.length; + if (buffer.length > 0) { + flushedBuffers.add(new LiteralByteString(buffer)); + flushedBuffersTotalBytes += buffer.length; + } // We want to increase our total capacity by 50%, but as a minimum, // the new buffer should also at least be >= minSize and // >= initial Capacity. @@ -1307,11 +1348,12 @@ * but it needs to be flushed. */ private void flushLastBuffer() { + if (bufferPos == 0) { + return; + } if (bufferPos < buffer.length) { - if (bufferPos > 0) { - byte[] bufferCopy = Arrays.copyOf(buffer, bufferPos); - flushedBuffers.add(new LiteralByteString(bufferCopy)); - } + byte[] bufferCopy = Arrays.copyOf(buffer, bufferPos); + flushedBuffers.add(new LiteralByteString(bufferCopy)); // We reuse this buffer for further writes. } else { // Buffer is completely full. Huzzah.
diff --git a/java/core/src/test/java/com/google/protobuf/ByteStringTest.java b/java/core/src/test/java/com/google/protobuf/ByteStringTest.java index c4e6c6a..ce9fec3 100644 --- a/java/core/src/test/java/com/google/protobuf/ByteStringTest.java +++ b/java/core/src/test/java/com/google/protobuf/ByteStringTest.java
@@ -22,6 +22,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.Comparator; import java.util.Iterator; import java.util.List; @@ -1121,4 +1122,70 @@ .that(byteArrayCopier.getClass().getSimpleName().endsWith("ArraysByteArrayCopier")) .isTrue(); } + + @Test + public void testCopyFrom_randomAccessListSizes() { + assertThat(ByteString.copyFrom(Collections.<ByteString>emptyList())) + .isSameInstanceAs(ByteString.EMPTY); + assertThat(ByteString.copyFrom(new ArrayList<ByteString>())) + .isSameInstanceAs(ByteString.EMPTY); + + ByteString single = ByteString.copyFromUtf8("singleton"); + assertThat(ByteString.copyFrom(Collections.singletonList(single))).isSameInstanceAs(single); + assertThat(ByteString.copyFrom(new ArrayList<>(Collections.singletonList(single)))) + .isSameInstanceAs(single); + + for (int numElements : new int[] {2, 7}) { + List<ByteString> list = new ArrayList<>(numElements); + ByteArrayOutputStream expected = new ByteArrayOutputStream(); + for (int i = 0; i < numElements; i++) { + byte[] chunk = getTestBytes(200 + i * 37, 1000L + i); + list.add(ByteString.copyFrom(chunk)); + expected.write(chunk, 0, chunk.length); + } + ByteString fromList = ByteString.copyFrom(list); + ByteString fromIterable = + ByteString.copyFrom( + new Iterable<ByteString>() { + @Override + public Iterator<ByteString> iterator() { + return list.iterator(); + } + }); + assertThat(fromList.toByteArray()).isEqualTo(expected.toByteArray()); + assertThat(fromList).isEqualTo(fromIterable); + assertThat(fromList.hashCode()).isEqualTo(fromIterable.hashCode()); + } + } + + @Test + public void testReadFrom_multiChunkMaxBufferReuse() throws IOException { + byte[] data = getTestBytes(ByteString.MAX_READ_FROM_CHUNK_SIZE * 4 + 321, 98765L); + ByteString result = ByteString.readFrom(new ByteArrayInputStream(data)); + assertThat(result.toByteArray()).isEqualTo(data); + + byte[] customData = getTestBytes(2000, 54321L); + ByteString customResult = ByteString.readFrom(new ByteArrayInputStream(customData), 64, 256); + assertThat(customResult.toByteArray()).isEqualTo(customData); + } + + @Test + public void testNewOutput_zeroInitialCapacityAndRepeatedToByteString() throws IOException { + ByteString.Output emptyOut = ByteString.newOutput(0); + assertThat(emptyOut.toByteString()).isSameInstanceAs(ByteString.EMPTY); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + emptyOut.writeTo(baos); + assertThat(baos.toByteArray()).isEmpty(); + + byte[] data = getTestBytes(64, 12345L); + ByteString.Output exactOut = ByteString.newOutput(data.length); + exactOut.write(data); + ByteString first = exactOut.toByteString(); + ByteString second = exactOut.toByteString(); + assertThat(second).isSameInstanceAs(first); + assertThat(second.toByteArray()).isEqualTo(data); + + exactOut.writeTo(baos); + assertThat(baos.toByteArray()).isEqualTo(data); + } }