diff --git a/minerva/docker.py b/minerva/docker.py index d06e2a0..97eefc4 100644 --- a/minerva/docker.py +++ b/minerva/docker.py @@ -1,7 +1,7 @@ import threading class Docker: - def __init__(self, machine, container, variables={}, stdout=None, stderr=None): + def __init__(self, container, machine=None, variables={}, stdout=None, stderr=None): self.machine = machine self.uri = container self.variables = variables diff --git a/minerva/machine.py b/minerva/machine.py index 1ebc513..5d7b68f 100644 --- a/minerva/machine.py +++ b/minerva/machine.py @@ -5,7 +5,9 @@ import datetime from fabric import Connection import shlex import threading +import os import minerva +import select class Machine: def __init__(self, @@ -143,8 +145,13 @@ class Machine: # Unfortunately, under the hood, it's running /bin/bash -c '...' # You stand informed # - # This is blocking - def cmd(self, command, hide=True, disown=False): + # This creates a pseudo-TTY on the other end + # + # `watch` means it'll print the output live, else it'll wait till it's finished + # `disown` means it'll run in the background + # + # TODO switch this to use `self.ssh.client.exec_command()` + def cmd(self, command, hide=True, disown=False, watch=False): res = self.ssh.run(f"{self.prep_variables()}; {command}", warn=True, hide=hide, @@ -152,6 +159,91 @@ class Machine: return res + # https://github.com/paramiko/paramiko/issues/593#issuecomment-145377328 + # https://stackoverflow.com/questions/23504126/do-you-have-to-check-exit-status-ready-if-you-are-going-to-check-recv-ready/32758464#32758464 + # + def cmd(self, command, hide=True, disown=False, watch=False): + print(self.ssh) + print(self.ssh.client) + + self.ssh.run("echo hello world", warn=True, hide=hide, disown=disown) + + stdin, stdout, stderr = self.ssh.client.exec_command(command) + + # this is the same for all three inputs + channel = stdin.channel + + out_r, out_w = os.pipe() + err_r, err_w = os.pipe() + + # https://stackoverflow.com/a/78765054 + + # 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 to prevent read block hangs + stdout_chunks = [] + stderr_chunks = [] + stdout_chunks.append(channel.recv(len(channel.in_buffer))) + + timeout = 60 + + def fill_buffers(): + # perform 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 buffers are empty + got_chunk = False + readq, _, _ = select.select([channel], [], [], timeout) + + # returns three empty lists on timeout + if not readq: + break + for c in readq: + if c.recv_ready(): + stdout_chunks.append(channel.recv(len(c.in_buffer))) + got_chunk = True + if c.recv_stderr_ready(): + stderr_chunks.append(channel.recv_stderr(len(c.in_stderr_buffer))) + got_chunk = True + # for c + + """ + 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 received exit code + 3) check if input buffers are empty + 4) exit the loop + """ + if (not got_chunk + and channel.exit_status_ready() + and not channel.recv_stderr_ready() + and not channel.recv_ready()): + # indicate that we're not going to read from this channel anymore + channel.shutdown_read() + # close the channel + channel.close() + # remote side is finished and our buffers are empty + break + # if + # while + + # close the pseudofiles + stdout.close() + stderr.close() + + thread = threading.Thread(target = fill_buffers) + thread.start() + + if not disown: + thread.join() + + return (stdout_chunks, stderr_chunks) + #return (os.fdopen(out_r), os.fdopen(err_r), thread) + def write_env_file(self, variables, fname="~/env.list"): vals = "\n".join([f"{var}={val}" for var, val in variables.items()]) @@ -159,6 +251,11 @@ class Machine: return fname + def load_docker(self, docker): + docker.machine = self + return docker + + def aws_docker_login(self, ecr): return self.cmd(f"aws ecr get-login-password --region {self.pier.session.region_name} | " + f"docker login --username AWS --password-stdin {ecr}"