diff --git a/datadog/dogshell/wrap.py b/datadog/dogshell/wrap.py index 99abc61a9..31f37517e 100644 --- a/datadog/dogshell/wrap.py +++ b/datadog/dogshell/wrap.py @@ -19,9 +19,11 @@ ''' import sys -import subprocess import time -from optparse import OptionParser +import optparse +import threading +import subprocess +import pkg_resources as pkg from datadog import initialize, api @@ -34,7 +36,49 @@ class Timeout(Exception): pass +class OutputReader(threading.Thread): + ''' + Thread collecting the output of a subprocess, optionally forwarding it to + a given file descriptor and storing it for further retrieval. + ''' + def __init__(self, proc_out, fwd_out=None): + ''' + Instantiates an OutputReader. + :param proc_out: the output to read + :type proc_out: file descriptor + :param fwd_out: the output to forward to (None to disable forwarding) + :type fwd_out: file descriptor or None + ''' + threading.Thread.__init__(self) + self.daemon = True + self._out_content = '' + self._out = proc_out + self._fwd_out = fwd_out + + def run(self): + ''' + Thread's main loop: collects the output optionnally forwarding it to + the file descriptor passed in the constructor. + ''' + for line in iter(self._out.readline, b''): + if self._fwd_out is not None: + self._fwd_out.write(line) + + self._out_content += line + self._out.close() + + @property + def content(self): + ''' + The content stored in out so far. (Not threadsafe, wait with .join()) + ''' + return self._out_content + + def poll_proc(proc, sleep_interval, timeout): + ''' + Polls the process until it returns or a given timeout has been reached + ''' start_time = time.time() returncode = None while returncode is None: @@ -47,7 +91,10 @@ def poll_proc(proc, sleep_interval, timeout): def execute(cmd, cmd_timeout, sigterm_timeout, sigkill_timeout, - proc_poll_interval): + proc_poll_interval, buffer_outs): + ''' + Launches the process and monitors its outputs + ''' start_time = time.time() returncode = -1 stdout = '' @@ -59,8 +106,24 @@ def execute(cmd, cmd_timeout, sigterm_timeout, sigkill_timeout, print >> sys.stderr, u"Failed to execute %s" % (repr(cmd)) raise try: + # Let's that the threads collecting the output from the command in the + # background + out_reader = OutputReader(proc.stdout, sys.stdout if not buffer_outs else None) + err_reader = OutputReader(proc.stderr, sys.stderr if not buffer_outs else None) + out_reader.start() + err_reader.start() + + # Let's quietly wait from the program's completion here et get the exit + # code when it finishes returncode = poll_proc(proc, proc_poll_interval, cmd_timeout) - stdout, stderr = proc.communicate() + + # Let's harvest the outputs collected by our background threads after + # making sure they're done reading it. + out_reader.join() + err_reader.join() + stdout = out_reader.content + stderr = err_reader.content + duration = time.time() - start_time except Timeout: duration = time.time() - start_time @@ -86,29 +149,58 @@ def execute(cmd, cmd_timeout, sigterm_timeout, sigkill_timeout, def main(): - parser = OptionParser() - parser.add_option('-n', '--name', action='store', type='string', help="The name of the event") - parser.add_option('-k', '--api_key', action='store', type='string') + parser = optparse.OptionParser(usage="%prog -n [event_name] -k [api_key] --submit_mode i\ +[ all | errors ] [options] \"command\". \n\nNote that you need to enclose your command in \ +quotes to prevent python as soon as there is a space in your command. \n \nNOTICE: In normal \ +mode, the whole stderr is printed before stdout, in flush_live mode they will be mixed but there \ +is not guarantee that messages sent by the command on both stderr and stdout are printed in the \ +order they were sent.", version="%prog {0}".format(pkg.require("datadog")[0].version)) + + parser.add_option('-n', '--name', action='store', type='string', help="the name of the event \ +as it should appear on your Datadog stream") + parser.add_option('-k', '--api_key', action='store', type='string', + help="your DataDog API Key") parser.add_option('-m', '--submit_mode', action='store', type='choice', - default='errors', choices=['errors', 'all']) - parser.add_option('-t', '--timeout', action='store', type='int', default=60 * 60 * 24) + default='errors', choices=['errors', 'all'], help="[ all | errors ] if set \ +to error, an event will be sent only of the command exits with a non zero exit status or if it \ +times out.") parser.add_option('-p', '--priority', action='store', type='choice', choices=['normal', 'low'], - help="The priority of the event (default: 'normal')") - parser.add_option('--sigterm_timeout', action='store', type='int', default=60 * 2) - parser.add_option('--sigkill_timeout', action='store', type='int', default=60) - parser.add_option('--proc_poll_interval', action='store', type='float', default=0.5) - parser.add_option('--notify_success', action='store', type='string', default='') - parser.add_option('--notify_error', action='store', type='string', default='') + help="the priority of the event (default: 'normal')") + parser.add_option('-t', '--timeout', action='store', type='int', default=60 * 60 * 24, + help="(in seconds) a timeout after which your command must be aborted. An \ +event will be sent to your DataDog stream (default: 24hours)") + parser.add_option('--sigterm_timeout', action='store', type='int', default=60 * 2, + help="(in seconds) When your command times out, the \ +process it triggers is sent a SIGTERM. If this sigterm_timeout is reached, it will be sent a \ +SIGKILL signal. (default: 2m)") + parser.add_option('--sigkill_timeout', action='store', type='int', default=60, + help="(in seconds) how long to wait at most after SIGKILL \ + has been sent (default: 60s)") + parser.add_option('--proc_poll_interval', action='store', type='float', default=0.5, + help="(in seconds). interval at which your command will be polled \ +(default: 500ms)") + parser.add_option('--notify_success', action='store', type='string', default='', + help="a message string and @people directives to send notifications in \ +case of success.") + parser.add_option('--notify_error', action='store', type='string', default='', + help="a message string and @people directives to send notifications in \ +case of error.") + parser.add_option('-b', '--buffer_outs', action='store_true', dest='buffer_outs', default=False, + help="displays the stderr and stdout of the command only once it has \ +returned (the command outputs remains buffered in dogwrap meanwhile)") options, args = parser.parse_args() cmd = [] for part in args: cmd.extend(part.split(' ')) + # If silent is checked we force the outputs to be buffered (and therefore + # not forwarded to the Terminal streams) and we just avoid printing the + # buffers at the end returncode, stdout, stderr, duration = execute( cmd, options.timeout, options.sigterm_timeout, options.sigkill_timeout, - options.proc_poll_interval) + options.proc_poll_interval, options.buffer_outs) initialize(api_key=options.api_key) host = api._host_name @@ -130,8 +222,8 @@ def main(): event_body = [u'%%%\n', u'commmand:\n```\n', u' '.join(cmd), u'\n```\n', - u'exit code: %s\n\n' % returncode, - ] + u'exit code: %s\n\n' % returncode, ] + if stdout: event_body.extend([u'stdout:\n```\n', stdout, u'\n```\n']) if stderr: @@ -156,8 +248,10 @@ def main(): 'host': host, 'priority': options.priority or event_priority, } - print >> sys.stderr, stderr.strip() - print >> sys.stdout, stdout.strip() + + if options.buffer_outs: + print >> sys.stderr, stderr.strip() + print >> sys.stdout, stdout.strip() if options.submit_mode == 'all' or returncode != 0: api.Event.create(title=event_title, text=event_body, **event)