Fix connection leaks on failed web socket upgrades.

Closes: #4258

okhttp-tests/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java

@@ -17,8 +17,10 @@
 
 import java.io.EOFException;
 import java.io.IOException;
+import java.net.HttpURLConnection;
 import java.net.ProtocolException;
 import java.net.SocketTimeoutException;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Random;
 import java.util.concurrent.CountDownLatch;
@@ -27,6 +29,7 @@
 import java.util.logging.Logger;
 import okhttp3.Interceptor;
 import okhttp3.OkHttpClient;
+import okhttp3.Protocol;
 import okhttp3.RecordingEventListener;
 import okhttp3.RecordingHostnameVerifier;
 import okhttp3.Request;
@@ -695,6 +698,35 @@ public final class WebSocketHttpTest {
     clientListener.assertTextMessage("Hello, WebSockets!");
   }
 
+  /**
+   * We had a bug where web socket connections were leaked if the HTTP connection upgrade was not
+   * successful. This test confirms that connections are released back to the connection pool!
+   * https://github.com/square/okhttp/issues/4258
+   */
+  @Test public void webSocketConnectionIsReleased() throws Exception {
+    // This test assumes HTTP/1.1 pooling semantics.
+    client = client.newBuilder()
+        .protocols(Arrays.asList(Protocol.HTTP_1_1))
+        .build();
+
+    webServer.enqueue(new MockResponse()
+        .setResponseCode(HttpURLConnection.HTTP_NOT_FOUND)
+        .setBody("not found!"));
+    webServer.enqueue(new MockResponse());
+
+    newWebSocket();
+    clientListener.assertFailure();
+
+    Request regularRequest = new Request.Builder()
+        .url(webServer.url("/"))
+        .build();
+    Response response = client.newCall(regularRequest).execute();
+    response.close();
+
+    assertEquals(0, webServer.takeRequest().getSequenceNumber());
+    assertEquals(1, webServer.takeRequest().getSequenceNumber());
+  }
+
   private MockResponse upgradeResponse(RecordedRequest request) {
     String key = request.getHeader("Sec-WebSocket-Key");
     return new MockResponse()

okhttp/src/main/java/okhttp3/internal/http/RetryAndFollowUpInterceptor.java

@@ -164,9 +164,7 @@ public StreamAllocation streamAllocation() {
       }
 
       if (followUp == null) {
-        if (!forWebSocket) {
-          streamAllocation.release();
-        }
+        streamAllocation.release();
         return response;
       }