Recover from an already-allocated connection that restricts new streams. (square#3525)

This is an edge case that can occur with HTTP/2. Since multiple requests
use the same connection, it's possible for one request to restrict the
connection from creating new streams during a follow-up request.

square#3521

Author: dave-r12

okhttp-tests/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java

@@ -29,7 +29,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.SynchronousQueue;
-
 import javax.net.ssl.HostnameVerifier;
 import okhttp3.Cache;
 import okhttp3.Call;
@@ -46,6 +45,7 @@
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
+import okhttp3.Route;
 import okhttp3.TestUtil;
 import okhttp3.internal.DoubleInetAddressDns;
 import okhttp3.internal.RecordingOkAuthenticator;
@@ -736,6 +736,74 @@ private void noRecoveryFromErrorWithRetryDisabled(ErrorCode errorCode) throws Ex
     }
   }
 
+  @Test public void recoverFromConnectionNoNewStreamsOnFollowUp() throws InterruptedException {
+    server.enqueue(new MockResponse()
+        .setResponseCode(401));
+    server.enqueue(new MockResponse()
+        .setSocketPolicy(SocketPolicy.RESET_STREAM_AT_START)
+        .setHttp2ErrorCode(ErrorCode.CANCEL.httpCode));
+    server.enqueue(new MockResponse()
+        .setBody("DEF"));
+    server.enqueue(new MockResponse()
+        .setResponseCode(301)
+        .addHeader("Location", "/foo"));
+    server.enqueue(new MockResponse()
+        .setBody("ABC"));
+
+    final CountDownLatch latch = new CountDownLatch(1);
+    final BlockingQueue<String> responses = new SynchronousQueue<>();
+    okhttp3.Authenticator authenticator = new okhttp3.Authenticator() {
+      @Override public Request authenticate(Route route, Response response) throws IOException {
+        responses.offer(response.body().string());
+        try {
+          latch.await();
+        } catch (InterruptedException e) {
+          throw new AssertionError();
+        }
+        return response.request();
+      }
+    };
+
+    OkHttpClient blockingAuthClient = client.newBuilder()
+        .authenticator(authenticator)
+        .build();
+
+    Callback callback = new Callback() {
+      @Override public void onFailure(Call call, IOException e) {
+        fail();
+      }
+
+      @Override public void onResponse(Call call, Response response) throws IOException {
+        responses.offer(response.body().string());
+      }
+    };
+
+    // Make the first request waiting until we get our auth challenge.
+    Request request = new Request.Builder()
+        .url(server.url("/"))
+        .build();
+    blockingAuthClient.newCall(request).enqueue(callback);
+    String response1 = responses.take();
+    assertEquals("", response1);
+    assertEquals(0, server.takeRequest().getSequenceNumber());
+
+    // Now make the second request which will restrict the first HTTP/2 connection from creating new
+    // streams.
+    client.newCall(request).enqueue(callback);
+    String response2 = responses.take();
+    assertEquals("DEF", response2);
+    assertEquals(1, server.takeRequest().getSequenceNumber());
+    assertEquals(0, server.takeRequest().getSequenceNumber());
+
+    // Let the first request proceed. It should discard the the held HTTP/2 connection and get a new
+    // one.
+    latch.countDown();
+    String response3 = responses.take();
+    assertEquals("ABC", response3);
+    assertEquals(1, server.takeRequest().getSequenceNumber());
+    assertEquals(2, server.takeRequest().getSequenceNumber());
+  }
+
   @Test public void nonAsciiResponseHeader() throws Exception {
     server.enqueue(new MockResponse()
         .addHeaderLenient("Alpha", "α")

okhttp/src/main/java/okhttp3/internal/connection/StreamAllocation.java

@@ -160,30 +160,44 @@ private RealConnection findConnection(int connectTimeout, int readTimeout, int w
     boolean foundPooledConnection = false;
     RealConnection result = null;
     Route selectedRoute = null;
+    Connection releasedConnection;
+    Socket toClose;
     synchronized (connectionPool) {
       if (released) throw new IllegalStateException("released");
       if (codec != null) throw new IllegalStateException("codec != null");
       if (canceled) throw new IOException("Canceled");
 
-      // Attempt to use an already-allocated connection.
-      RealConnection allocatedConnection = this.connection;
-      if (allocatedConnection != null && !allocatedConnection.noNewStreams) {
-        return allocatedConnection;
+      // Attempt to use an already-allocated connection. We need to be careful here because our
+      // already-allocated connection may have been restricted from creating new streams.
+      releasedConnection = this.connection;
+      toClose = releaseIfNoNewStreams();
+      if (this.connection != null) {
+        // We had an already-allocated connection and it's good.
+        result = this.connection;
+        releasedConnection = null;
       }
 
-      // Attempt to get a connection from the pool.
-      Internal.instance.get(connectionPool, address, this, null);
-      if (connection != null) {
-        foundPooledConnection = true;
-        result = connection;
-      } else {
-        selectedRoute = route;
+      if (result == null) {
+        // Attempt to get a connection from the pool.
+        Internal.instance.get(connectionPool, address, this, null);
+        if (connection != null) {
+          foundPooledConnection = true;
+          result = connection;
+        } else {
+          selectedRoute = route;
+        }
       }
     }
+    closeQuietly(toClose);
 
-    // If we found a pooled connection, we're done.
+    if (releasedConnection != null) {
+      eventListener.connectionReleased(call, releasedConnection);
+    }
     if (foundPooledConnection) {
       eventListener.connectionAcquired(call, result);
+    }
+    if (result != null) {
+      // If we found an already-allocated or pooled connection, we're done.
       return result;
     }
 
@@ -257,6 +271,21 @@ private RealConnection findConnection(int connectTimeout, int readTimeout, int w
     return result;
   }
 
+  /**
+   * Releases the currently held connection and returns a socket to close if the held connection
+   * restricts new streams from being created. With HTTP/2 multiple requests share the same
+   * connection so it's possible that our connection is restricted from creating new streams during
+   * a follow-up request.
+   */
+  private Socket releaseIfNoNewStreams() {
+    assert (Thread.holdsLock(connectionPool));
+    RealConnection allocatedConnection = this.connection;
+    if (allocatedConnection != null && allocatedConnection.noNewStreams) {
+      return deallocate(false, false, true);
+    }
+    return null;
+  }
+
   public void streamFinished(boolean noNewStreams, HttpCodec codec, long bytesRead, IOException e) {
     eventListener.responseBodyEnd(call, bytesRead);