netty: create adaptive cumulator

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

@@ -31,6 +31,7 @@
  */
 @Internal
 public abstract class GrpcHttp2ConnectionHandler extends Http2ConnectionHandler {
+  protected static final int CUMULATOR_COMPOSE_MIN_SIZE = 1024;
 
   @Nullable
   protected final ChannelPromise channelUnused;
@@ -42,6 +43,7 @@ protected GrpcHttp2ConnectionHandler(
       Http2Settings initialSettings) {
     super(decoder, encoder, initialSettings);
     this.channelUnused = channelUnused;
+    setCumulator(new NettyAdaptiveCumulator(CUMULATOR_COMPOSE_MIN_SIZE));
   }
 
   /**

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

@@ -0,0 +1,183 @@
+/*
+ * 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.netty;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.CompositeByteBuf;
+
+class NettyAdaptiveCumulator implements io.netty.handler.codec.ByteToMessageDecoder.Cumulator {
+  private final int composeMinSize;
+
+  NettyAdaptiveCumulator(int composeMinSize) {
+    Preconditions.checkArgument(composeMinSize >= 0, "composeMinSize must be non-negative");
+    this.composeMinSize = composeMinSize;
+  }
+
+  /**
+   * "Adaptive" cumulator: cumulate {@link ByteBuf}s by dynamically switching between merge and
+   * compose strategies.
+   *
+   * <p>This cumulator applies a heuristic to make a decision whether to track a reference to the
+   * buffer with bytes received from the network stack in an array ("zero-copy"), or to merge into
+   * the last component (the tail) by performing a memory copy.
+   *
+   * <p>It is necessary as a protection from a potential attack on the {@link
+   * io.netty.handler.codec.ByteToMessageDecoder#COMPOSITE_CUMULATOR}. Consider a pathological case
+   * when an attacker sends TCP packages containing a single byte of data, and forcing the cumulator
+   * to track each one in a separate buffer. The cost is memory overhead for each buffer, and extra
+   * compute to read the cumulation.
+   *
+   * <p>Implemented heuristic establishes a minimal threshold for the total size of the tail and
+   * incoming buffer, below which they are merged. The sum of the tail and the incoming buffer is
+   * used to avoid a case where attacker alternates the size of data packets to trick the cumulator
+   * into always selecting compose strategy.
+   *
+   * <p>Merging strategy attempts to minimize unnecessary memory writes. When possible, it expands
+   * the tail capacity and only copies the incoming buffer into available memory. Otherwise, when
+   * both tail and the buffer must be copied, the tail is reallocated (or fully replaced) with a new
+   * buffer of exponentially increasing capacity (bounded to {@link #composeMinSize}) to ensure
+   * runtime {@code O(n^2)} is amortized to {@code O(n)}.
+   */
+  @Override
+  @SuppressWarnings("ReferenceEquality")
+  public final ByteBuf cumulate(ByteBufAllocator alloc, ByteBuf cumulation, ByteBuf in) {
+    if (!cumulation.isReadable()) {
+      cumulation.release();
+      return in;
+    }
+    CompositeByteBuf composite = null;
+    try {
+      if (cumulation instanceof CompositeByteBuf && cumulation.refCnt() == 1) {
+        composite = (CompositeByteBuf) cumulation;
+        // Writer index must equal capacity if we are going to "write"
+        // new components to the end
+        if (composite.writerIndex() != composite.capacity()) {
+          composite.capacity(composite.writerIndex());
+        }
+      } else {
+        composite = alloc.compositeBuffer(Integer.MAX_VALUE)
+            .addFlattenedComponents(true, cumulation);
+      }
+      addInput(alloc, composite, in);
+      in = null;
+      return composite;
+    } finally {
+      if (in != null) {
+        // We must release if the ownership was not transferred as otherwise it may produce a leak
+        in.release();
+        // Also release any new buffer allocated if we're not returning it
+        if (composite != null && composite != cumulation) {
+          composite.release();
+        }
+      }
+    }
+  }
+
+  @VisibleForTesting
+  void addInput(ByteBufAllocator alloc, CompositeByteBuf composite, ByteBuf in) {
+    if (shouldCompose(composite, in, composeMinSize)) {
+      composite.addFlattenedComponents(true, in);
+    } else {
+      // The total size of the new data and the last component are below the threshold. Merge them.
+      mergeWithCompositeTail(alloc, composite, in);
+    }
+  }
+
+  @VisibleForTesting
+  static boolean shouldCompose(CompositeByteBuf composite, ByteBuf in, int composeMinSize) {
+    int componentCount = composite.numComponents();
+    if (composite.numComponents() == 0) {
+      return true;
+    }
+    int tailSize = composite.capacity() - composite.toByteIndex(componentCount - 1);
+    return tailSize + in.readableBytes() >= composeMinSize;
+  }
+
+  /**
+   * Append the given {@link ByteBuf} {@code in} to {@link CompositeByteBuf} {@code composite} by
+   * expanding or replacing the tail component of the {@link CompositeByteBuf}.
+   *
+   * <p>The goal is to prevent {@code O(n^2)} runtime in a pathological case, that forces copying
+   * the tail component into a new buffer, for each incoming single-byte buffer. We append the new
+   * bytes to the tail, when a write (or a fast write) is possible.
+   *
+   * <p>Otherwise, the tail is replaced with a new buffer, with the capacity increased enough to
+   * achieve runtime amortization.
+   *
+   * <p>We assume that implementations of {@link ByteBufAllocator#calculateNewCapacity(int, int)},
+   * are similar to {@link io.netty.buffer.AbstractByteBufAllocator#calculateNewCapacity(int, int)},
+   * which doubles buffer capacity by normalizing it to the closest power of two. This assumption
+   * is verified in unit tests for this method.
+   */
+  @VisibleForTesting
+  static void mergeWithCompositeTail(ByteBufAllocator alloc, CompositeByteBuf composite,
+      ByteBuf in) {
+
+    int newBytes = in.readableBytes();
+    int tailIndex = composite.numComponents() - 1;
+    int tailStart = composite.toByteIndex(tailIndex);
+    int tailBytes = composite.capacity() - tailStart;
+    int totalBytes = newBytes + tailBytes;
+
+    ByteBuf tail = composite.component(tailIndex);
+    ByteBuf merged = null;
+
+    try {
+      if (tail.refCnt() == 1 && !tail.isReadOnly() && totalBytes <= tail.maxCapacity()) {
+        // Ideal case: the tail isn't shared, and can be expanded to the required capacity.
+        // Take ownership of the tail.
+        merged = tail.retainedDuplicate();
+        merged = merged.unwrap();
+        // The tail is a readable non-composite buffer, so writeBytes() handles everything for us.
+        // - ensureWritable() performs a fast resize when possible (f.e. PooledByteBuf's simply
+        //   updates its boundary to the end of consecutive memory run assigned to this buffer)
+        // - when the required size doesn't fit into maxFastWritableBytes(), a new buffer is
+        //   allocated, and the capacity calculated with alloc.calculateNewCapacity()
+        merged.writeBytes(in);
+      } else {
+        // The tail is shared, or not expandable. Replace it with a new buffer of desired capacity.
+        merged = alloc.buffer(alloc.calculateNewCapacity(totalBytes, Integer.MAX_VALUE));
+        merged.setBytes(0, composite, tailStart, tailBytes)
+            .setBytes(tailBytes, in, in.readerIndex(), newBytes)
+            .writerIndex(totalBytes);
+        in.readerIndex(in.writerIndex());
+      }
+      // Store readerIndex to avoid out of bounds writerIndex during component replacement.
+      int prevReader = composite.readerIndex();
+      // Remove the tail, reset writer index, add merged component.
+      composite.removeComponent(tailIndex).setIndex(0, tailStart)
+          .addFlattenedComponents(true, merged);
+      merged = null;
+      in.release();
+      in = null;
+      // Restore the reader.
+      composite.readerIndex(prevReader);
+    } finally {
+      // Input buffer was merged with the tail.
+      if (in != null) {
+        in.release();
+      }
+      // If merge's ownership isn't transferred to the composite buf, release it to prevent a leak.
+      if (merged != null) {
+        merged.release();
+      }
+    }
+  }
+}