Also poll logs before closing Jetty in case of shutdown racing with log handlers

1 files changed, 28 insertions, 20 deletions

diff --git a/container-core/src/test/java/com/yahoo/jdisc/http/server/jetty/ProxyProtocolTest.java b/container-core/src/test/java/com/yahoo/jdisc/http/server/jetty/ProxyProtocolTest.java
index 6cd6f05933a..246b7875692 100644
--- a/container-core/src/test/java/com/yahoo/jdisc/http/server/jetty/ProxyProtocolTest.java
+++ b/container-core/src/test/java/com/yahoo/jdisc/http/server/jetty/ProxyProtocolTest.java
@@ -9,7 +9,6 @@ import com.yahoo.jdisc.http.ConnectorConfig;
 import com.yahoo.jdisc.http.ServerConfig;
 import com.yahoo.security.SslContextBuilder;
 import org.apache.hc.client5.http.ssl.NoopHostnameVerifier;
-import org.assertj.core.api.Assertions;
 import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.client.api.ContentResponse;
 import org.eclipse.jetty.client.http.HttpClientTransportOverHTTP;
@@ -23,13 +22,14 @@ import org.junit.jupiter.api.io.TempDir;
 import java.io.IOException;
 import java.net.URI;
 import java.nio.file.Path;
-import java.util.Collection;
 import java.util.concurrent.ExecutionException;
