diff --git a/util/src/main/java/io/kubernetes/client/util/WebSocketStreamHandler.java b/util/src/main/java/io/kubernetes/client/util/WebSocketStreamHandler.java index e98b8e092e..b15abef7bb 100644 --- a/util/src/main/java/io/kubernetes/client/util/WebSocketStreamHandler.java +++ b/util/src/main/java/io/kubernetes/client/util/WebSocketStreamHandler.java @@ -128,7 +128,11 @@ public synchronized boolean isClosed() { @Override public synchronized void close() { - if (state != State.CLOSED) { + if (state == State.CLOSED) { + notifyAll(); + return; + } + try { // Close all output streams. Caller of getInputStream(int) is responsible // for closing returned input streams for (PipedOutputStream out : pipedOutput.values()) { @@ -155,15 +159,20 @@ public synchronized void close() { log.error("Error on close", ex); } } - + } finally { + // Always transition to CLOSED and wake up any writers blocked waiting for the socket + // to open, even if closing one of the streams above failed unexpectedly. state = State.CLOSED; - if (null != socket) { - // code 1000 means "Normal Closure" - socket.close(1000, "Triggered client-side terminate"); - log.debug("Successfully closed socket."); + try { + if (null != socket) { + // code 1000 means "Normal Closure" + socket.close(1000, "Triggered client-side terminate"); + log.debug("Successfully closed socket."); + } + } finally { + notifyAll(); } } - notifyAll(); } /** @@ -273,8 +282,13 @@ public void flush() throws IOException { if (state == State.CLOSED) { throw new IOException("Socket is closed!"); } + WebSocket ws = WebSocketStreamHandler.this.socket; + if (ws == null) { + // The socket was never opened, so there is nothing queued to flush. + return; + } int i = 0; - while (WebSocketStreamHandler.this.socket.queueSize() > 0) { + while (ws.queueSize() > 0) { try { Thread.sleep(FLUSH_WAIT_MILLIS); } catch (InterruptedException ex) { @@ -300,9 +314,10 @@ public void write(byte[] b) throws IOException { public void write(byte[] b, int offset, int length) throws IOException { if (WebSocketStreamHandler.this.socket == null) { synchronized (WebSocketStreamHandler.this) { - if (state == State.CLOSED) throw new IllegalStateException(); - if (WebSocketStreamHandler.this.socket == null) { - // wait for the websocket to be opened + // wait for the websocket to be opened, failing fast if the handler is closed + // (e.g. because the connection failed) before or while waiting. + while (WebSocketStreamHandler.this.socket == null) { + if (state == State.CLOSED) throw new IllegalStateException(); try { WebSocketStreamHandler.this.wait(); } catch (InterruptedException e) { diff --git a/util/src/test/java/io/kubernetes/client/CopyTest.java b/util/src/test/java/io/kubernetes/client/CopyTest.java index d759cda5a0..cd5511a684 100644 --- a/util/src/test/java/io/kubernetes/client/CopyTest.java +++ b/util/src/test/java/io/kubernetes/client/CopyTest.java @@ -18,13 +18,11 @@ import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor; import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo; import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static org.awaitility.Awaitility.await; -import com.github.tomakehurst.wiremock.core.Admin; -import com.github.tomakehurst.wiremock.extension.Parameters; -import com.github.tomakehurst.wiremock.extension.PostServeAction; import com.github.tomakehurst.wiremock.junit5.WireMockExtension; import com.github.tomakehurst.wiremock.matching.AnythingPattern; -import com.github.tomakehurst.wiremock.stubbing.ServeEvent; +import com.github.tomakehurst.wiremock.matching.RequestPatternBuilder; import io.kubernetes.client.openapi.ApiClient; import io.kubernetes.client.openapi.ApiException; import io.kubernetes.client.openapi.models.V1ObjectMeta; @@ -40,7 +38,6 @@ import java.nio.file.Paths; import java.util.HashMap; import java.util.Map; -import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -125,25 +122,59 @@ public InputStream copyFileFromPod(String namespace, String pod, String srcPath) } } - public static class CountDownLatchAction extends PostServeAction { - @Override - public String getName() { - return "countdown"; - } - - @Override - public void doAction(ServeEvent serveEvent, Admin admin, Parameters parameters) { - CountDownLatch latch = (CountDownLatch) parameters.get("latch"); - latch.countDown(); - } - } - @RegisterExtension static WireMockExtension apiServer = WireMockExtension.newInstance() - .options(wireMockConfig().dynamicPort().extensions(new CountDownLatchAction())) + .options(wireMockConfig().dynamicPort()) .build(); + private RequestPatternBuilder execRequest() { + return getRequestedFor( + urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")); + } + + private void stubExecNotFound() { + apiServer.stubFor( + get(urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) + .willReturn( + aResponse() + .withStatus(404) + .withHeader("Content-Type", "application/json") + .withBody("{}"))); + } + + /** + * Runs the given action on a background thread, waits until WireMock has recorded an exec + * request, then stops the thread and waits for it to terminate so it cannot leak into later + * tests. + */ + private void runUntilExecRequested(ThrowingRunnable action) throws InterruptedException { + Thread t = + new Thread( + () -> { + try { + action.run(); + } catch (Exception ex) { + // Expected: the mock server rejects the websocket upgrade. + } + }); + t.start(); + try { + await() + .atMost(30, TimeUnit.SECONDS) + .until(() -> !apiServer.findAll(execRequest()).isEmpty()); + } finally { + t.interrupt(); + t.join(TimeUnit.SECONDS.toMillis(30)); + } + assertFalse(t.isAlive(), "background copy thread did not terminate"); + } + + @FunctionalInterface + private interface ThrowingRunnable { + void run() throws Exception; + } + @BeforeEach void setup() { client = new ClientBuilder().setBasePath("http://localhost:" + apiServer.getPort()).build(); @@ -187,46 +218,19 @@ void url() { @Test void copyFileToPod(@TempDir Path tempDir) throws Exception { - Path testFile = Files.createTempFile(tempDir, "testfile", null); Copy copy = new Copy(client); - CountDownLatch latch = new CountDownLatch(1); - Parameters params = new Parameters(); - params.put("latch", latch); + stubExecNotFound(); - apiServer.stubFor( - get(urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) - .withPostServeAction("countdown", params) - .willReturn( - aResponse() - .withStatus(404) - .withHeader("Content-Type", "application/json") - .withBody("{}"))); - - // When attempting to write to the process outputstream in copyFileToPod, the - // WebSocketStreamHandler is in a wait state because no websocket is created by mock, which - // blocks the main thread. So here we execute the method in a thread. - Thread t = - new Thread( - new Runnable() { - public void run() { - try { - copy.copyFileToPod( - namespace, podName, "", testFile, Paths.get("/copied-testfile")); - } catch (IOException | ApiException ex) { - ex.printStackTrace(); - } - } - }); - t.start(); - assertTrue(latch.await(30, TimeUnit.SECONDS)); - t.interrupt(); + // The mock server never upgrades to a websocket, so run the copy in the background and only + // wait for the exec request to reach the server. + runUntilExecRequested( + () -> copy.copyFileToPod(namespace, podName, "", testFile, Paths.get("/copied-testfile"))); apiServer.verify( - getRequestedFor( - urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) + execRequest() .withQueryParam("stdin", equalTo("true")) .withQueryParam("stdout", equalTo("true")) .withQueryParam("stderr", equalTo("true")) @@ -238,46 +242,17 @@ public void run() { @Test void copyBinaryDataToPod() throws InterruptedException { - byte[] testSrc = new byte[0]; Copy copy = new Copy(client); - CountDownLatch latch = new CountDownLatch(1); - Parameters params = new Parameters(); - params.put("latch", latch); - - apiServer.stubFor( - get(urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) - .withPostServeAction("countdown", params) - .willReturn( - aResponse() - .withStatus(404) - .withHeader("Content-Type", "application/json") - .withBody("{}"))); + stubExecNotFound(); - // When attempting to write to the process outputstream in copyFileToPod, the - // WebSocketStreamHandler is in a wait state because no websocket is created by mock, which - // blocks the main thread. So here we execute the method in a thread. - Thread t = - new Thread( - new Runnable() { - public void run() { - try { - copy.copyFileToPod( - namespace, podName, "", testSrc, Paths.get("/copied-binarydata")); - } catch (IOException | ApiException ex) { - ex.printStackTrace(); - } - } - }); - t.start(); - assertTrue(latch.await(30, TimeUnit.SECONDS)); - t.interrupt(); + runUntilExecRequested( + () -> copy.copyFileToPod(namespace, podName, "", testSrc, Paths.get("/copied-binarydata"))); apiServer.verify( - getRequestedFor( - urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) + execRequest() .withQueryParam("stdin", equalTo("true")) .withQueryParam("stdout", equalTo("true")) .withQueryParam("stderr", equalTo("true")) @@ -291,42 +266,15 @@ public void run() { void testCopyDirectoryFromPod(@TempDir Path tempDir) throws Exception { Copy copy = new Copy(client); - CountDownLatch latch = new CountDownLatch(1); - Parameters params = new Parameters(); - params.put("latch", latch); - - apiServer.stubFor( - get(urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) - .withPostServeAction("countdown", params) - .willReturn( - aResponse() - .withStatus(404) - .withHeader("Content-Type", "application/json") - .withBody("{}"))); + stubExecNotFound(); - Thread t = - new Thread( - new Runnable() { - public void run() { - try { - copy.copyDirectoryFromPod( - namespace, - podName, - "", - tempDir.toString(), - Paths.get("/copied-testDir")); - } catch (IOException | ApiException | CopyNotSupportedException ex) { - ex.printStackTrace(); - } - } - }); - t.start(); - assertTrue(latch.await(30, TimeUnit.SECONDS)); - t.interrupt(); + runUntilExecRequested( + () -> + copy.copyDirectoryFromPod( + namespace, podName, "", tempDir.toString(), Paths.get("/copied-testDir"))); apiServer.verify( - getRequestedFor( - urlPathEqualTo("/api/v1/namespaces/" + namespace + "/pods/" + podName + "/exec")) + execRequest() .withQueryParam("stdin", equalTo("false")) .withQueryParam("stdout", equalTo("true")) .withQueryParam("stderr", equalTo("true")) diff --git a/util/src/test/java/io/kubernetes/client/WebsocketStreamHandlerTest.java b/util/src/test/java/io/kubernetes/client/WebsocketStreamHandlerTest.java index d901ba6d23..2dd575a6ee 100644 --- a/util/src/test/java/io/kubernetes/client/WebsocketStreamHandlerTest.java +++ b/util/src/test/java/io/kubernetes/client/WebsocketStreamHandlerTest.java @@ -14,6 +14,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.awaitility.Awaitility.await; import io.kubernetes.client.util.WebSocketStreamHandler; import java.io.ByteArrayInputStream; @@ -22,6 +23,8 @@ import java.io.OutputStream; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import okhttp3.Request; import okhttp3.WebSocket; import okio.ByteString; @@ -186,6 +189,53 @@ void handlerIsClosed() throws IOException { assertThat(mockWebSocket.closed).isTrue(); } + @Test + void handlerClosedBeforeOpenWithOutputStream() { + WebSocketStreamHandler handler = new WebSocketStreamHandler(); + handler.getOutputStream(0); + + // Simulates a connection failure (e.g. a non-101 upgrade response) before the socket opens. + handler.close(); + + assertThat(handler.isClosed()).isTrue(); + } + + @Test + void writeFailsWhenHandlerClosedBeforeOpen() throws Exception { + WebSocketStreamHandler handler = new WebSocketStreamHandler(); + OutputStream outputStream = handler.getOutputStream(0); + handler.close(); + + assertThatThrownBy(() -> outputStream.write(new byte[] {1})) + .isInstanceOf(IllegalStateException.class); + } + + @Test + void blockedWriteIsReleasedWhenHandlerClosedBeforeOpen() throws Exception { + WebSocketStreamHandler handler = new WebSocketStreamHandler(); + OutputStream outputStream = handler.getOutputStream(0); + + AtomicReference writeError = new AtomicReference<>(); + Thread writer = + new Thread( + () -> { + try { + outputStream.write(new byte[] {1}); + } catch (Throwable t) { + writeError.set(t); + } + }); + writer.start(); + + // Wait until the writer is blocked waiting for the socket to be opened. + await().atMost(10, TimeUnit.SECONDS).until(() -> writer.getState() == Thread.State.WAITING); + handler.close(); + + writer.join(TimeUnit.SECONDS.toMillis(10)); + assertThat(writer.isAlive()).isFalse(); + assertThat(writeError.get()).isInstanceOf(IllegalStateException.class); + } + private static class MockWebSocket implements WebSocket { byte[] data; private boolean closed = false;