Implementation for composite.

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();
+      }
     }
   }