protobuf: zero copy into protobuf

api/src/main/java/io/grpc/ByteBufferReadable.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2020 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc;
+
+import java.nio.ByteBuffer;
+import javax.annotation.Nullable;
+
+/**
+ * An {@link java.io.InputStream} or alike that supports the operation of reading its content into
+ * {@link ByteBuffer}s.
+ *
+ * <p>Usually used for implementations (directly or indirectly) backed by {@link ByteBuffer}s so
+ * that read operations avoid making an extra copy by returning {@link ByteBuffer}s sharing content
+ * with the backing {@link ByteBuffer}s.
+ */
+public interface ByteBufferReadable {
+
+
+  /**
+   * Reads up to a total of {@code length} bytes as {@link ByteBuffer}s.
+   *
+   * @param length the maximum number of bytes to be read.
+   * @return {@link ByteBuffer}s that contains the bytes being read or {@code null} if no more
+   *     bytes can be read.
+   */
+  @Nullable
+  Iterable<ByteBuffer> readByteBuffers(int length);
+}

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

@@ -16,6 +16,9 @@
 
 package io.grpc.internal;
 
+import java.nio.ByteBuffer;
+import java.util.List;
+
 /**
  * Abstract base class for {@link ReadableBuffer} implementations.
  */
@@ -30,6 +33,16 @@ public final int readInt() {
     return (b1 << 24) | (b2 << 16) | (b3 << 8) | b4;
   }
 
