netty: improve flushing and object allocations in write queue.

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

@@ -38,7 +38,7 @@
 /**
  * Command sent from a Netty client stream to the handler to cancel the stream.
  */
-class CancelClientStreamCommand {
+class CancelClientStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final NettyClientStream stream;
   private final Status reason;
 

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

@@ -40,7 +40,7 @@
 /**
  * Command sent from a Netty server stream to the handler to cancel the stream.
  */
-class CancelServerStreamCommand {
+class CancelServerStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final NettyServerStream.TransportState stream;
   private final Status reason;
 

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

@@ -39,7 +39,7 @@
  * A command to create a new stream. This is created by {@link NettyClientStream} and passed to the
  * {@link NettyClientHandler} for processing in the Channel thread.
  */
-class CreateStreamCommand {
+class CreateStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final Http2Headers headers;
   private final NettyClientStream stream;
 

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

@@ -37,7 +37,7 @@
  * A command to trigger close and close all streams. It is buffered differently than normal close
  * and also includes reason for closure.
  */
-class ForcefulCloseCommand {
+class ForcefulCloseCommand extends WriteQueue.AbstractQueuedCommand {
   private final Status status;
 
   public ForcefulCloseCommand(Status status) {

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

@@ -37,7 +37,7 @@
  * A command to trigger close. It is buffered differently than normal close and also includes
  * reason for closure.
  */
-class GracefulCloseCommand {
+class GracefulCloseCommand extends WriteQueue.AbstractQueuedCommand {
   private final Status status;
 
   public GracefulCloseCommand(Status status) {

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

@@ -37,7 +37,7 @@
 /**
  * Command which requests messages from the deframer.
  */
-class RequestMessagesCommand {
+class RequestMessagesCommand extends WriteQueue.AbstractQueuedCommand {
 
   private final int numMessages;
   private final Stream stream;

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

@@ -34,14 +34,17 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufHolder;
 import io.netty.buffer.DefaultByteBufHolder;
+import io.netty.channel.ChannelPromise;
 
 /**
  * Command sent from the transport to the Netty channel to send a GRPC frame to the remote endpoint.
  */
-class SendGrpcFrameCommand extends DefaultByteBufHolder {
+class SendGrpcFrameCommand extends DefaultByteBufHolder implements WriteQueue.QueuedCommand {
   private final StreamIdHolder stream;
   private final boolean endStream;
 
+  private ChannelPromise promise;
+
   SendGrpcFrameCommand(StreamIdHolder stream, ByteBuf content, boolean endStream) {
     super(content);
     this.stream = stream;
@@ -116,4 +119,14 @@ public int hashCode() {
     }
     return hash;
   }
+
+  @Override
+  public ChannelPromise promise() {
+    return promise;
+  }
+
+  @Override
+  public void promise(ChannelPromise promise) {
+    this.promise = promise;
+  }
 }

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

@@ -38,7 +38,7 @@
 /**
  * Command sent from the transport to the Netty channel to send a PING frame.
  */
-class SendPingCommand {
+class SendPingCommand extends WriteQueue.AbstractQueuedCommand {
   private final PingCallback callback;
   private final Executor executor;
 

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

@@ -38,7 +38,7 @@
 /**
  * Command sent from the transport to the Netty channel to send response headers to the client.
  */
-class SendResponseHeadersCommand {
+class SendResponseHeadersCommand extends WriteQueue.AbstractQueuedCommand {
   private final StreamIdHolder stream;
   private final Http2Headers headers;
   private final boolean endOfStream;

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

@@ -31,6 +31,7 @@
 
 package io.grpc.netty;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 
 import io.netty.channel.Channel;
@@ -49,6 +50,9 @@ class WriteQueue {
 
   private static final int DEQUE_CHUNK_SIZE = 128;
 
+  @VisibleForTesting
+  static final int MAX_WRITES_BEFORE_FLUSH = 8192;
+
   /**
    * {@link Runnable} used to schedule work onto the tail of the event loop.
    */
@@ -60,18 +64,18 @@ public void run() {
   };
 
   private final Channel channel;
-  private final BlockingQueue<Runnable> queue;
+  private final BlockingQueue<QueuedCommand> queue;
   private final AtomicBoolean scheduled = new AtomicBoolean();
 
   /**
    * ArrayDeque to copy queue into when flushing in event loop.
    */
-  private final ArrayDeque<Runnable> writeChunk =
-      new ArrayDeque<Runnable>(DEQUE_CHUNK_SIZE);
+  private final ArrayDeque<QueuedCommand> writeChunk =
+      new ArrayDeque<QueuedCommand>(DEQUE_CHUNK_SIZE);
 
   public WriteQueue(Channel channel) {
     this.channel = Preconditions.checkNotNull(channel, "channel");
-    queue = new LinkedBlockingQueue<Runnable>();
+    queue = new LinkedBlockingQueue<QueuedCommand>();
   }
 
   /**
@@ -93,7 +97,7 @@ void scheduleFlush() {
    * @param flush true if a flush of the write should be schedule, false if a later call to
    *              enqueue will schedule the flush.
    */
-  ChannelFuture enqueue(Object command, boolean flush) {
+  ChannelFuture enqueue(QueuedCommand command, boolean flush) {
     return enqueue(command, channel.newPromise(), flush);
   }
 
@@ -105,8 +109,12 @@ ChannelFuture enqueue(Object command, boolean flush) {
    * @param flush true if a flush of the write should be schedule, false if a later call to
    *              enqueue will schedule the flush.
    */
-  ChannelFuture enqueue(Object command, ChannelPromise promise, boolean flush) {
-    queue.add(new QueuedCommand(command, promise));
+  ChannelFuture enqueue(QueuedCommand command, ChannelPromise promise, boolean flush) {
+    // Detect errornous code that tries to reuse command objects.
+    Preconditions.checkNotNull(command.promise() == null, "promise must not be set on command");
+
+    command.promise(promise);
+    queue.add(command);
     if (flush) {
       scheduleFlush();
     }
@@ -118,20 +126,34 @@ ChannelFuture enqueue(Object command, ChannelPromise promise, boolean flush) {
    * called in the event loop
    */
   private void flush() {
+    assert channel.eventLoop().inEventLoop();
+
     try {
       boolean flushed = false;
+      // We can't just call flush after having completely drained the queue, as new objects might
+      // be added to the queue continuously and so it may never get empty. If we never flushed in
+      // that case we would be guaranteed to OOM.
+      // However, we also don't want to flush too often as flushing invokes expensive system
+      // calls and we want to feed Netty with enough data so as to maximize the data written per
+      // system call.
+      int writesBeforeFlush = MAX_WRITES_BEFORE_FLUSH;
+      // Always dequeue in chunks, in order to not acquire the queue's lock too often.
       while (queue.drainTo(writeChunk, DEQUE_CHUNK_SIZE) > 0) {
-        while (writeChunk.size() > 0) {
-          writeChunk.poll().run();
+        writesBeforeFlush -= writeChunk.size();
+        while (!writeChunk.isEmpty()) {
+          QueuedCommand cmd = writeChunk.poll();
+          channel.write(cmd, cmd.promise());
+        }
+        flushed = false;
+        if (writesBeforeFlush <= 0) {
+          writesBeforeFlush = MAX_WRITES_BEFORE_FLUSH;
+          flushed = true;
+          channel.flush();
         }
-        // Flush each chunk so we are releasing buffers periodically. In theory this loop
-        // might never end as new events are continuously added to the queue, if we never
-        // flushed in that case we would be guaranteed to OOM.
-        flushed = true;
-        channel.flush();
       }
+
       if (!flushed) {
-        // Must flush at least once
+        // In case there were no items in the queue, we must flush at least once
         channel.flush();
       }
     } finally {
@@ -143,22 +165,33 @@ private void flush() {
     }
   }
 
-  /**
-   * Simple wrapper type around a command and its optional completion listener.
-   */
-  private final class QueuedCommand implements Runnable {
-    private final Object command;
-    private final ChannelPromise promise;
+  abstract static class AbstractQueuedCommand implements QueuedCommand {
 
-    private QueuedCommand(Object command, ChannelPromise promise) {
-      this.command = command;
+    private ChannelPromise promise;
+
+    @Override
+    public final void promise(ChannelPromise promise) {
       this.promise = promise;
     }
 
     @Override
-    public void run() {
-      channel.write(command, promise);
+    public final ChannelPromise promise() {
+      return promise;
     }
   }
-}
 
+  /**
+   * Simple wrapper type around a command and its optional completion listener.
+   */
+  interface QueuedCommand {
+    /**
+     * Returns the promise beeing notified of the success/failure of the write.
+     */
+    ChannelPromise promise();
+
+    /**
+     * Sets the promise.
+     */
+    void promise(ChannelPromise promise);
+  }
+}

netty/src/test/java/io/grpc/netty/NettyClientHandlerTest.java

@@ -373,12 +373,11 @@ public void connectionWindowShouldBeOverridden() throws Exception {
   @Test
   public void createIncrementsIdsForActualAndBufferdStreams() throws Exception {
     receiveMaxConcurrentStreams(2);
-    CreateStreamCommand command = new CreateStreamCommand(grpcHeaders, stream);
-    enqueue(command);
+    enqueue(new CreateStreamCommand(grpcHeaders, stream));
     verify(stream).id(eq(3));
-    enqueue(command);
+    enqueue(new CreateStreamCommand(grpcHeaders, stream));
     verify(stream).id(eq(5));
-    enqueue(command);
+    enqueue(new CreateStreamCommand(grpcHeaders, stream));
     verify(stream).id(eq(7));
   }
 

netty/src/test/java/io/grpc/netty/NettyClientStreamTest.java

@@ -60,6 +60,7 @@
 import io.grpc.Status;
 import io.grpc.internal.ClientStreamListener;
 import io.grpc.internal.GrpcUtil;
+import io.grpc.netty.WriteQueue.QueuedCommand;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
@@ -380,7 +381,7 @@ public void removeUserAgentFromApplicationHeaders() {
     metadata.put(GrpcUtil.USER_AGENT_KEY, "bad agent");
     listener = mock(ClientStreamListener.class);
     Mockito.reset(writeQueue);
-    when(writeQueue.enqueue(any(), any(boolean.class))).thenReturn(future);
+    when(writeQueue.enqueue(any(QueuedCommand.class), any(boolean.class))).thenReturn(future);
 
     stream = new NettyClientStreamImpl(methodDescriptor, new Metadata(), channel, handler,
         DEFAULT_MAX_MESSAGE_SIZE, AsciiString.of("localhost"), AsciiString.of("http"),
@@ -404,8 +405,8 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
         }
         return null;
       }
-    }).when(writeQueue).enqueue(any(), any(ChannelPromise.class), anyBoolean());
-    when(writeQueue.enqueue(any(), anyBoolean())).thenReturn(future);
+    }).when(writeQueue).enqueue(any(QueuedCommand.class), any(ChannelPromise.class), anyBoolean());
+    when(writeQueue.enqueue(any(QueuedCommand.class), anyBoolean())).thenReturn(future);
     NettyClientStream stream = new NettyClientStreamImpl(methodDescriptor, new Metadata(), channel,
         handler, DEFAULT_MAX_MESSAGE_SIZE, AsciiString.of("localhost"), AsciiString.of("http"),
         AsciiString.of("agent"));

netty/src/test/java/io/grpc/netty/NettyHandlerTestBase.java

@@ -194,7 +194,7 @@ protected final Http2Connection connection() {
     return handler().connection();
   }
 
-  protected final ChannelFuture enqueue(Object command) {
+  protected final ChannelFuture enqueue(WriteQueue.QueuedCommand command) {
     ChannelFuture future = writeQueue.enqueue(command, newPromise(), true);
     channel.runPendingTasks();
     return future;

netty/src/test/java/io/grpc/netty/NettyServerStreamTest.java

@@ -54,6 +54,7 @@
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.internal.ServerStreamListener;
+import io.grpc.netty.WriteQueue.QueuedCommand;
 import io.netty.buffer.EmptyByteBuf;
 import io.netty.buffer.UnpooledByteBufAllocator;
 import io.netty.channel.ChannelPromise;
@@ -287,8 +288,8 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
         }
         return null;
       }
-    }).when(writeQueue).enqueue(any(), any(ChannelPromise.class), anyBoolean());
-    when(writeQueue.enqueue(any(), anyBoolean())).thenReturn(future);
+    }).when(writeQueue).enqueue(any(QueuedCommand.class), any(ChannelPromise.class), anyBoolean());
+    when(writeQueue.enqueue(any(QueuedCommand.class), anyBoolean())).thenReturn(future);
     NettyServerStream.TransportState state =
         new NettyServerStream.TransportState(handler, http2Stream, DEFAULT_MAX_MESSAGE_SIZE);
     NettyServerStream stream = new NettyServerStream(channel, state);