[JENKINS-64294] Buffer stdin for jenkins-cli.jar

dengbuladuo · Apr 30, 2021 · 82f3bfd · 82f3bfd
1 parent 5382345
commit 82f3bfd
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 4 deletions.
diff --git a/cli/src/main/java/hudson/cli/CLI.java b/cli/src/main/java/hudson/cli/CLI.java
@@ -422,10 +422,14 @@ void start(List<String> args) throws IOException {
                 public void run() {
                     try {
                         final OutputStream stdin = streamStdin();
-                        int c;
-                        // TODO check available to avoid sending lots of one-byte frames
-                        while (!complete && (c = System.in.read()) != -1) {
-                           stdin.write(c);
+                        byte[] buf = new byte[60_000]; // less than 64Kb frame size for WS
+                        while (!complete) {
+                            int len = System.in.read(buf);
+                            if (len == -1) {
+                                break;
+                            } else {
+                                stdin.write(buf, 0, len);
+                            }
                         }
                         sendEndStdin();
                     } catch (IOException x) {

diff --git a/core/src/main/java/hudson/cli/CLIAction.java b/core/src/main/java/hudson/cli/CLIAction.java
@@ -124,10 +124,13 @@ public HttpResponse doWs() {
         Authentication authentication = Jenkins.getAuthentication2();
         return WebSockets.upgrade(new WebSocketSession() {
             ServerSideImpl connection;
+            long sentBytes, sentCount, receivedBytes, receivedCount;
             class OutputImpl implements PlainCLIProtocol.Output {
                 @Override
                 public void send(byte[] data) throws IOException {
                     sendBinary(ByteBuffer.wrap(data));
+                    sentBytes += data.length;
+                    sentCount++;
                 }
                 @Override
                 public void close() throws IOException {
@@ -161,6 +164,8 @@ protected void opened() {
             protected void binary(byte[] payload, int offset, int len) {
                 try {
                     connection.handle(new DataInputStream(new ByteArrayInputStream(payload, offset, len)));
+                    receivedBytes += len;
+                    receivedCount++;
                 } catch (IOException x) {
                     error(x);
                 }
@@ -172,6 +177,7 @@ protected void error(Throwable cause) {
             @Override
             protected void closed(int statusCode, String reason) {
                 LOGGER.fine(() -> "closed: " + statusCode + ": " + reason);
+                LOGGER.fine(() -> "received " + receivedCount + " packets of " + receivedBytes + " bytes; sent " + sentCount + " packets of " + sentBytes + " bytes");
                 connection.handleClose();
             }
         });

diff --git a/test/src/test/java/hudson/cli/CLIActionTest.java b/test/src/test/java/hudson/cli/CLIActionTest.java
@@ -8,9 +8,13 @@
 import hudson.model.User;
 import hudson.util.ProcessTree;
 import hudson.util.StreamTaskListener;
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.io.PrintWriter;
@@ -27,6 +31,10 @@
 import jenkins.util.FullDuplexHttpService;
 import jenkins.util.Timer;
 import org.apache.commons.io.FileUtils;
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.io.input.NullInputStream;
+import org.apache.commons.io.output.CountingOutputStream;
+import org.apache.commons.io.output.NullOutputStream;
 import org.apache.commons.io.output.TeeOutputStream;
 import static org.junit.Assert.assertEquals;
 
@@ -38,6 +46,7 @@
 import org.jvnet.hudson.test.LoggerRule;
 import org.jvnet.hudson.test.MockAuthorizationStrategy;
 import org.jvnet.hudson.test.TestExtension;
+import org.kohsuke.args4j.Option;
 
 public class CLIActionTest {
     @Rule
@@ -209,4 +218,73 @@ protected int run() throws Exception {
 
     }
 
+    @Issue("JENKINS-64294")
+    @Test
+    public void largeTransferWebSocket() throws Exception {
+        logging.record(CLIAction.class, Level.FINE);
+        File jar = tmp.newFile("jenkins-cli.jar");
+        FileUtils.copyURLToFile(j.jenkins.getJnlpJars("jenkins-cli.jar").getURL(), jar);
+        CountingOutputStream cos = new CountingOutputStream(NullOutputStream.NULL_OUTPUT_STREAM);
+        long size = /*999_*/999_999;
+        // Download:
+        assertEquals(0, new Launcher.LocalLauncher(StreamTaskListener.fromStderr()).launch().cmds(
+            "java", "-jar", jar.getAbsolutePath(),
+                "-webSocket",
+                "-s", j.getURL().toString(),
+                "large-download",
+                "-size", Long.toString(size)).
+            stdout(cos).stderr(System.err).join());
+        assertEquals(size, cos.getByteCount());
+        // Upload:
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        assertEquals(0, new Launcher.LocalLauncher(StreamTaskListener.fromStderr()).launch().cmds(
+            "java", "-jar", jar.getAbsolutePath(),
+                "-webSocket",
+                "-s", j.getURL().toString(),
+                "large-upload").
+            stdin(new NullInputStream(size)).
+            stdout(baos).stderr(System.err).join());
+        assertEquals("received " + size + " bytes", baos.toString().trim());
+    }
+
+    @TestExtension("largeTransferWebSocket")
+    public static final class LargeUploadCommand extends CLICommand {
+        @Override
+        protected int run() throws Exception {
+            try (InputStream is = new BufferedInputStream(stdin); CountingOutputStream cos = new CountingOutputStream(NullOutputStream.NULL_OUTPUT_STREAM)) {
+                System.err.println("starting upload");
+                long start = System.nanoTime();
+                IOUtils.copyLarge(is, cos);
+                System.err.printf("finished upload in %.1fs%n", (System.nanoTime() - start) / 1_000_000_000.0);
+                stdout.println("received " + cos.getByteCount() + " bytes");
+                stdout.flush();
+            }
+            return 0;
+        }
+        @Override
+        public String getShortDescription() {
+            return "";
+        }
+    }
+
+    @TestExtension("largeTransferWebSocket")
+    public static final class LargeDownloadCommand extends CLICommand {
+        @Option(name = "-size", required = true)
+        public int size;
+        @Override
+        protected int run() throws Exception {
+            try (OutputStream os = new BufferedOutputStream(stdout)) {
+                System.err.println("starting download");
+                long start = System.nanoTime();
+                IOUtils.copyLarge(new NullInputStream(size), os);
+                System.err.printf("finished download in %.1fs%n", ((System.nanoTime() - start) / 1_000_000_000.0));
+            }
+            return 0;
+        }
+        @Override
+        public String getShortDescription() {
+            return "";
+        }
+    }
+
 }