+  @Override
+  public boolean shouldUseByteBuffer() {
+    return false;
+  }
+
+  @Override
+  public List<ByteBuffer> readByteBuffers(int length) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public boolean hasArray() {
     return false;

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

@@ -21,6 +21,8 @@
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Queue;
 
 /**
@@ -59,6 +61,35 @@ public void addBuffer(ReadableBuffer buffer) {
     compositeBuffer.close();
   }
 
+  @Override
+  public boolean shouldUseByteBuffer() {
+    for (ReadableBuffer buf : buffers) {
+      if (!buf.shouldUseByteBuffer()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @Override
+  public List<ByteBuffer> readByteBuffers(int length) {
+    checkReadable(length);
+    readableBytes -= length;
+
+    List<ByteBuffer> res = new ArrayList<>();
+    while (length > 0) {
+      ReadableBuffer buffer = buffers.peek();
+      int readLength = length;
+      if (buffer.readableBytes() <= length) {
+        readLength = buffer.readableBytes();
+        buffers.poll();
+      }
+      res.addAll(buffer.readByteBuffers(readLength));
+      length -= readLength;
+    }
+    return res;
+  }
+
   @Override
   public int readableBytes() {
     return readableBytes;

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

@@ -21,6 +21,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.List;
 
 /**
  * Base class for a wrapper around another {@link ReadableBuffer}.
@@ -81,6 +82,16 @@ public ReadableBuffer readBytes(int length) {
     return buf.readBytes(length);
   }
 
+  @Override
+  public boolean shouldUseByteBuffer() {
+    return buf.shouldUseByteBuffer();
+  }
+
+  @Override
+  public List<ByteBuffer> readByteBuffers(int length) {
+    return buf.readByteBuffers(length);
+  }
+
   @Override
   public boolean hasArray() {
     return buf.hasArray();

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

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.List;
 
 /**
  * Interface for an abstract byte buffer. Buffers are intended to be a read-only, except for the
@@ -102,6 +103,22 @@ public interface ReadableBuffer extends Closeable {
    */
   ReadableBuffer readBytes(int length);
 
+  /**
+   * Indicates whether or not this buffer supports {@link #readByteBuffers} operation that returns
+   * buffer's content as {@link ByteBuffer}s without making an extra copy.
+   */
+  boolean shouldUseByteBuffer();
+
+  /**
+   * Reads {@code length} bytes as {@link ByteBuffer}s. This is an optional method, so callers
+   * should first check {@link #shouldUseByteBuffer}.
+   *
+   * @param length the total number of bytes to contain in the returned {@link ByteBuffer}s.
+   * @throws UnsupportedOperationException the buffer does not support this method
+   * @throws IndexOutOfBoundsException if required bytes are not readable
+   */
+  List<ByteBuffer> readByteBuffers(int length);
+
   /**
    * Indicates whether or not this buffer exposes a backing array.
    */

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

@@ -19,13 +19,15 @@
 import static com.google.common.base.Charsets.UTF_8;
 
 import com.google.common.base.Preconditions;
+import io.grpc.ByteBufferReadable;
 import io.grpc.KnownLength;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
+import javax.annotation.Nullable;
 
 /**
  * Utility methods for creating {@link ReadableBuffer} instances.
@@ -103,7 +105,11 @@ public static String readAsStringUtf8(ReadableBuffer buffer) {
    * @param owner if {@code true}, the returned stream will close the buffer when closed.
    */
   public static InputStream openStream(ReadableBuffer buffer, boolean owner) {
-    return new BufferInputStream(owner ? buffer : ignoreClose(buffer));
+    if (!owner) {
+      buffer = ignoreClose(buffer);
+    }
+    return buffer.shouldUseByteBuffer()
+        ? new ByteBufferReadableInputStream(buffer) : new BufferInputStream(buffer);
   }
 
   /**
@@ -297,7 +303,7 @@ public int arrayOffset() {
   /**
    * An {@link InputStream} that is backed by a {@link ReadableBuffer}.
    */
-  private static final class BufferInputStream extends InputStream implements KnownLength {
+  private static class BufferInputStream extends InputStream implements KnownLength {
     final ReadableBuffer buffer;
 
     public BufferInputStream(ReadableBuffer buffer) {
@@ -336,5 +342,25 @@ public void close() throws IOException {
     }
   }
 
+  private static final class ByteBufferReadableInputStream extends BufferInputStream
+      implements ByteBufferReadable {
+
+    ByteBufferReadableInputStream(ReadableBuffer buffer) {
+      super(buffer);
+    }
+
+    @Nullable
+    @Override
+    public Iterable<ByteBuffer> readByteBuffers(int length) {
+      if (buffer.readableBytes() == 0) {
+        // EOF.
+        return null;
+      }
+
+      length = Math.min(buffer.readableBytes(), length);
+      return buffer.readByteBuffers(length);
+    }
+  }
+
   private ReadableBuffers() {}
 }

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

@@ -22,6 +22,8 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
 
 /**
  * A {@link java.nio.Buffer} implementation that is backed by a Netty {@link ByteBuf}. This class
@@ -79,6 +81,21 @@ public NettyReadableBuffer readBytes(int length) {
     return new NettyReadableBuffer(buffer.readRetainedSlice(length));
   }
 
+  @Override
+  public boolean shouldUseByteBuffer() {
+    return buffer.nioBufferCount() > 0;
+  }
+
+  @Override
+  public List<ByteBuffer> readByteBuffers(int length) {
+    if (buffer.readableBytes() < length) {
+      throw new IndexOutOfBoundsException();
+    }
+    List<ByteBuffer> res = Arrays.asList(buffer.nioBuffers(buffer.readerIndex(), length));
+    buffer.skipBytes(length);
+    return res;
+  }
+
   @Override
   public boolean hasArray() {
     return buffer.hasArray();

protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java

@@ -24,6 +24,7 @@
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import com.google.protobuf.Parser;
+import io.grpc.ByteBufferReadable;
 import io.grpc.ExperimentalApi;
 import io.grpc.KnownLength;
 import io.grpc.Metadata;
@@ -173,7 +174,15 @@ public T parse(InputStream stream) {
       try {
         if (stream instanceof KnownLength) {
           int size = stream.available();
-          if (size > 0 && size <= DEFAULT_MAX_MESSAGE_SIZE) {
+          if (size <= 0) {
+            return defaultInstance;
+          }
+          // TODO(chengyuanzhang): we may still want to go with the byte array approach for small
+          //  messages.
+          if (stream instanceof ByteBufferReadable) {
+            cis = CodedInputStream.newInstance(
+                ((ByteBufferReadable) stream).readByteBuffers(size));
+          } else if (size < DEFAULT_MAX_MESSAGE_SIZE) {
             Reference<byte[]> ref;
             // buf should not be used after this method has returned.
             byte[] buf;
@@ -197,8 +206,6 @@ public T parse(InputStream stream) {
               throw new RuntimeException("size inaccurate: " + size + " != " + position);
             }
             cis = CodedInputStream.newInstance(buf, 0, size);
-          } else if (size == 0) {
-            return defaultInstance;
           }
         }
       } catch (IOException e) {