Merge pull request #647 from yue9944882/bugfix/flush-websocket-before-closing

Release/flush websocket buffers on abnormal closing

Author: k8s-ci-robot

util/src/main/java/io/kubernetes/client/util/WebSocketStreamHandler.java

@@ -92,16 +92,36 @@ protected void handleMessage(int stream, InputStream inStream) throws IOExceptio
   public synchronized void close() {
     if (state != State.CLOSED) {
       state = State.CLOSED;
+      try {
+        if (null != socket) {
+          // code 1000 means "Normal Closure"
+          socket.close(1000, "Triggered client-side terminate");
+          log.debug("Successfully closed socket.");
+        }
+      } catch (IOException ex) {
+        log.error("Error on close socket", ex);
+        // continue to flush and release the buffers
+      }
       // Close all output streams.  Caller of getInputStream(int) is responsible
       // for closing returned input streams
       for (PipedOutputStream out : pipedOutput.values()) {
+        try {
+          out.flush();
+        } catch (IOException ex) {
+          log.error("Error on flush", ex);
+        }
         try {
           out.close();
         } catch (IOException ex) {
           log.error("Error on close", ex);
         }
       }
       for (OutputStream out : output.values()) {
+        try {
+          out.flush();
+        } catch (IOException ex) {
+          log.error("Error on flush", ex);
+        }
         try {
           out.close();
         } catch (IOException ex) {

util/src/test/java/io/kubernetes/client/WebsocketStreamHandlerTest.java

@@ -0,0 +1,68 @@
+package io.kubernetes.client;
+
+import static org.junit.Assert.*;
+
+import com.squareup.okhttp.RequestBody;
+import com.squareup.okhttp.ws.WebSocket;
+import io.kubernetes.client.util.WebSocketStreamHandler;
+import java.io.*;
+import okio.Buffer;
+import okio.BufferedSink;
+import okio.Okio;
+import org.junit.Test;
+
+public class WebsocketStreamHandlerTest {
+
+  private static String testProtocol = "test-protocol";
+
+  @Test
+  public void testHandlerReceivingData() throws IOException {
+    int testStreamId = 0;
+    byte testData = 1;
+    byte[] testDatas =
+        new byte[] {(byte) testStreamId, testData, testData}; // first byte stands for stream id,
+    ByteArrayInputStream testBytesInputStream = new ByteArrayInputStream(testDatas);
+
+    WebSocketStreamHandler handler = new WebSocketStreamHandler();
+    MockWebSocket mockWebSocket = new MockWebSocket();
+
+    handler.open(testProtocol, mockWebSocket);
+
+    InputStream inputStream = handler.getInputStream(testStreamId);
+
+    // handler receiving
+    handler.bytesMessage(testBytesInputStream);
+
+    byte[] receivingData = new byte[16];
+    inputStream.read(receivingData);
+    handler.close();
+
+    assertEquals(testData, receivingData[0]);
+    assertEquals(testData, receivingData[1]);
+    assertTrue(mockWebSocket.closed);
+  }
+
+  private class MockWebSocket implements WebSocket {
+    private byte[] data;
+    private boolean pinged = false;
+    private boolean closed = false;
+
+    @Override
+    public void sendMessage(RequestBody message) throws IOException {
+      ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+      BufferedSink sink = Okio.buffer(Okio.sink(outputStream));
+      message.writeTo(sink);
+      this.data = outputStream.toByteArray();
+    }
+
+    @Override
+    public void sendPing(Buffer payload) throws IOException {
+      this.pinged = true;
+    }
+
+    @Override
+    public void close(int code, String reason) throws IOException {
+      this.closed = true;
+    }
+  }
+}