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

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

@@ -75,6 +75,10 @@ public class MessageFramer implements Framer {
   // effectively final.  Can only be set once.
   private int maxOutboundMessageSize = NO_MAX_OUTBOUND_MESSAGE_SIZE;
   private WritableBuffer buffer;
+  /**
+   * if > 0 - the number of bytes to allocate for the current known-length message.
+   */
+  private int knownLengthPendingAllocation;
   private Compressor compressor = Codec.Identity.NONE;
   private boolean messageCompression = true;
   private final OutputStreamAdapter outputStreamAdapter = new OutputStreamAdapter();
@@ -222,9 +226,7 @@ public class MessageFramer implements Framer {
     headerScratch.put(UNCOMPRESSED).putInt(messageLength);
     // Allocate the initial buffer chunk based on frame header + payload length.
     // Note that the allocator may allocate a buffer larger or smaller than this length
-    if (buffer == null) {
+    knownLengthPendingAllocation = HEADER_LENGTH + messageLength;
-      buffer = bufferAllocator.allocate(headerScratch.position() + messageLength);
-    }
     writeRaw(headerScratch.array(), 0, headerScratch.position());
     return writeToOutputStream(message, outputStreamAdapter);
   }
@@ -288,8 +290,9 @@ public class MessageFramer implements Framer {
         commitToSink(false, false);
       }
       if (buffer == null) {
-        // Request a buffer allocation using the message length as a hint.
+        checkState(knownLengthPendingAllocation > 0, "knownLengthPendingAllocation reached 0");
-        buffer = bufferAllocator.allocate(len);
+        buffer = bufferAllocator.allocate(knownLengthPendingAllocation);
+        knownLengthPendingAllocation -= min(knownLengthPendingAllocation, buffer.writableBytes());
       }
       int toWrite = min(len, buffer.writableBytes());
       buffer.write(b, off, toWrite);
@@ -388,6 +391,8 @@ public class MessageFramer implements Framer {
    * {@link OutputStream}.
    */
   private final class BufferChainOutputStream extends OutputStream {
+    private static final int FIRST_BUFFER_SIZE = 4096;
+
     private final List<WritableBuffer> bufferList = new ArrayList<>();
     private WritableBuffer current;
 
@@ -397,7 +402,7 @@ public class MessageFramer implements Framer {
      * {@link #write(byte[], int, int)}.
      */
     @Override
-    public void write(int b) throws IOException {
+    public void write(int b) {
       if (current != null && current.writableBytes() > 0) {
         current.write((byte)b);
         return;
@@ -410,7 +415,7 @@ public class MessageFramer implements Framer {
     public void write(byte[] b, int off, int len) {
       if (current == null) {
         // Request len bytes initially from the allocator, it may give us more.
-        current = bufferAllocator.allocate(len);
+        current = bufferAllocator.allocate(Math.max(FIRST_BUFFER_SIZE, len));
         bufferList.add(current);
       }
       while (len > 0) {

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

@@ -24,6 +24,8 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
 import com.google.protobuf.ByteString;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -53,8 +55,6 @@ import io.grpc.testing.integration.Messages.SimpleResponse;
 import io.grpc.testing.integration.TestServiceGrpc.TestServiceBlockingStub;
 import io.grpc.testing.integration.TransportCompressionTest.Fzip;
 import java.nio.charset.Charset;
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -146,25 +146,16 @@ public class CompressionTest {
    * Parameters for test.
    */
   @Parameters
-  public static Collection<Object[]> params() {
+  public static Iterable<Object[]> params() {
-    boolean[] bools = new boolean[]{false, true};
+    List<Boolean> bools = Lists.newArrayList(false, true);
-    List<Object[]> combos = new ArrayList<>(64);
+    return Iterables.transform(Lists.cartesianProduct(
-    for (boolean enableClientMessageCompression : bools) {
+        bools, // enableClientMessageCompression
-      for (boolean clientAcceptEncoding : bools) {
+        bools, // clientAcceptEncoding
-        for (boolean clientEncoding : bools) {
+        bools, // clientEncoding
-          for (boolean enableServerMessageCompression : bools) {
+        bools, // enableServerMessageCompression
-            for (boolean serverAcceptEncoding : bools) {
+        bools, // serverAcceptEncoding
-              for (boolean serverEncoding : bools) {
+        bools  // serverEncoding
-                combos.add(new Object[] {
+    ), List::toArray);
-                    enableClientMessageCompression, clientAcceptEncoding, clientEncoding,
-                    enableServerMessageCompression, serverAcceptEncoding, serverEncoding});
-              }
-            }
-          }
-        }
-      }
-    }
-    return combos;
   }
 
   @Test

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

@@ -233,7 +233,7 @@ public class NettyClientStreamTest extends NettyStreamTestBase<NettyClientStream
     // Verify that failed SendGrpcFrameCommand results in immediate CancelClientStreamCommand.
     inOrder.verify(writeQueue).enqueue(any(CancelClientStreamCommand.class), eq(true));
     // Verify that any other failures do not produce another CancelClientStreamCommand in the queue.
-    inOrder.verify(writeQueue, atLeast(1)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
+    inOrder.verify(writeQueue, atLeast(0)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
     inOrder.verify(writeQueue).enqueue(any(SendGrpcFrameCommand.class), eq(true));
     inOrder.verifyNoMoreInteractions();
 

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

@@ -158,7 +158,7 @@ public class NettyServerStreamTest extends NettyStreamTestBase<NettyServerStream
     // Verify that failed SendGrpcFrameCommand results in immediate CancelServerStreamCommand.
     inOrder.verify(writeQueue).enqueue(any(CancelServerStreamCommand.class), eq(true));
     // Verify that any other failures do not produce another CancelServerStreamCommand in the queue.
-    inOrder.verify(writeQueue, atLeast(1)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
+    inOrder.verify(writeQueue, atLeast(0)).enqueue(any(SendGrpcFrameCommand.class), eq(false));
     inOrder.verify(writeQueue).enqueue(any(SendGrpcFrameCommand.class), eq(true));
     inOrder.verifyNoMoreInteractions();
   }

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

@@ -19,6 +19,7 @@ package io.grpc.okhttp;
 import io.grpc.internal.WritableBuffer;
 import io.grpc.internal.WritableBufferAllocator;
 import okio.Buffer;
+import okio.Segment;
 
 /**
  * The default allocator for {@link OkHttpWritableBuffer}s used by the OkHttp transport. OkHttp
@@ -27,9 +28,6 @@ import okio.Buffer;
  */
 class OkHttpWritableBufferAllocator implements WritableBufferAllocator {
 
-  // Use 4k as our minimum buffer size.
-  private static final int MIN_BUFFER = 4096;
-
   // Set the maximum buffer size to 1MB
   private static final int MAX_BUFFER = 1024 * 1024;
 
@@ -45,7 +43,9 @@ class OkHttpWritableBufferAllocator implements WritableBufferAllocator {
    */
   @Override
   public WritableBuffer allocate(int capacityHint) {
-    capacityHint = Math.min(MAX_BUFFER, Math.max(MIN_BUFFER, capacityHint));
+    // okio buffer uses fixed size Segments, round capacityHint up
+    capacityHint = Math.min(MAX_BUFFER,
+        (capacityHint + Segment.SIZE - 1) / Segment.SIZE * Segment.SIZE);
     return new OkHttpWritableBuffer(new Buffer(), capacityHint);
   }
 }

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

@@ -21,6 +21,7 @@ import static org.junit.Assert.assertEquals;
 import io.grpc.internal.WritableBuffer;
 import io.grpc.internal.WritableBufferAllocator;
 import io.grpc.internal.WritableBufferAllocatorTestBase;
+import okio.Segment;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -42,7 +43,7 @@ public class OkHttpWritableBufferAllocatorTest extends WritableBufferAllocatorTe
   public void testCapacity() {
     WritableBuffer buffer = allocator().allocate(4096);
     assertEquals(0, buffer.readableBytes());
-    assertEquals(4096, buffer.writableBytes());
+    assertEquals(Segment.SIZE, buffer.writableBytes());
   }
 
   @Test
@@ -54,8 +55,8 @@ public class OkHttpWritableBufferAllocatorTest extends WritableBufferAllocatorTe
 
   @Test
   public void testIsExactBelowMaxCapacity() {
-    WritableBuffer buffer = allocator().allocate(4097);
+    WritableBuffer buffer = allocator().allocate(Segment.SIZE + 1);
     assertEquals(0, buffer.readableBytes());
-    assertEquals(4097, buffer.writableBytes());
+    assertEquals(Segment.SIZE * 2, buffer.writableBytes());
   }
 }