Merge pull request #2388 from dani8art/wait-before-adding-data-to-queue

fix: WebSocketOutputStream wait to write until WebSocket queue is released

Author: k8s-ci-robot

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

@@ -22,6 +22,8 @@
 import java.io.PipedOutputStream;
 import java.io.Reader;
 import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.HashMap;
 import java.util.Map;
 import okhttp3.WebSocket;
@@ -210,6 +212,13 @@ private synchronized OutputStream getSocketInputOutputStream(int stream) {
   }
 
   private class WebSocketOutputStream extends OutputStream {
+
+    private static final long MAX_QUEUE_SIZE = 16L * 1024 * 1024;
+
+    private static final int MAX_WAIT_MILLIS = 10000;
+
+    private static final int WAIT_MILLIS = 10;
+
     private final byte stream;
 
     public WebSocketOutputStream(int stream) {
@@ -266,10 +275,28 @@ public void write(byte[] b, int offset, int length) throws IOException {
         int bufferSize = Math.min(remaining, 15 * 1024 * 1024);
         byte[] buffer = new byte[bufferSize + 1];
         buffer[0] = stream;
+
         System.arraycopy(b, offset + bytesWritten, buffer, 1, bufferSize);
-        if (!WebSocketStreamHandler.this.socket.send(ByteString.of(buffer))) {
-          throw new IOException("WebSocket has closed.");
+        ByteString byteString = ByteString.of(buffer);
+
+        final Instant start = Instant.now();
+        synchronized (WebSocketOutputStream.this) {
+          while (WebSocketStreamHandler.this.socket.queueSize() + byteString.size() > MAX_QUEUE_SIZE
+              && Instant.now().isBefore(start.plus(MAX_WAIT_MILLIS, ChronoUnit.MILLIS))) {
+            try {
+              wait(WAIT_MILLIS);
+            } catch (InterruptedException e) {
+              throw new IOException("Error waiting web socket queue", e);
+            }
+          }
+
+          if (!WebSocketStreamHandler.this.socket.send(byteString)) {
+            throw new IOException("WebSocket has closed.");
+          }
+
+          notifyAll();
         }
+
         bytesWritten += bufferSize;
         remaining -= bufferSize;
       }