[stdout stream] Added an option to stream outputs

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)