protobuf, api, core, netty: zero copy into protobuf

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

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+/**
+ * Extension to an {@link java.io.InputStream} whose content can be accessed as {@link
+ * ByteBuffer}s.
+ *
+ * <p>This can be used for optimizing the case for the consumer of a {@link ByteBuffer}-backed
+ * input stream supports efficient reading from {@link ByteBuffer}s directly. This turns the reader
+ * interface from an {@link java.io.InputStream} to {@link ByteBuffer}s, without copying the
+ * content to a byte array and read from it.
+ */
+// TODO(chengyuanzhang): add ExperimentalApi annotation.
+public interface HasByteBuffer {
+
+  /**
+   * Gets a {@link ByteBuffer} containing up to {@code length} bytes of the content, or {@code
+   * null} if has reached end of the content.
+   * @param length the maximum number of bytes to contain in returned {@link ByteBuffer}.
+   */
+  @Nullable
+  ByteBuffer getByteBuffer(int length);
+}

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

@@ -16,6 +16,8 @@
 
 package io.grpc.internal;
 
+import java.nio.ByteBuffer;
+
 /**
  * Abstract base class for {@link ReadableBuffer} implementations.
  */
@@ -45,6 +47,24 @@ public int arrayOffset() {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public void mark() {}
+
+  @Override
+  public void reset() {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public boolean canUseByteBuffer() {
+    return false;
+  }
+
+  @Override
+  public ByteBuffer getByteBuffer(int length) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public void close() {}
 

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

@@ -20,8 +20,10 @@
 import java.io.OutputStream;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
+import java.nio.InvalidMarkException;
 import java.util.ArrayDeque;
 import java.util.Queue;
+import javax.annotation.Nullable;
 
 /**
  * A {@link ReadableBuffer} that is composed of 0 or more {@link ReadableBuffer}s. This provides a
@@ -34,7 +36,9 @@
 public class CompositeReadableBuffer extends AbstractReadableBuffer {
 
   private int readableBytes;
-  private final Queue<ReadableBuffer> buffers = new ArrayDeque<>();
+  private final Queue<ReadableBuffer> readableBuffers = new ArrayDeque<>();
+  private final Queue<ReadableBuffer> rewindableBuffers = new ArrayDeque<>();
+  private boolean marked;
 
   /**
    * Adds a new {@link ReadableBuffer} at the end of the buffer list. After a buffer is added, it is
@@ -43,16 +47,24 @@ public class CompositeReadableBuffer extends AbstractReadableBuffer {
    * this {@code CompositeBuffer}.
    */
   public void addBuffer(ReadableBuffer buffer) {
+    boolean markHead = marked && readableBuffers.isEmpty();
+    enqueueBuffer(buffer);
+    if (markHead) {
+      readableBuffers.peek().mark();
+    }
+  }
+
+  private void enqueueBuffer(ReadableBuffer buffer) {
     if (!(buffer instanceof CompositeReadableBuffer)) {
-      buffers.add(buffer);
+      readableBuffers.add(buffer);
       readableBytes += buffer.readableBytes();
       return;
     }
 
     CompositeReadableBuffer compositeBuffer = (CompositeReadableBuffer) buffer;
-    while (!compositeBuffer.buffers.isEmpty()) {
-      ReadableBuffer subBuffer = compositeBuffer.buffers.remove();
-      buffers.add(subBuffer);
+    while (!compositeBuffer.readableBuffers.isEmpty()) {
+      ReadableBuffer subBuffer = compositeBuffer.readableBuffers.remove();
+      readableBuffers.add(subBuffer);
     }
     readableBytes += compositeBuffer.readableBytes;
     compositeBuffer.readableBytes = 0;
@@ -138,27 +150,78 @@ public int readInternal(ReadableBuffer buffer, int length) throws IOException {
 
   @Override
   public CompositeReadableBuffer readBytes(int length) {
-    checkReadable(length);
-    readableBytes -= length;
-
-    CompositeReadableBuffer newBuffer = new CompositeReadableBuffer();
-    while (length > 0) {
-      ReadableBuffer buffer = buffers.peek();
-      if (buffer.readableBytes() > length) {
+    final CompositeReadableBuffer newBuffer = new CompositeReadableBuffer();
+    execute(new ReadOperation() {
+      @Override
+      int readInternal(ReadableBuffer buffer, int length) {
         newBuffer.addBuffer(buffer.readBytes(length));
-        length = 0;
-      } else {
-        newBuffer.addBuffer(buffers.poll());
-        length -= buffer.readableBytes();
+        return 0;
       }
-    }
+    }, length);
     return newBuffer;
   }
 
+  @Override
+  public void mark() {
+    while (!rewindableBuffers.isEmpty()) {
+      rewindableBuffers.remove().close();
+    }
+    marked = true;
+    ReadableBuffer buffer = readableBuffers.peek();
+    if (buffer != null) {
+      buffer.mark();
+    }
+  }
+
+  @Override
+  public void reset() {
+    if (!marked) {
+      throw new InvalidMarkException();
+    }
+    ReadableBuffer buffer;
+    if ((buffer = readableBuffers.peek()) != null) {
+      int currentRemain = buffer.readableBytes();
+      buffer.reset();
+      readableBytes += (buffer.readableBytes() - currentRemain);
+    }
+    int size = readableBuffers.size();
+    while ((buffer = rewindableBuffers.poll()) != null) {
+      buffer.reset();
+      readableBuffers.add(buffer);
+      readableBytes += buffer.readableBytes();
+    }
+    for (int i = 0; i < size; i++) {
+      readableBuffers.add(readableBuffers.remove());
+    }
+  }
+
+  @Override
+  public boolean canUseByteBuffer() {
+    for (ReadableBuffer buffer : readableBuffers) {
+      if (!buffer.canUseByteBuffer()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @Nullable
+  @Override
+  public ByteBuffer getByteBuffer(int length) {
+    if (readableBuffers.isEmpty()) {
+      return null;
+    }
+    ReadableBuffer buffer = readableBuffers.peek();
+    return buffer.getByteBuffer(Math.min(length, buffer.readableBytes()));
+  }
+
   @Override
   public void close() {
-    while (!buffers.isEmpty()) {
-      buffers.remove().close();
+    while (!readableBuffers.isEmpty()) {
+      readableBuffers.remove().close();
+    }
+    while (!rewindableBuffers.isEmpty()) {
+      rewindableBuffers.remove().close();
     }
   }
 
@@ -169,12 +232,12 @@ public void close() {
   private void execute(ReadOperation op, int length) {
     checkReadable(length);
 
-    if (!buffers.isEmpty()) {
+    if (!readableBuffers.isEmpty()) {
       advanceBufferIfNecessary();
     }
 
-    for (; length > 0 && !buffers.isEmpty(); advanceBufferIfNecessary()) {
-      ReadableBuffer buffer = buffers.peek();
+    for (; length > 0 && !readableBuffers.isEmpty(); advanceBufferIfNecessary()) {
+      ReadableBuffer buffer = readableBuffers.peek();
       int lengthToCopy = Math.min(length, buffer.readableBytes());
 
       // Perform the read operation for this buffer.
@@ -197,9 +260,17 @@ private void execute(ReadOperation op, int length) {
    * If the current buffer is exhausted, removes and closes it.
    */
   private void advanceBufferIfNecessary() {
-    ReadableBuffer buffer = buffers.peek();
+    ReadableBuffer buffer = readableBuffers.peek();
     if (buffer.readableBytes() == 0) {
-      buffers.remove().close();
+      if (marked) {
+        rewindableBuffers.add(readableBuffers.remove());
+        ReadableBuffer next = readableBuffers.peek();
+        if (next != null) {
+          next.mark();
+        }
+      } else {
+        readableBuffers.remove().close();
+      }
     }
   }
 

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 javax.annotation.Nullable;
 
 /**
  * Base class for a wrapper around another {@link ReadableBuffer}.
@@ -96,6 +97,27 @@ public int arrayOffset() {
     return buf.arrayOffset();
   }
 
+  @Override
+  public void mark() {
+    buf.mark();
+  }
+
+  @Override
+  public void reset() {
+    buf.reset();
+  }
+
+  @Override
+  public boolean canUseByteBuffer() {
+    return buf.canUseByteBuffer();
+  }
+
+  @Nullable
+  @Override
+  public ByteBuffer getByteBuffer(int length) {
+    return buf.getByteBuffer(length);
+  }
+
   @Override
   public void close() {
     buf.close();

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 javax.annotation.Nullable;
 
 /**
  * Interface for an abstract byte buffer. Buffers are intended to be a read-only, except for the
@@ -123,6 +124,39 @@ public interface ReadableBuffer extends Closeable {
    */
   int arrayOffset();
 
+  /**
+   * Marks the current position in this buffer. A subsequent call to the {@link #reset} method
+   * repositions this stream at the last marked position so that subsequent reads re-read the same
+   * bytes.
+   */
+  void mark();
+
+  /**
+   * Repositions this buffer to the position at the time {@link #mark} was last called on this
+   * buffer.
+   */
+  void reset();
+
+  /**
+   * Indicates whether or not {@link #getByteBuffer} operation is supported for this buffer.
+   */
+  boolean canUseByteBuffer();
+
+  /**
+   * Gets a {@link ByteBuffer} that contains up to {@code length} bytes of this buffer's content,
+   * or {@code null} if this buffer has been exhausted. The position of this buffer is unchanged
+   * after calling this method. The returned buffer's content should not be modified, but the
+   * position, limit, and mark may be changed. Operations for changing the position, limit, and
+   * mark of the returned {@link ByteBuffer} does not affect the position, limit, and mark of
+   * this buffer. {@link ByteBuffer}s returned by this method have independent position, limit
+   * and mark. This is an optional method, so callers should first check {@link #canUseByteBuffer}.
+   *
+   * @param length the maximum number of bytes to contain in returned {@link ByteBuffer}.
+   * @throws UnsupportedOperationException the buffer does not support this method.
+   */
+  @Nullable
+  ByteBuffer getByteBuffer(int length);
+
   /**
    * Closes this buffer and releases any resources.
    */