Skip to content

Commit 5a7f350

Browse files
authored
optimize number of buffer allocations (#11879)
Currently this improves 2 flows 1. Known length message which length is greater than 1Mb. Previously the first buffer was 1Mb, and then many buffers of 4096 bytes (from CodedOutputStream), now subsequent buffers are also up to 1Mb 2. In case of compression, the first write is always 10 bytes buffer (gzip header), but worth allocating more space
1 parent 7585b16 commit 5a7f350

File tree

6 files changed

+34
-37
lines changed

6 files changed

+34
-37
lines changed

core/src/main/java/io/grpc/internal/MessageFramer.java

+12-7
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,10 @@ void deliverFrame(
7575
// effectively final. Can only be set once.
7676
private int maxOutboundMessageSize = NO_MAX_OUTBOUND_MESSAGE_SIZE;
7777
private WritableBuffer buffer;
78+
/**
79+
* if > 0 - the number of bytes to allocate for the current known-length message.
80+
*/
81+
private int knownLengthPendingAllocation;
7882
private Compressor compressor = Codec.Identity.NONE;
7983
private boolean messageCompression = true;
8084
private final OutputStreamAdapter outputStreamAdapter = new OutputStreamAdapter();
@@ -222,9 +226,7 @@ private int writeKnownLengthUncompressed(InputStream message, int messageLength)
222226
headerScratch.put(UNCOMPRESSED).putInt(messageLength);
223227
// Allocate the initial buffer chunk based on frame header + payload length.
224228
// Note that the allocator may allocate a buffer larger or smaller than this length
225-
if (buffer == null) {
226-
buffer = bufferAllocator.allocate(headerScratch.position() + messageLength);
227-
}
229+
knownLengthPendingAllocation = HEADER_LENGTH + messageLength;
228230
writeRaw(headerScratch.array(), 0, headerScratch.position());
229231
return writeToOutputStream(message, outputStreamAdapter);
230232
}
@@ -288,8 +290,9 @@ private void writeRaw(byte[] b, int off, int len) {
288290
commitToSink(false, false);
289291
}
290292
if (buffer == null) {
291-
// Request a buffer allocation using the message length as a hint.
292-
buffer = bufferAllocator.allocate(len);
293+
checkState(knownLengthPendingAllocation > 0, "knownLengthPendingAllocation reached 0");
294+
buffer = bufferAllocator.allocate(knownLengthPendingAllocation);
295+
knownLengthPendingAllocation -= min(knownLengthPendingAllocation, buffer.writableBytes());
293296
}
294297
int toWrite = min(len, buffer.writableBytes());
295298
buffer.write(b, off, toWrite);
@@ -388,6 +391,8 @@ public void write(byte[] b, int off, int len) {
388391
* {@link OutputStream}.
389392
*/
390393
private final class BufferChainOutputStream extends OutputStream {
394+
private static final int FIRST_BUFFER_SIZE = 4096;
395+
391396
private final List<WritableBuffer> bufferList = new ArrayList<>();
392397
private WritableBuffer current;
393398

@@ -397,7 +402,7 @@ private final class BufferChainOutputStream extends OutputStream {
397402
* {@link #write(byte[], int, int)}.
398403
*/
399404
@Override
400-
public void write(int b) throws IOException {
405+
public void write(int b) {
401406
if (current != null && current.writableBytes() > 0) {
402407
current.write((byte)b);
403408
return;
@@ -410,7 +415,7 @@ public void write(int b) throws IOException {
410415
public void write(byte[] b, int off, int len) {
411416
if (current == null) {
412417
// Request len bytes initially from the allocator, it may give us more.
413-
current = bufferAllocator.allocate(len);
418+
current = bufferAllocator.allocate(Math.max(FIRST_BUFFER_SIZE, len));
414419
bufferList.add(current);
415420
}
416421
while (len > 0) {

interop-testing/src/test/java/io/grpc/testing/integration/CompressionTest.java

+12-21
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
import static org.junit.Assert.assertNull;
2525
import static org.junit.Assert.assertTrue;
2626

27+
import com.google.common.collect.Iterables;
28+
import com.google.common.collect.Lists;
2729
import com.google.protobuf.ByteString;
2830
import io.grpc.CallOptions;
2931
import io.grpc.Channel;
@@ -53,8 +55,6 @@
5355
import io.grpc.testing.integration.TestServiceGrpc.TestServiceBlockingStub;
5456
import io.grpc.testing.integration.TransportCompressionTest.Fzip;
5557
import java.nio.charset.Charset;
56-
import java.util.ArrayList;
57-
import java.util.Collection;
5858
import java.util.List;
5959
import java.util.concurrent.Executors;
6060
import java.util.concurrent.ScheduledExecutorService;
@@ -146,25 +146,16 @@ public void tearDown() {
146146
* Parameters for test.
147147
*/
148148
@Parameters
149-
public static Collection<Object[]> params() {
150-
boolean[] bools = new boolean[]{false, true};
151-
List<Object[]> combos = new ArrayList<>(64);
152-
for (boolean enableClientMessageCompression : bools) {
153-
for (boolean clientAcceptEncoding : bools) {
154-
for (boolean clientEncoding : bools) {
155-
for (boolean enableServerMessageCompression : bools) {
156-
for (boolean serverAcceptEncoding : bools) {
157-
for (boolean serverEncoding : bools) {
158-
combos.add(new Object[] {
159-
enableClientMessageCompression, clientAcceptEncoding, clientEncoding,
160-
enableServerMessageCompression, serverAcceptEncoding, serverEncoding});
161-
}
162-
}
163-
}
164-
}
165-
}
166-
}
167-
return combos;
149+
public static Iterable<Object[]> params() {
150+
List<Boolean> bools = Lists.newArrayList(false, true);
151+
return Iterables.transform(Lists.cartesianProduct(
152+
bools, // enableClientMessageCompression
153+
bools, // clientAcceptEncoding
154+
bools, // clientEncoding
155+
bools, // enableServerMessageCompression
156+
bools, // serverAcceptEncoding
157+
bools // serverEncoding
158+
), List::toArray);
168159
}
169160

170161
@Test

netty/src/test/java/io/grpc/netty/NettyClientStreamTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ public void writeFrameFutureFailedShouldCancelRpc() {
233233
// Verify that failed SendGrpcFrameCommand results in immediate CancelClientStreamCommand.
234234
inOrder.verify(writeQueue).enqueue(any(CancelClientStreamCommand.class), eq(true));
235235
// Verify that any other failures do not produce another CancelClientStreamCommand in the queue.
236-
inOrder.verify(writeQueue, atLeast(1)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
236+
inOrder.verify(writeQueue, atLeast(0)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
237237
inOrder.verify(writeQueue).enqueue(any(SendGrpcFrameCommand.class), eq(true));
238238
inOrder.verifyNoMoreInteractions();
239239

netty/src/test/java/io/grpc/netty/NettyServerStreamTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ public void writeFrameFutureFailedShouldCancelRpc() {
158158
// Verify that failed SendGrpcFrameCommand results in immediate CancelServerStreamCommand.
159159
inOrder.verify(writeQueue).enqueue(any(CancelServerStreamCommand.class), eq(true));
160160
// Verify that any other failures do not produce another CancelServerStreamCommand in the queue.
161-
inOrder.verify(writeQueue, atLeast(1)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
161+
inOrder.verify(writeQueue, atLeast(0)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
162162
inOrder.verify(writeQueue).enqueue(any(SendGrpcFrameCommand.class), eq(true));
163163
inOrder.verifyNoMoreInteractions();
164164
}

okhttp/src/main/java/io/grpc/okhttp/OkHttpWritableBufferAllocator.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import io.grpc.internal.WritableBuffer;
2020
import io.grpc.internal.WritableBufferAllocator;
2121
import okio.Buffer;
22+
import okio.Segment;
2223

2324
/**
2425
* The default allocator for {@link OkHttpWritableBuffer}s used by the OkHttp transport. OkHttp
@@ -27,9 +28,6 @@
2728
*/
2829
class OkHttpWritableBufferAllocator implements WritableBufferAllocator {
2930

30-
// Use 4k as our minimum buffer size.
31-
private static final int MIN_BUFFER = 4096;
32-
3331
// Set the maximum buffer size to 1MB
3432
private static final int MAX_BUFFER = 1024 * 1024;
3533

@@ -45,7 +43,9 @@ class OkHttpWritableBufferAllocator implements WritableBufferAllocator {
4543
*/
4644
@Override
4745
public WritableBuffer allocate(int capacityHint) {
48-
capacityHint = Math.min(MAX_BUFFER, Math.max(MIN_BUFFER, capacityHint));
46+
// okio buffer uses fixed size Segments, round capacityHint up
47+
capacityHint = Math.min(MAX_BUFFER,
48+
(capacityHint + Segment.SIZE - 1) / Segment.SIZE * Segment.SIZE);
4949
return new OkHttpWritableBuffer(new Buffer(), capacityHint);
5050
}
5151
}

okhttp/src/test/java/io/grpc/okhttp/OkHttpWritableBufferAllocatorTest.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import io.grpc.internal.WritableBuffer;
2222
import io.grpc.internal.WritableBufferAllocator;
2323
import io.grpc.internal.WritableBufferAllocatorTestBase;
24+
import okio.Segment;
2425
import org.junit.Test;
2526
import org.junit.runner.RunWith;
2627
import org.junit.runners.JUnit4;
@@ -42,7 +43,7 @@ protected WritableBufferAllocator allocator() {
4243
public void testCapacity() {
4344
WritableBuffer buffer = allocator().allocate(4096);
4445
assertEquals(0, buffer.readableBytes());
45-
assertEquals(4096, buffer.writableBytes());
46+
assertEquals(Segment.SIZE, buffer.writableBytes());
4647
}
4748

4849
@Test
@@ -54,8 +55,8 @@ public void testInitialCapacityHasMaximum() {
5455

5556
@Test
5657
public void testIsExactBelowMaxCapacity() {
57-
WritableBuffer buffer = allocator().allocate(4097);
58+
WritableBuffer buffer = allocator().allocate(Segment.SIZE + 1);
5859
assertEquals(0, buffer.readableBytes());
59-
assertEquals(4097, buffer.writableBytes());
60+
assertEquals(Segment.SIZE * 2, buffer.writableBytes());
6061
}
6162
}

0 commit comments

Comments
 (0)