Skip to content

Commit 8c01855

Browse files
Danpkittenis
authored andcommitted
Don't close channel on get_stout, check if command has finished instead and return buffers if not. Channel is now also returned from get_stdout to allow user to get exit status on command completion. This resolves #11
1 parent 5847030 commit 8c01855

File tree

1 file changed

+30
-5
lines changed

1 file changed

+30
-5
lines changed

pssh.py

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -430,23 +430,48 @@ def get_stdout(self, greenlet, return_buffers=False):
430430
:rtype: Dictionary containing ``{host: {'exit_code': exit code}}`` entry \
431431
for example ``{'myhost1': {'exit_code': 0}}``
432432
:rtype: With ``return_buffers=True``: ``{'myhost1': {'exit_code': 0,
433+
'channel' : None or SSH channel of command if command is still executing,
433434
'stdout' : <iterable>,
434435
'stderr' : <iterable>,}}``
435436
"""
436437
channel, host, _stdout, _stderr = greenlet.get()
437-
stdout = (line.strip() for line in _stdout)
438-
stderr = (line.strip() for line in _stderr)
439-
channel.close()
440-
if not return_buffers:
438+
stdout = self._read_output_buffer(_stdout)
439+
stderr = self._read_output_buffer(_stderr)
440+
if channel.exit_status_ready():
441+
channel.close()
442+
else:
443+
logger.debug("Command still executing on get_stdout call - not closing channel and returning None as exit code.")
444+
# If channel is not closed we cannot get full stdout/stderr so must return buffers
445+
return_buffers = True
446+
# Channel must be closed or reading stdout/stderr will block forever
447+
if not return_buffers and channel.closed:
441448
for line in stdout:
442449
host_logger.info("[%s]\t%s", host, line,)
443450
for line in stderr:
444451
host_logger.info("[%s] [err] %s", host, line,)
445452
return {host: {'exit_code': channel.recv_exit_status(),}}
446-
return {host: {'exit_code': channel.recv_exit_status(),
453+
gevent.sleep(.2)
454+
return {host: {'exit_code': channel.recv_exit_status() if channel.exit_status_ready() else None,
455+
'channel' : channel if not channel.closed else None,
447456
'stdout' : stdout,
448457
'stderr' : stderr, }}
449458

459+
def wait_for_exit_status(self, channel):
460+
"""Block and wait for exit status on channel.
461+
WARNING - this will block forever if the command executed never exits
462+
:rtype: int - Exit code of command executed"""
463+
while not channel.exit_status_ready():
464+
gevent.sleep(1)
465+
channel.close()
466+
return channel.recv_exit_status()
467+
468+
def _read_output_buffer(self, output_buffer):
469+
"""Read from output buffers,
470+
allowing coroutines to execute in between reading"""
471+
for line in output_buffer:
472+
gevent.sleep(1)
473+
yield line.strip()
474+
450475
def copy_file(self, local_file, remote_file):
451476
"""Copy local file to remote file in parallel
452477

0 commit comments

Comments
 (0)