+import java.util.function.Predicate;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static com.yahoo.jdisc.http.server.jetty.Utils.generatePrivateKeyAndCertificate;
 import static com.yahoo.yolean.Exceptions.uncheckInterrupted;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.eclipse.jetty.client.ProxyProtocolClientConnectionFactory.V1;
 import static org.eclipse.jetty.client.ProxyProtocolClientConnectionFactory.V2;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -70,12 +70,11 @@ class ProxyProtocolTest {
         JettyTestDriver driver = createSslWithProxyProtocolTestDriver(certificateFile, privateKeyFile, requestLogMock, connectionLog, false);
         sendJettyClientRequest(driver, certificateFile, new V1.Tag(proxiedRemoteAddress, proxiedRemotePort));
         sendJettyClientRequest(driver, certificateFile, new V2.Tag(proxiedRemoteAddress, proxiedRemotePort));
-        assertTrue(driver.close());
 
-        assertLogSize(2, requestLogMock.entries());
+        assertLogSizeAndCloseDriver(driver, requestLogMock, 2, connectionLog, 2);
+
         assertLogEntryHasRemote(requestLogMock.entries().get(0), proxiedRemoteAddress, proxiedRemotePort);
         assertLogEntryHasRemote(requestLogMock.entries().get(1), proxiedRemoteAddress, proxiedRemotePort);
-        Assertions.assertThat(connectionLog.logEntries()).hasSize(2);
         assertLogEntryHasRemote(connectionLog.logEntries().get(0), proxiedRemoteAddress, proxiedRemotePort);
         assertEquals("v1", connectionLog.logEntries().get(0).proxyProtocolVersion().get());
         assertLogEntryHasRemote(connectionLog.logEntries().get(1), proxiedRemoteAddress, proxiedRemotePort);
@@ -90,13 +89,12 @@ class ProxyProtocolTest {
         sendJettyClientRequest(driver, certificateFile, null);
         sendJettyClientRequest(driver, certificateFile, new V1.Tag(proxiedRemoteAddress, 12345));
         sendJettyClientRequest(driver, certificateFile, new V2.Tag(proxiedRemoteAddress, 12345));
-        assertTrue(driver.close());
 
-        assertLogSize(3, requestLogMock.entries());
+        assertLogSizeAndCloseDriver(driver, requestLogMock, 3, connectionLog, 3);
+
         assertLogEntryHasRemote(requestLogMock.entries().get(0), "127.0.0.1", 0);
         assertLogEntryHasRemote(requestLogMock.entries().get(1), proxiedRemoteAddress, 0);
         assertLogEntryHasRemote(requestLogMock.entries().get(2), proxiedRemoteAddress, 0);
-        Assertions.assertThat(connectionLog.logEntries()).hasSize(3);
         assertLogEntryHasRemote(connectionLog.logEntries().get(0), null, 0);
         assertLogEntryHasRemote(connectionLog.logEntries().get(1), proxiedRemoteAddress, 12345);
         assertLogEntryHasRemote(connectionLog.logEntries().get(2), proxiedRemoteAddress, 12345);
@@ -113,11 +111,10 @@ class ProxyProtocolTest {
                 V2.Tag.Command.PROXY, null, V2.Tag.Protocol.STREAM, proxiedRemoteAddress,
                 proxiedRemotePort, proxyLocalAddress, proxyLocalPort, null);
         ContentResponse response = sendJettyClientRequest(driver, certificateFile, v2Tag);
-        assertTrue(driver.close());
 
         int clientPort = Integer.parseInt(response.getHeaders().get("Jdisc-Local-Port"));
         assertNotEquals(proxyLocalPort, clientPort);
-        assertLogSize(1, connectionLog.logEntries());
+        assertLogSizeAndCloseDriver(driver, requestLogMock, 1, connectionLog, 1);
         assertNotEquals(proxyLocalPort, connectionLog.logEntries().get(0).localPort().get().intValue());
     }
 
@@ -162,7 +159,7 @@ class ProxyProtocolTest {
         HttpClient client = createJettyHttpClient(certificateFile);
         ExecutionException cause = null;
         try {
-            int maxAttempts = 3;
+            int maxAttempts = 10;
             for (int attempt = 0; attempt < maxAttempts; attempt++) {
                 try {
                     ContentResponse response = client.newRequest(URI.create("https://localhost:" + testDriver.server().getListenPort() + "/"))
@@ -209,23 +206,34 @@ class ProxyProtocolTest {
 
     private static void assertLogEntryHasRemote(ConnectionLogEntry entry, String expectedAddress, int expectedPort) {
         if (expectedAddress != null) {
-            Assertions.assertThat(entry.remoteAddress()).hasValue(expectedAddress);
+            assertThat(entry.remoteAddress()).hasValue(expectedAddress);
         } else {
-            Assertions.assertThat(entry.remoteAddress()).isEmpty();
+            assertThat(entry.remoteAddress()).isEmpty();
         }
         if (expectedPort > 0) {
-            Assertions.assertThat(entry.remotePort()).hasValue(expectedPort);
+            assertThat(entry.remotePort()).hasValue(expectedPort);
         } else {
-            Assertions.assertThat(entry.remotePort()).isEmpty();
+            assertThat(entry.remotePort()).isEmpty();
         }
     }
 
 
-    private static void assertLogSize(int expectedItems, Collection<?> items) {
-        for (int attempt = 0; attempt < 10; attempt++) {
-            if (items.size() >= expectedItems) break;
-            uncheckInterrupted(() -> Thread.sleep(200));
+    /* Don't close Jetty to early ensuring that the request log is written */
+    private static void assertLogSizeAndCloseDriver(
+            JettyTestDriver driver, InMemoryRequestLog reqLog, int expectedReqLogSize, InMemoryConnectionLog connLog,
+            int expectedConnLogSize) {
+        Predicate<Void> waitCondition = __ ->
+                reqLog.entries().size() < expectedConnLogSize && connLog.logEntries().size() < expectedConnLogSize;
+        await(waitCondition);
+        assertTrue(driver.close());
+        if (waitCondition.test(null)) await(waitCondition);
+        assertThat(reqLog.entries()).hasSize(expectedReqLogSize);
+        assertThat(connLog.logEntries()).hasSize(expectedConnLogSize);
+    }
+
+    private static void await(Predicate<Void> waitCondition) {
+        for (int attempt = 0; attempt < 1000 && waitCondition.test(null); attempt++) {
+            uncheckInterrupted(() -> Thread.sleep(10));
         }
-        Assertions.assertThat(items).hasSize(expectedItems);
     }
 }