address comments and add tests

buchgr · buchgr · commit fb8454b80d99 · 2016-06-29T15:52:48.000+02:00
diff --git a/netty/src/main/java/io/grpc/netty/CancelClientStreamCommand.java b/netty/src/main/java/io/grpc/netty/CancelClientStreamCommand.java
@@ -34,7 +34,6 @@
 import com.google.common.base.Preconditions;
 
 import io.grpc.Status;
-import io.netty.channel.ChannelPromise;
 
 /**
  * Command sent from a Netty client stream to the handler to cancel the stream.
@@ -43,8 +42,6 @@ class CancelClientStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final NettyClientStream stream;
   private final Status reason;
 
-  private ChannelPromise promise;
-
   CancelClientStreamCommand(NettyClientStream stream, Status reason) {
     this.stream = Preconditions.checkNotNull(stream, "stream");
     Preconditions.checkNotNull(reason);
diff --git a/netty/src/main/java/io/grpc/netty/CancelServerStreamCommand.java b/netty/src/main/java/io/grpc/netty/CancelServerStreamCommand.java
@@ -36,7 +36,6 @@
 import com.google.common.base.Preconditions;
 
 import io.grpc.Status;
-import io.netty.channel.ChannelPromise;
 
 /**
  * Command sent from a Netty server stream to the handler to cancel the stream.
@@ -45,8 +44,6 @@ class CancelServerStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final NettyServerStream.TransportState stream;
   private final Status reason;
 
-  private ChannelPromise promise;
-
   CancelServerStreamCommand(NettyServerStream.TransportState stream, Status reason) {
     this.stream = Preconditions.checkNotNull(stream);
     this.reason = Preconditions.checkNotNull(reason);
diff --git a/netty/src/main/java/io/grpc/netty/CreateStreamCommand.java b/netty/src/main/java/io/grpc/netty/CreateStreamCommand.java
@@ -33,7 +33,6 @@
 
 import com.google.common.base.Preconditions;
 
-import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.http2.Http2Headers;
 
 /**
@@ -44,8 +43,6 @@ class CreateStreamCommand extends WriteQueue.AbstractQueuedCommand {
   private final Http2Headers headers;
   private final NettyClientStream stream;
 
-  private ChannelPromise promise;
-
   CreateStreamCommand(Http2Headers headers,
                       NettyClientStream stream) {
     this.stream = Preconditions.checkNotNull(stream, "stream");
diff --git a/netty/src/main/java/io/grpc/netty/ForcefulCloseCommand.java b/netty/src/main/java/io/grpc/netty/ForcefulCloseCommand.java
@@ -32,7 +32,6 @@
 package io.grpc.netty;
 
 import io.grpc.Status;
-import io.netty.channel.ChannelPromise;
 
 /**
  * A command to trigger close and close all streams. It is buffered differently than normal close
@@ -41,8 +40,6 @@
 class ForcefulCloseCommand extends WriteQueue.AbstractQueuedCommand {
   private final Status status;
 
-  private ChannelPromise promise;
-
   public ForcefulCloseCommand(Status status) {
     this.status = status;
   }
diff --git a/netty/src/main/java/io/grpc/netty/GracefulCloseCommand.java b/netty/src/main/java/io/grpc/netty/GracefulCloseCommand.java
@@ -32,7 +32,6 @@
 package io.grpc.netty;
 
 import io.grpc.Status;
-import io.netty.channel.ChannelPromise;
 
 /**
  * A command to trigger close. It is buffered differently than normal close and also includes
@@ -41,8 +40,6 @@
 class GracefulCloseCommand extends WriteQueue.AbstractQueuedCommand {
   private final Status status;
 
-  private ChannelPromise promise;
-
   public GracefulCloseCommand(Status status) {
     this.status = status;
   }
diff --git a/netty/src/main/java/io/grpc/netty/RequestMessagesCommand.java b/netty/src/main/java/io/grpc/netty/RequestMessagesCommand.java
@@ -33,7 +33,6 @@
 
 import io.grpc.internal.AbstractStream2;
 import io.grpc.internal.Stream;
-import io.netty.channel.ChannelPromise;
 
 /**
  * Command which requests messages from the deframer.
@@ -44,8 +43,6 @@ class RequestMessagesCommand extends WriteQueue.AbstractQueuedCommand {
   private final Stream stream;
   private final AbstractStream2.TransportState state;
 
-  private ChannelPromise promise;
-
   public RequestMessagesCommand(Stream stream, int numMessages) {
     this.state = null;
     this.numMessages = numMessages;
diff --git a/netty/src/main/java/io/grpc/netty/SendGrpcFrameCommand.java b/netty/src/main/java/io/grpc/netty/SendGrpcFrameCommand.java
@@ -120,11 +120,6 @@ public int hashCode() {
     return hash;
   }
 
-  @Override
-  public Object command() {
-    return this;
-  }
-
   @Override
   public ChannelPromise promise() {
     return promise;
diff --git a/netty/src/main/java/io/grpc/netty/SendPingCommand.java b/netty/src/main/java/io/grpc/netty/SendPingCommand.java
@@ -32,7 +32,6 @@
 package io.grpc.netty;
 
 import io.grpc.internal.ClientTransport.PingCallback;
-import io.netty.channel.ChannelPromise;
 
 import java.util.concurrent.Executor;
 
@@ -43,8 +42,6 @@ class SendPingCommand extends WriteQueue.AbstractQueuedCommand {
   private final PingCallback callback;
   private final Executor executor;
 
-  private ChannelPromise promise;
-
   SendPingCommand(PingCallback callback, Executor executor) {
     this.callback = callback;
     this.executor = executor;
diff --git a/netty/src/main/java/io/grpc/netty/SendResponseHeadersCommand.java b/netty/src/main/java/io/grpc/netty/SendResponseHeadersCommand.java
@@ -33,7 +33,6 @@
 
 import com.google.common.base.Preconditions;
 
-import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.http2.Http2Headers;
 
 /**
@@ -44,8 +43,6 @@ class SendResponseHeadersCommand extends WriteQueue.AbstractQueuedCommand {
   private final Http2Headers headers;
   private final boolean endOfStream;
 
-  private ChannelPromise promise;
-
   SendResponseHeadersCommand(StreamIdHolder stream, Http2Headers headers, boolean endOfStream) {
     this.stream = Preconditions.checkNotNull(stream);
     this.headers = Preconditions.checkNotNull(headers);
diff --git a/netty/src/main/java/io/grpc/netty/WriteQueue.java b/netty/src/main/java/io/grpc/netty/WriteQueue.java
@@ -99,7 +99,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);
   }
 
@@ -111,15 +111,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) {
-    final QueuedCommand queuedCommand;
-    if (command instanceof QueuedCommand) {
-      queuedCommand = (QueuedCommand) command;
-      queuedCommand.promise(promise);
-    } else {
-      queuedCommand = new InternalQueuedCommand(command, promise);
-    }
-    queue.add(queuedCommand);
+  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();
     }
@@ -145,16 +142,19 @@ private void flush() {
         writesBeforeFlush -= writeChunk.size();
         while (!writeChunk.isEmpty()) {
           QueuedCommand cmd = writeChunk.poll();
-          channel.write(cmd.command(), cmd.promise());
+          channel.write(cmd, cmd.promise());
         }
         if (writesBeforeFlush <= 0) {
           writesBeforeFlush = min(queue.size(), MAX_WRITES_BEFORE_FLUSH);
           flushed = true;
           channel.flush();
         }
       }
+
+      assert writesBeforeFlush == 0;
+
       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 {
@@ -166,15 +166,10 @@ private void flush() {
     }
   }
 
-  static class AbstractQueuedCommand implements QueuedCommand {
+  abstract static class AbstractQueuedCommand implements QueuedCommand {
 
     private ChannelPromise promise;
 
-    @Override
-    public Object command() {
-      return this;
-    }
-
     @Override
     public final void promise(ChannelPromise promise) {
       this.promise = promise;
@@ -190,11 +185,6 @@ public final ChannelPromise promise() {
    * Simple wrapper type around a command and its optional completion listener.
    */
   interface QueuedCommand {
-    /**
-     * Returns the object to write to the channel.
-     */
-    Object command();
-
     /**
      * Returns the promise beeing notified of the success/failure of the write.
      */
@@ -205,18 +195,4 @@ interface QueuedCommand {
      */
     void promise(ChannelPromise promise);
   }
-
-  private static final class InternalQueuedCommand extends AbstractQueuedCommand {
-    private final Object command;
-
-    private InternalQueuedCommand(Object command, ChannelPromise promise) {
-      this.command = command;
-      promise(promise);
-    }
-
-    @Override
-    public Object command() {
-      return command;
-    }
-  }
 }
diff --git a/netty/src/test/java/io/grpc/netty/NettyClientHandlerTest.java b/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));
   }
 
diff --git a/netty/src/test/java/io/grpc/netty/NettyClientStreamTest.java b/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"));
diff --git a/netty/src/test/java/io/grpc/netty/NettyHandlerTestBase.java b/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;
diff --git a/netty/src/test/java/io/grpc/netty/NettyServerStreamTest.java b/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);
diff --git a/netty/src/test/java/io/grpc/netty/WriteQueueTest.java b/netty/src/test/java/io/grpc/netty/WriteQueueTest.java
