Paramiko sometimes hangs if the buffer is too big.

If the data coming from the node tool command is too large, paramiko can sometimes stall and hang. This pull request loops reads from the buffers to prevent that from happening.

Author: nicholaspeshek

cstar/remote_paramiko.py

@@ -14,6 +14,7 @@
 
 import paramiko.client
 import re
+import select
 
 from cstar.output import err, debug, msg
 from cstar.exceptions import BadSSHHost, BadEnvironmentVariable, NoHostsSpecified
@@ -105,6 +106,49 @@ def run_job(self, file, jobid, timeout=None, env={}):
     """ % (self.escape(dir),)
 
             stdin, stdout, stderr = self.client.exec_command(cmd, timeout=timeout)
+            # get the shared channel for stdout/stderr/stdin
+            channel = stdout.channel
+
+            stdin.close()
+            # indicate that we're not going to write to that channel anymore
+            channel.shutdown_write()
+
+            # read stdout/stderr in order to prevent read block hangs
+            stdout_chunks = []
+            stdout_chunks.append(stdout.channel.recv(len(stdout.channel.in_buffer)))
+            # chunked read to prevent stalls
+            while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
+                # stop if channel was closed prematurely, and there is no data in the buffers.
+                got_chunk = False
+                readq, _, _ = select.select([stdout.channel], [], [], timeout)
+                for c in readq:
+                    if c.recv_ready():
+                        stdout_chunks.append(stdout.channel.recv(len(c.in_buffer)))
+                        got_chunk = True
+                    if c.recv_stderr_ready():
+                        # make sure to read stderr to prevent stall
+                        stderr.channel.recv_stderr(len(c.in_stderr_buffer))
+                        got_chunk = True
+                '''
+                1) make sure that there are at least 2 cycles with no data in the input buffers in order to not exit too early (i.e. cat on a >200k file).
+                2) if no data arrived in the last loop, check if we already received the exit code
+                3) check if input buffers are empty
+                4) exit the loop
+                '''
+                if not got_chunk \
+                    and stdout.channel.exit_status_ready() \
+                    and not stderr.channel.recv_stderr_ready() \
+                    and not stdout.channel.recv_ready():
+                    # indicate that we're not going to read from this channel anymore
+                    stdout.channel.shutdown_read()
+                    # close the channel
+                    stdout.channel.close()
+                    break    # exit as remote side is finished and our bufferes are empty
+
+            # close all the pseudofiles
+            stdout.close()
+            stderr.close()
+
             stdout.channel.recv_exit_status()
             real_output = self.read_file(dir + "/stdout")
             real_error = self.read_file(dir + "/stderr")
@@ -120,11 +164,55 @@ def run(self, argv):
         try:
             self._connect()
             cmd = " ".join(self.escape(s) for s in argv)
-
             stdin, stdout, stderr = self.client.exec_command(cmd)
+            # get the shared channel for stdout/stderr/stdin
+            channel = stdout.channel
+
+            # we do not need stdin.
+            stdin.close()
+            # indicate that we're not going to write to that channel anymore
+            channel.shutdown_write()
+
+            # read stdout/stderr in order to prevent read block hangs
+            stdout_chunks = []
+            stderr_chunks = []
+            stdout_chunks.append(stdout.channel.recv(len(stdout.channel.in_buffer)))
+            # chunked read to prevent stalls
+            while not channel.closed or channel.recv_ready() or channel.recv_stderr_ready():
+                # stop if channel was closed prematurely, and there is no data in the buffers.
+                got_chunk = False
+                readq, _, _ = select.select([stdout.channel], [], [], 30)
+                for c in readq:
+                    if c.recv_ready():
+                        stdout_chunks.append(stdout.channel.recv(len(c.in_buffer)))
+                        got_chunk = True
+                    if c.recv_stderr_ready():
+                        # make sure to read stderr to prevent stall
+                        stderr_chunks.append(stderr.channel.recv_stderr(len(c.in_stderr_buffer)))
+                        got_chunk = True
+                '''
+                1) make sure that there are at least 2 cycles with no data in the input buffers in order to not exit too early (i.e. cat on a >200k file).
+                2) if no data arrived in the last loop, check if we already received the exit code
+                3) check if input buffers are empty
+                4) exit the loop
+                '''
+                if not got_chunk \
+                    and stdout.channel.exit_status_ready() \
+                    and not stderr.channel.recv_stderr_ready() \
+                    and not stdout.channel.recv_ready():
+                    # indicate that we're not going to read from this channel anymore
+                    stdout.channel.shutdown_read()
+                    # close the channel
+                    stdout.channel.close()
+                    break    # exit as remote side is finished and our bufferes are empty
+
+            # close all the pseudofiles
+            stdout.close()
+            stderr.close()
+
+            out = b''.join(stdout_chunks)
+            error = b''.join(stderr_chunks)
             status = stdout.channel.recv_exit_status()
-            out = stdout.read()
-            error = stderr.read()
             if status != 0:
                 err("Command %s failed with status %d on host %s" % (cmd, status, self.hostname))
             else: