optimize number of buffer allocations

Author: panchenko

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

@@ -75,6 +75,10 @@ void deliverFrame(
   // 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();
@@ -120,6 +124,8 @@ public void setMaxOutboundMessageSize(int maxSize) {
     maxOutboundMessageSize = maxSize;
   }
 
+  private static final int COMPRESSION_MIN_LENGTH = 20;
+
   /**
    * Writes out a payload message.
    *
@@ -137,7 +143,7 @@ public void writePayload(InputStream message) {
     int messageLength = -2;
     try {
       messageLength = getKnownLength(message);
-      if (messageLength != 0 && compressed) {
+      if (messageLength >= COMPRESSION_MIN_LENGTH && compressed) {
         written = writeCompressed(message, messageLength);
       } else {
         written = writeUncompressed(message, messageLength);
@@ -222,13 +228,25 @@ private int writeKnownLengthUncompressed(InputStream message, int messageLength)
     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
+    knownLengthPendingAllocation = HEADER_LENGTH + messageLength;
     if (buffer == null) {
-      buffer = bufferAllocator.allocate(headerScratch.position() + messageLength);
+      buffer = allocateKnownLength();
     }
     writeRaw(headerScratch.array(), 0, headerScratch.position());
     return writeToOutputStream(message, outputStreamAdapter);
   }
 
+  /**
+   * Allocate buffer according to {@link #knownLengthPendingAllocation} which is decremented after
+   * that.
+   */
+  private WritableBuffer allocateKnownLength() {
+    WritableBuffer newBuffer = bufferAllocator.allocateKnownLength(knownLengthPendingAllocation);
+    knownLengthPendingAllocation -= Math.min(knownLengthPendingAllocation,
+        newBuffer.writableBytes());
+    return newBuffer;
+  }
+
   /**
    * Write a message that has been serialized to a sequence of buffers.
    */
@@ -243,7 +261,7 @@ private void writeBufferChain(BufferChainOutputStream bufferChain, boolean compr
     }
     headerScratch.clear();
     headerScratch.put(compressed ? COMPRESSED : UNCOMPRESSED).putInt(messageLength);
-    WritableBuffer writeableHeader = bufferAllocator.allocate(HEADER_LENGTH);
+    WritableBuffer writeableHeader = bufferAllocator.allocateKnownLength(HEADER_LENGTH);
     writeableHeader.write(headerScratch.array(), 0, headerScratch.position());
     if (messageLength == 0) {
       // the payload had 0 length so make the header the current buffer.
@@ -289,7 +307,9 @@ private void writeRaw(byte[] b, int off, int len) {
       }
       if (buffer == null) {
         // Request a buffer allocation using the message length as a hint.
-        buffer = bufferAllocator.allocate(len);
+        buffer = knownLengthPendingAllocation > 0
+            ? allocateKnownLength()
+            : bufferAllocator.allocate(len);
       }
       int toWrite = min(len, buffer.writableBytes());
       buffer.write(b, off, toWrite);
@@ -397,7 +417,7 @@ private final class BufferChainOutputStream extends OutputStream {
      * {@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;

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

@@ -27,4 +27,10 @@ public interface WritableBufferAllocator {
    * free to return a buffer with a greater or lesser capacity.
    */
   WritableBuffer allocate(int capacityHint);
+
+  /**
+   * Request a new {@link WritableBuffer} with the given {@code capacityHint}. This method is
+   * similar to {@link #allocate(int)}, but there is no need to allocate greater capacity.
+   */
+  WritableBuffer allocateKnownLength(int capacityHint);
 }

netty/src/main/java/io/grpc/netty/NettyWritableBufferAllocator.java

@@ -33,6 +33,9 @@
  */
 class NettyWritableBufferAllocator implements WritableBufferAllocator {
 
+  // Use 4k as our minimum buffer size.
+  private static final int MIN_BUFFER = 4 * 1024;
+
   // Set the maximum buffer size to 1MB.
   private static final int MAX_BUFFER = 1024 * 1024;
 
@@ -44,6 +47,12 @@ class NettyWritableBufferAllocator implements WritableBufferAllocator {
 
   @Override
   public WritableBuffer allocate(int capacityHint) {
+    capacityHint = Math.min(MAX_BUFFER, Math.max(MIN_BUFFER, capacityHint));
+    return new NettyWritableBuffer(allocator.buffer(capacityHint, capacityHint));
+  }
+
+  @Override
+  public WritableBuffer allocateKnownLength(int capacityHint) {
     capacityHint = Math.min(MAX_BUFFER, capacityHint);
     return new NettyWritableBuffer(allocator.buffer(capacityHint, capacityHint));
   }

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

@@ -19,6 +19,7 @@
 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
@@ -42,6 +43,12 @@ class OkHttpWritableBufferAllocator implements WritableBufferAllocator {
    */
   @Override
   public WritableBuffer allocate(int capacityHint) {
+    capacityHint = Math.min(MAX_BUFFER, Math.max(Segment.SIZE, capacityHint));
+    return new OkHttpWritableBuffer(new Buffer(), capacityHint);
+  }
+
+  @Override
+  public WritableBuffer allocateKnownLength(int capacityHint) {
     capacityHint = Math.min(MAX_BUFFER, capacityHint);
     return new OkHttpWritableBuffer(new Buffer(), capacityHint);
   }

servlet/src/main/java/io/grpc/servlet/ServletServerStream.java

@@ -38,6 +38,7 @@
 import io.grpc.internal.TransportFrameUtil;
 import io.grpc.internal.TransportTracer;
 import io.grpc.internal.WritableBuffer;
+import io.grpc.internal.WritableBufferAllocator;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
@@ -72,7 +73,7 @@ final class ServletServerStream extends AbstractServerStream {
       Attributes attributes,
       String authority,
       InternalLogId logId) throws IOException {
-    super(ByteArrayWritableBuffer::new, statsTraceCtx);
+    super(ALLOCATOR, statsTraceCtx);
     transportState =
         new ServletTransportState(maxInboundMessageSize, statsTraceCtx, new TransportTracer());
     this.attributes = attributes;
@@ -161,6 +162,20 @@ public void deframeFailed(Throwable cause) {
     }
   }
 
+  private static final WritableBufferAllocator ALLOCATOR = new WritableBufferAllocator() {
+    private static final int MIN_BUFFER = 4096;
+
+    @Override
+    public WritableBuffer allocate(int capacityHint) {
+      return new ByteArrayWritableBuffer(max(MIN_BUFFER, capacityHint));
+    }
+
+    @Override
+    public WritableBuffer allocateKnownLength(int capacityHint) {
+      return new ByteArrayWritableBuffer(capacityHint);
+    }
+  };
+
   private static final class ByteArrayWritableBuffer implements WritableBuffer {
 
     private final int capacity;