From 2b5c294c4959f689990ddb9748d3524a8fa433bb Mon Sep 17 00:00:00 2001
From: Yuri Schimke <yuri@schimke.ee>
Date: Tue, 3 Aug 2021 10:54:48 +0300
Subject: [PATCH] Rewrite Http2 connectionNotReusedAfterShutdown test to be
 less flaky (#6819)

---
 .../internal/http2/HttpOverHttp2Test.java     | 39 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/okhttp/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java b/okhttp/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
index f6c5d7b91..2d5b90380 100644
--- a/okhttp/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
+++ b/okhttp/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java
@@ -21,6 +21,7 @@ import java.io.InputStream;
 import java.io.InterruptedIOException;
 import java.net.Authenticator;
 import java.net.HttpURLConnection;
+import java.net.Socket;
 import java.net.SocketTimeoutException;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -33,8 +34,10 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
+import javax.net.SocketFactory;
 import javax.net.ssl.SSLException;
 import mockwebserver3.Dispatcher;
 import mockwebserver3.MockResponse;
@@ -49,6 +52,7 @@ import okhttp3.Callback;
 import okhttp3.Connection;
 import okhttp3.Cookie;
 import okhttp3.Credentials;
+import okhttp3.DelegatingSocketFactory;
 import okhttp3.EventListener;
 import okhttp3.Headers;
 import okhttp3.Interceptor;
@@ -446,6 +450,9 @@ public final class HttpOverHttp2Test {
         .setBody(new Buffer().write(new byte[Http2Connection.OKHTTP_CLIENT_WINDOW_SIZE + 1])));
     server.enqueue(new MockResponse()
         .setBody("abc"));
+    // Enqueue an additional response that show if we burnt a good prior response.
+    server.enqueue(new MockResponse()
+      .setBody("XXX"));
 
     Call call1 = client.newCall(new Request.Builder()
         .url(server.url("/"))
@@ -1723,14 +1730,29 @@ public final class HttpOverHttp2Test {
         .setBody("ABC"));
     server.enqueue(new MockResponse()
         .setBody("DEF"));
+    // Enqueue an additional response that show if we burnt a good prior response.
+    server.enqueue(new MockResponse()
+      .setBody("XXX"));
 
-    Call call1 = client.newCall(new Request.Builder()
+    List<RealConnection> connections = new ArrayList<>();
+
+    OkHttpClient localClient = client.newBuilder().eventListener(new EventListener() {
+      @Override
+      public void connectionAcquired(@NotNull Call call, @NotNull Connection connection) {
+        connections.add((RealConnection) connection);
+      }
+    }).build();
+
+    Call call1 = localClient.newCall(new Request.Builder()
         .url(server.url("/"))
         .build());
     Response response1 = call1.execute();
     assertThat(response1.body().string()).isEqualTo("ABC");
 
-    Call call2 = client.newCall(new Request.Builder()
+    // Add delays for DISCONNECT_AT_END to propogate
+    waitForConnectionShutdown(connections.get(0));
+
+    Call call2 = localClient.newCall(new Request.Builder()
         .url(server.url("/"))
         .build());
     Response response2 = call2.execute();
@@ -1739,6 +1761,19 @@ public final class HttpOverHttp2Test {
     assertThat(server.takeRequest().getSequenceNumber()).isEqualTo(0);
   }
 
+  private void waitForConnectionShutdown(RealConnection connection)
+    throws InterruptedException, TimeoutException {
+    if (connection.isHealthy(false)) {
+      Thread.sleep(100L);
+    }
+    if (connection.isHealthy(false)) {
+      Thread.sleep(2000L);
+    }
+    if (connection.isHealthy(false)) {
+      throw new TimeoutException("connection didn't shutdown within timeout");
+    }
+  }
+
   /**
    * This simulates a race condition where we receive a healthy HTTP/2 connection and just prior to
    * writing our request, we get a GOAWAY frame from the server.