[JENKINS-61409] Websockets: Use AbstractByteBufferCommandTransport to transport messages (#4596)

* [JENKINS-61409] Websockets: Use AbstractByteBufferCommandTransport to transport messages

* Depend on timestamped snapshot since incrementals is broken

* Fix reviews

* Update snapshot dep

* Print directly

core/src/main/java/jenkins/agents/WebSocketAgents.java

@@ -30,17 +30,16 @@ import hudson.ExtensionList;
 import hudson.model.Computer;
 import hudson.model.InvisibleAction;
 import hudson.model.UnprotectedRootAction;
-import hudson.remoting.AbstractByteArrayCommandTransport;
+import hudson.remoting.AbstractByteBufferCommandTransport;
 import hudson.remoting.Capability;
-import hudson.remoting.Channel;
 import hudson.remoting.ChannelBuilder;
+import hudson.remoting.ChunkHeader;
 import hudson.remoting.Engine;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ClosedChannelException;
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -106,7 +105,7 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
         private final JnlpConnectionState state;
         private final String agent;
         private final Capability remoteCapability;
-        private AbstractByteArrayCommandTransport.ByteArrayReceiver receiver;
+        private Transport transport;
 
         Session(JnlpConnectionState state, String agent, Capability remoteCapability) {
             this.state = state;
@@ -119,7 +118,8 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
             Computer.threadPoolForRemoting.submit(() -> {
                 LOGGER.fine(() -> "setting up channel for " + agent);
                 state.fireBeforeChannel(new ChannelBuilder(agent, Computer.threadPoolForRemoting));
-                state.fireAfterChannel(state.getChannelBuilder().build(new Transport()));
+                transport = new Transport();
+                state.fireAfterChannel(state.getChannelBuilder().build(transport));
                 LOGGER.fine(() -> "set up channel for " + agent);
                 return null;
             });
@@ -128,10 +128,10 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
         @Override
         protected void binary(byte[] payload, int offset, int len) {
             LOGGER.finest(() -> "reading block of length " + len + " from " + agent);
-            if (offset == 0 && len == payload.length) {
-                receiver.handle(payload);
-            } else {
-                receiver.handle(Arrays.copyOfRange(payload, offset, offset + len));
+            try {
+                transport.receive(ByteBuffer.wrap(payload, offset, len));
+            } catch (IOException | InterruptedException e) {
+                error(e);
             }
         }
 
@@ -139,7 +139,7 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
         protected void closed(int statusCode, String reason) {
             LOGGER.finest(() -> "closed " + statusCode + " " + reason);
             IOException x = new ClosedChannelException();
-            receiver.terminate(x);
+            transport.terminate(x);
             state.fireChannelClosed(x);
             state.fireAfterDisconnect();
         }
@@ -149,22 +149,13 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
             LOGGER.log(Level.WARNING, null, cause);
         }
 
-        class Transport extends AbstractByteArrayCommandTransport {
+        class Transport extends AbstractByteBufferCommandTransport {
 
             @Override
-            public void setup(AbstractByteArrayCommandTransport.ByteArrayReceiver bar) {
-                receiver = bar;
-            }
-
-            @Override
-            public void writeBlock(Channel chnl, byte[] bytes) throws IOException {
-                LOGGER.finest(() -> "writing block of length " + bytes.length + " to " + agent);
-                try {
-                    sendBinary(ByteBuffer.wrap(bytes)).get();
-                } catch (Exception x) {
-                    x.printStackTrace();
-                    throw new IOException(x);
-                }
+            protected void write(ByteBuffer header, ByteBuffer data) throws IOException {
+                LOGGER.finest(() -> "sending message of length " + ChunkHeader.length(ChunkHeader.peek(header)));
+                sendBinary(header, false);
+                sendBinary(data, true);
             }
 
             @Override
@@ -183,7 +174,6 @@ public final class WebSocketAgents extends InvisibleAction implements Unprotecte
                 LOGGER.finest(() -> "closeRead");
                 close();
             }
-
         }
 
     }

core/src/main/java/jenkins/websocket/WebSocketSession.java

@@ -60,6 +60,7 @@ public abstract class WebSocketSession {
     private static final Logger LOGGER = Logger.getLogger(WebSocketSession.class.getName());
 
     private Object session;
+    // https://www.eclipse.org/jetty/javadoc/9.4.24.v20191120/org/eclipse/jetty/websocket/common/WebSocketRemoteEndpoint.html
     private Object remoteEndpoint;
     private ScheduledFuture<?> pings;
 
@@ -130,6 +131,15 @@ public abstract class WebSocketSession {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    protected final void sendBinary(ByteBuffer partialByte, boolean isLast) {
+        try {
+            remoteEndpoint.getClass().getMethod("sendPartialBytes", ByteBuffer.class, boolean.class).invoke(remoteEndpoint, partialByte, isLast);
+        } catch (Exception x) {
+            throw new RuntimeException(x);
+        }
+    }
+
     @SuppressWarnings("unchecked")
     protected final Future<Void> sendText(String text) {
         try {

pom.xml

@@ -101,7 +101,7 @@ THE SOFTWARE.
     <maven-war-plugin.version>3.2.3</maven-war-plugin.version>
 
     <!-- Bundled Remoting version -->
-    <remoting.version>4.2</remoting.version>
+    <remoting.version>4.3-20200323.083513-2</remoting.version><!-- https://github.com/jenkinsci/remoting/pull/373 -->
     <!-- Minimum Remoting version, which is tested for API compatibility -->
     <remoting.minimum.supported.version>3.14</remoting.minimum.supported.version>
 

test/src/test/java/jenkins/agents/WebSocketAgentsTest.java

@@ -35,6 +35,7 @@ import hudson.slaves.SlaveComputer;
 import hudson.tasks.BatchFile;
 import hudson.tasks.Shell;
 import java.io.File;
+import java.util.Random;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -99,9 +100,10 @@ public class WebSocketAgentsTest {
             ).stdout(System.out).start());
             r.waitOnline(s);
             assertEquals("response", s.getChannel().call(new DummyTask()));
+            assertNotNull(s.getChannel().call(new FatTask()));
             FreeStyleProject p = r.createFreeStyleProject();
             p.setAssignedNode(s);
-            p.getBuildersList().add(Functions.isWindows() ? new BatchFile("echo hello") : new Shell("echo hello"));
+            p.getBuildersList().add(Functions.isWindows() ? new BatchFile("echo hello") : new Shell("dd if=/dev/random count=1024 bs=200"));
             r.buildAndAssertSuccess(p);
             s.toComputer().getLogText().writeLogTo(0, System.out);
         } finally {
@@ -122,4 +124,18 @@ public class WebSocketAgentsTest {
         }
     }
 
+    private static class FatTask extends SlaveToMasterCallable<String, RuntimeException> {
+        private byte[] payload;
+
+        private FatTask() {
+            payload = new byte[1024 * 1024];
+            new Random().nextBytes(payload);
+        }
+
+        @Override
+        public String call() {
+            return new String(payload);
+        }
+    }
+
 }