Fix test socket reuse during shutdown

Fixes: 264790610
Test: atest CtsLibcoreOkHttpTestCases:com.squareup.okhttp.ConnectionReuseTest#staleConnectionNotReusedForNonIdempotentRequest
Test: atest CtsLibcoreOkHttpTestCases:com.squareup.okhttp.URLConnectionTest
Change-Id: Ibc766ff1b67f26d0b4d6bfd67bdf3be6af5bfc3e

6 files changed, 98 insertions, 18 deletions

diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java
index db68595..c46d96d 100644
--- a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java
+++ b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java
@@ -21,6 +21,7 @@ import com.squareup.okhttp.internal.framed.Settings;
 import com.squareup.okhttp.ws.WebSocketListener;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.concurrent.TimeUnit;
 import okio.Buffer;
 
@@ -38,6 +39,7 @@ public final class MockResponse implements Cloneable {
   private TimeUnit throttlePeriodUnit = TimeUnit.SECONDS;
 
   private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
+  private Consumer<SocketPolicy> socketShutdownListener = null;
 
   private long bodyDelayAmount = 0;
   private TimeUnit bodyDelayUnit = TimeUnit.MILLISECONDS;
@@ -195,6 +197,28 @@ public final class MockResponse implements Cloneable {
   }
 
   /**
+   * Sets a listener that can wait until a socket is closed. This is important if the
+   * {@link SocketPolicy} gets set to something that shuts down the socket after a transaction and
+   * that socket may get reused in subsequent calls if they happen too fast.
+   *
+   * @param listener The listener that will be notified when a socket is closed. This could be an
+   *                 instance of {@link SocketShutdownListener}.
+   *
+   * @see SocketPolicy
+   * @see SocketShutdownListener
+   */
+  public MockResponse setSocketShutdownListener(Consumer<SocketPolicy> listener) {
+      this.socketShutdownListener = listener;
+      return this;
+  }
+
+  void notifyShutdown(SocketPolicy reason) {
+      if (socketShutdownListener != null) {
+          socketShutdownListener.accept(reason);
+      }
+  }
+
+  /**
    * Throttles the response body writer to sleep for the given period after each
    * series of {@code bytesPerPeriod} bytes are written. Use this to simulate
    * network behavior.
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
index b15e92f..7a3deee 100644
--- a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
+++ b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
@@ -554,11 +554,14 @@ public final class MockWebServer implements TestRule {
         // See warnings associated with these socket policies in SocketPolicy.
         if (response.getSocketPolicy() == SocketPolicy.DISCONNECT_AT_END) {
           socket.close();
+          response.notifyShutdown(SocketPolicy.DISCONNECT_AT_END);
           return false;
         } else if (response.getSocketPolicy() == SocketPolicy.SHUTDOWN_INPUT_AT_END) {
           socket.shutdownInput();
+          response.notifyShutdown(SocketPolicy.SHUTDOWN_INPUT_AT_END);
         } else if (response.getSocketPolicy() == SocketPolicy.SHUTDOWN_OUTPUT_AT_END) {
           socket.shutdownOutput();
+          response.notifyShutdown(SocketPolicy.SHUTDOWN_OUTPUT_AT_END);
         }
 
         sequenceNumber++;
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java
index cdbbf72..6f4602f 100644
--- a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java
+++ b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java
@@ -26,8 +26,9 @@ package com.squareup.okhttp.mockwebserver;
  * server may not have had time to close the socket. The socket will be closed at an indeterminate
  * point before or during the second request. It may be closed after client has started sending the
  * request body. If a request body is not retryable then the client may fail the request, making
- * client behavior non-deterministic. Add delays in the client to improve the chances that the
- * server has closed the socket before follow up requests are made.
+ * client behavior non-deterministic. In order to prevent that {@link
+ * MockResponse#setSocketShutdownListener} can be used. Otherwise, add delays in the client to
+ * improve the chances that the server has closed the socket before follow up requests are made.
  */
 public enum SocketPolicy {
 
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketShutdownListener.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketShutdownListener.java
new file mode 100644
index 0000000..b83ce4a
--- /dev/null
+++ b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketShutdownListener.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2023 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.okhttp.mockwebserver;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.function.Consumer;
+
+/**
+ * Utility class for registering as a listener to {@link MockResponse#setSocketShutdownListener}.
+ *
+ * An instance will wait for a single shutdown callback. Once triggered, the listener will stay in a
+ * "shutdown occurred" state.
+ *
+ * @see SocketPolicy
+ */
+public class SocketShutdownListener implements Consumer<SocketPolicy> {
+
+  final CountDownLatch shutdownLatch = new CountDownLatch(1);
+
+  @Override
+  public void accept(SocketPolicy reason) {
+    shutdownLatch.countDown();
+  }
+
+  public void waitForSocketShutdown() {
+    try {
+      shutdownLatch.await();
+    } catch (InterruptedException e) {
+    }
+  }
+}
diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/ConnectionReuseTest.java b/okhttp-tests/src/test/java/com/squareup/okhttp/ConnectionReuseTest.java
index ec2a4fd..93219bf 100644
--- a/okhttp-tests/src/test/java/com/squareup/okhttp/ConnectionReuseTest.java
+++ b/okhttp-tests/src/test/java/com/squareup/okhttp/ConnectionReuseTest.java
@@ -19,11 +19,11 @@ import com.squareup.okhttp.internal.SslContextBuilder;
 import com.squareup.okhttp.mockwebserver.MockResponse;
 import com.squareup.okhttp.mockwebserver.MockWebServer;
 import com.squareup.okhttp.mockwebserver.SocketPolicy;
+import com.squareup.okhttp.mockwebserver.SocketShutdownListener;
 import com.squareup.okhttp.testing.RecordingHostnameVerifier;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 import javax.net.ssl.SSLContext;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
@@ -166,10 +166,11 @@ public final class ConnectionReuseTest {
     assertEquals(0, server.takeRequest().getSequenceNumber());
   }
 
-  @Ignore("b/264790610")
   @Test public void staleConnectionNotReusedForNonIdempotentRequest() throws Exception {
+    SocketShutdownListener shutdownListener = new SocketShutdownListener();
     server.enqueue(new MockResponse().setBody("a")
-        .setSocketPolicy(SocketPolicy.SHUTDOWN_OUTPUT_AT_END));
+        .setSocketPolicy(SocketPolicy.SHUTDOWN_OUTPUT_AT_END)
+        .setSocketShutdownListener(shutdownListener));
     server.enqueue(new MockResponse().setBody("b"));
 
     Request requestA = new Request.Builder()
@@ -179,6 +180,8 @@ public final class ConnectionReuseTest {
     assertEquals("a", responseA.body().string());
     assertEquals(0, server.takeRequest().getSequenceNumber());
 
+    shutdownListener.waitForSocketShutdown();
+
     Request requestB = new Request.Builder()
         .url(server.url("/"))
         .post(RequestBody.create(MediaType.parse("text/plain"), "b"))
diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/URLConnectionTest.java b/okhttp-tests/src/test/java/com/squareup/okhttp/URLConnectionTest.java
index a3e1450..d20ab9a 100644
--- a/okhttp-tests/src/test/java/com/squareup/okhttp/URLConnectionTest.java
+++ b/okhttp-tests/src/test/java/com/squareup/okhttp/URLConnectionTest.java
@@ -28,6 +28,7 @@ import com.squareup.okhttp.mockwebserver.MockResponse;
 import com.squareup.okhttp.mockwebserver.MockWebServer;
 import com.squareup.okhttp.mockwebserver.RecordedRequest;
 import com.squareup.okhttp.mockwebserver.SocketPolicy;
+import com.squareup.okhttp.mockwebserver.SocketShutdownListener;
 import com.squareup.okhttp.testing.RecordingHostnameVerifier;
 import java.io.IOException;
 import java.io.InputStream;
@@ -409,8 +410,10 @@ public final class URLConnectionTest {
   }
 
   private void testServerClosesOutput(SocketPolicy socketPolicy) throws Exception {
+    SocketShutdownListener shutdownListener = new SocketShutdownListener();
     server.enqueue(new MockResponse().setBody("This connection won't pool properly")
-        .setSocketPolicy(socketPolicy));
+        .setSocketPolicy(socketPolicy)
+        .setSocketShutdownListener(shutdownListener));
     MockResponse responseAfter = new MockResponse().setBody("This comes after a busted connection");
     server.enqueue(responseAfter);
     server.enqueue(responseAfter); // Enqueue 2x because the broken connection may be reused.
@@ -420,9 +423,7 @@ public final class URLConnectionTest {
     assertContent("This connection won't pool properly", connection1);
     assertEquals(0, server.takeRequest().getSequenceNumber());
 
-    // Give the server time to enact the socket policy if it's one that could happen after the
-    // client has received the response.
-    Thread.sleep(500);
+    shutdownListener.waitForSocketShutdown();
 
     HttpURLConnection connection2 = client.open(server.getUrl("/b"));
     connection2.setReadTimeout(100);
@@ -641,10 +642,12 @@ public final class URLConnectionTest {
    * https://github.com/square/okhttp/issues/515
    */
   @Test public void sslFallbackNotUsedWhenRecycledConnectionFails() throws Exception {
+    SocketShutdownListener shutdownListener = new SocketShutdownListener();
     server.useHttps(sslContext.getSocketFactory(), false);
     server.enqueue(new MockResponse()
         .setBody("abc")
-        .setSocketPolicy(SocketPolicy.DISCONNECT_AT_END));
+        .setSocketPolicy(SocketPolicy.DISCONNECT_AT_END)
+        .setSocketShutdownListener(shutdownListener));
     server.enqueue(new MockResponse().setBody("def"));
 
     suppressTlsFallbackScsv(client.client());
@@ -652,8 +655,7 @@ public final class URLConnectionTest {
 
     assertContent("abc", client.open(server.getUrl("/")));
 
-    // Give the server time to disconnect.
-    Thread.sleep(500);
+    shutdownListener.waitForSocketShutdown();
 
     assertContent("def", client.open(server.getUrl("/")));
 
@@ -1273,9 +1275,11 @@ public final class URLConnectionTest {
   }
 
   @Test public void transparentGzipWorksAfterExceptionRecovery() throws Exception {
+    SocketShutdownListener shutdownListener = new SocketShutdownListener();
     server.enqueue(new MockResponse()
         .setBody("a")
-        .setSocketPolicy(SHUTDOWN_INPUT_AT_END));
+        .setSocketPolicy(SHUTDOWN_INPUT_AT_END)
+        .setSocketShutdownListener(shutdownListener));
     server.enqueue(new MockResponse()
         .addHeader("Content-Encoding: gzip")
         .setBody(gzip("b")));
@@ -1283,8 +1287,7 @@ public final class URLConnectionTest {
     // Seed the pool with a bad connection.
     assertContent("a", client.open(server.getUrl("/")));
 
-    // Give the server time to disconnect.
-    Thread.sleep(500);
+    shutdownListener.waitForSocketShutdown();
 
     // This connection will need to be recovered. When it is, transparent gzip should still work!
     assertContent("b", client.open(server.getUrl("/")));
@@ -2674,14 +2677,16 @@ public final class URLConnectionTest {
 
   private void reusedConnectionFailsWithPost(TransferKind transferKind, int requestSize)
       throws Exception {
-    server.enqueue(new MockResponse().setBody("A").setSocketPolicy(DISCONNECT_AT_END));
+    SocketShutdownListener shutdownListener = new SocketShutdownListener();
+    server.enqueue(new MockResponse().setBody("A")
+            .setSocketPolicy(DISCONNECT_AT_END)
+            .setSocketShutdownListener(shutdownListener));
     server.enqueue(new MockResponse().setBody("B"));
     server.enqueue(new MockResponse().setBody("C"));
 
     assertContent("A", client.open(server.getUrl("/a")));
 
-    // Give the server time to disconnect.
-    Thread.sleep(500);
+    shutdownListener.waitForSocketShutdown();
 
     // If the request body is larger than OkHttp's replay buffer, the failure may still occur.
     byte[] requestBody = new byte[requestSize];