OkHttp: make the pending stream cancellable.

Xudong Ma · Xudong Ma · commit ba103fb87143 · 2015-08-07T11:25:28.000-07:00
diff --git a/okhttp/src/main/java/io/grpc/transport/okhttp/OkHttpClientStream.java b/okhttp/src/main/java/io/grpc/transport/okhttp/OkHttpClientStream.java
@@ -37,6 +37,7 @@
 import com.squareup.okhttp.internal.spdy.ErrorCode;
 import com.squareup.okhttp.internal.spdy.Header;
 
+import io.grpc.Metadata;
 import io.grpc.MethodDescriptor.MethodType;
 import io.grpc.Status;
 import io.grpc.transport.ClientStreamListener;
@@ -90,6 +91,8 @@ static OkHttpClientStream newStream(ClientStreamListener listener,
   private List<Header> requestHeaders;
   @GuardedBy("lock")
   private Queue<PendingData> pendingData = new LinkedList<PendingData>();
+  @GuardedBy("lock")
+  private boolean cancelSent = false;
 
 
   private OkHttpClientStream(ClientStreamListener listener,
@@ -131,7 +134,7 @@ public Integer id() {
   @GuardedBy("lock")
   public void start(Integer id) {
     checkNotNull(id, "id");
-    checkState(this.id == null, "Can only set id once");
+    checkState(this.id == null, "the stream has been started with id %s", this.id);
     this.id = id;
     frameWriter.synStream(false, false, id, 0, requestHeaders);
     requestHeaders = null;
@@ -228,7 +231,25 @@ protected void returnProcessedBytes(int processedBytes) {
 
   @Override
   protected void sendCancel(Status reason) {
-    transport.finishStream(id(), reason, ErrorCode.CANCEL);
+    synchronized (lock) {
+      if (cancelSent) {
+        return;
+      }
+      cancelSent = true;
+      if (pendingData != null) {
+        // stream is pending.
+        transport.removePendingStream(this);
+        // release holding data, so they can be GCed or returned to pool earlier.
+        requestHeaders = null;
+        for (PendingData data : pendingData) {
+          data.buffer.clear();
+        }
+        pendingData = null;
+        transportReportStatus(reason, true, new Metadata.Trailers());
+      } else {
+        transport.finishStream(id(), reason, ErrorCode.CANCEL);
+      }
+    }
   }
 
   @Override
diff --git a/okhttp/src/main/java/io/grpc/transport/okhttp/OkHttpClientTransport.java b/okhttp/src/main/java/io/grpc/transport/okhttp/OkHttpClientTransport.java
@@ -293,6 +293,14 @@ private boolean startPendingStreams() {
     return hasStreamStarted;
   }
 
+  /**
+   * Removes given pending stream, used when a pending stream is cancelled.
+   */
+  @GuardedBy("lock")
+  void removePendingStream(OkHttpClientStream pendingStream) {
+    pendingStreams.remove(pendingStream);
+  }
+
   @Override
   public void start(Listener listener) {
     this.listener = Preconditions.checkNotNull(listener, "listener");
diff --git a/okhttp/src/test/java/io/grpc/transport/okhttp/OkHttpClientTransportTest.java b/okhttp/src/test/java/io/grpc/transport/okhttp/OkHttpClientTransportTest.java
@@ -641,6 +641,22 @@ public void pendingStreamSucceed() throws Exception {
     stream2.cancel(Status.CANCELLED);
   }
 
+  @Test
+  public void pendingStreamCancelled() throws Exception {
+    initTransport();
+    setMaxConcurrentStreams(0);
+    MockStreamListener listener = new MockStreamListener();
+    OkHttpClientStream stream
+        = clientTransport.newStream(method, new Metadata.Headers(), listener);
+    waitForStreamPending(1);
+    stream.sendCancel(Status.CANCELLED);
+    // The second cancel should be an no-op.
+    stream.sendCancel(Status.UNKNOWN);
+    listener.waitUntilStreamClosed();
+    assertEquals(0, clientTransport.getPendingStreamSize());
+    assertEquals(Status.CANCELLED.getCode(), listener.status.getCode());
+  }
+
   @Test
   public void pendingStreamFailedByGoAway() throws Exception {
     initTransport();
