Add LogDog / annotation protobuf support.

recipe_engine/stream.py

 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Abstract stream interface for representing recipe runs.
 
 We need to create streams for steps (and substeps) and also LOG_LINE steps.
 LogDog will implement LOG_LINE steps as real logs (i.e. uniformly), but
 annotations will implement them differently from normal logs, so we need
 a way to distinguish.
 
 StreamEngine will coordinate the multiplexing of streams.  In the case of
 annotations, this involves keeping track of the STEP_CURSOR and setting it
 accordingly, as well as filtering @@@ lines.
 
 Stream is a virtual well-behaved stream (associated with an Engine) which you
 can just write to without worrying.
 """
 
 import json
 import time
 
 from . import env
 from . import recipe_api
+from . import util
 
 
 class StreamEngine(object):
   class Stream(object):
     def write_line(self, line):
       raise NotImplementedError()
 
     def write_split(self, string):
       """Write a string (which may contain newlines) to the stream.  It will
       be terminated by a newline."""
@@ skipping 62 matching lines @@
     pass
 
   def __enter__(self):
     self.open()
     return self
 
   def __exit__(self, _exc_type, _exc_val, _exc_tb):
     self.close()
 
 
-# Because StreamEngine has no observations (i.e. it is an F-Algebra), we can
-# form products.  This code is entirely mechanical from the types (if we
-# had them formalized...).
 class ProductStreamEngine(StreamEngine):
+  """A StreamEngine that forms the non-commutative product of two other
+  StreamEngines.
+
+  Because StreamEngine has no observations (i.e. it is an F-Algebra), we can
+  form products.  This code is entirely mechanical from the types (if we
+  had them formalized...).
+
+  This product is non-commutative, meaning order matters. Specifically, an
+  exception in "engine_a" will prevent "engine_b" from being evaluated.
+  """
+
   def __init__(self, engine_a, engine_b):
     assert engine_a and engine_b
     self._engine_a = engine_a
     self._engine_b = engine_b
 
   class Stream(StreamEngine.Stream):
     def __init__(self, stream_a, stream_b):
       assert stream_a and stream_b
       self._stream_a = stream_a
       self._stream_b = stream_b
 
     def write_line(self, line):
       self._stream_a.write_line(line)
       self._stream_b.write_line(line)
 
     def close(self):
       self._stream_a.close()
       self._stream_b.close()
 
   class StepStream(Stream):
+    # pylint: disable=no-self-argument
     def _void_product(method_name):
       def inner(self, *args):
         getattr(self._stream_a, method_name)(*args)
         getattr(self._stream_b, method_name)(*args)
       return inner
 
     def new_log_stream(self, log_name):
       return ProductStreamEngine.Stream(
           self._stream_a.new_log_stream(log_name),
           self._stream_b.new_log_stream(log_name))
@@ skipping 13 matching lines @@
 
   def open(self):
     self._engine_a.open()
     self._engine_b.open()
 
   def close(self):
     self._engine_a.close()
     self._engine_b.close()
 
 
-def _noop(*args, **kwargs):
+class MultiStreamEngine(StreamEngine):
+  """A StreamEngine consisting of one or more inner StreamEngines.
+
+  A call to this StreamEngine will be distributed to the inner StreamEngines.
+  Any exceptions that are caught during an inner handler will be deferred until
+  all inner handlers have been executed.
+  """
+
+  def __init__(self, base, *engines):
+    self._engines = (base,) + engines
+    assert None not in self._engines
+
+  @classmethod
+  def create(cls, *engines):
+    assert len(engines) > 0, 'At least one engine must be provided.'
+    if len(engines) == 1:
+      return engines[0]
+    return cls(engines[0], *engines[1:])
+
+  class Stream(StreamEngine.Stream):
+    def __init__(self, *streams):
+      assert all(streams)
+      self._streams = streams
+
+    def write_line(self, line):
+      util.map_defer_exceptions(lambda s: s.write_line(line), self._streams)
+
+    def close(self):
+      util.map_defer_exceptions(lambda s: s.close(), self._streams)
+
+  class StepStream(Stream):
+    # pylint: disable=no-self-argument
+    def _multiplex(method_name):
+      def inner(self, *args):
+        util.map_defer_exceptions(lambda s: getattr(s, method_name)(*args),
+                                  self._streams)
+      return inner
+
+    def new_log_stream(self, log_name):
+      log_streams = []
+      try:
+        for s in self._streams:
+          log_streams.append(s.new_log_stream(log_name))
+        return MultiStreamEngine.Stream(*log_streams)
+      except Exception:
+        # Close any opened log streams.
+        util.map_defer_exceptions(lambda ls: ls.close(), log_streams)
+        raise
+
+    add_step_text = _multiplex('add_step_text')
+    add_step_summary_text = _multiplex('add_step_summary_text')
+    add_step_link = _multiplex('add_step_link')
+    reset_subannotation_state = _multiplex('reset_subannotation_state')
+    set_step_status = _multiplex('set_step_status')
+    set_build_property = _multiplex('set_build_property')
+    trigger = _multiplex('trigger')
+
+  def new_step_stream(self, step_config):
+    return self.StepStream(
+        *(se.new_step_stream(step_config)
+          for se in self._engines))
+
+  def open(self):
+    for se in self._engines:
+      se.open()
+
+  def close(self):
+    util.map_defer_exceptions(lambda se: se.close(), self._engines)
+
+  def append_stream_engine(self, se):
+    assert isinstance(se, StreamEngine)
+    self._engines.append(se)
+
+
+def _noop(*_args, **_kwargs):
   pass
 
 class NoopStreamEngine(StreamEngine):
   class Stream(StreamEngine.Stream):
     write_line = _noop
     close = _noop
 
   class StepStream(Stream):
-    def new_log_stream(self, log_name):
+    def new_log_stream(self, _log_name):
       return NoopStreamEngine.Stream()
     add_step_text = _noop
     add_step_summary_text = _noop
     add_step_link = _noop
     reset_subannotation_state = _noop
     set_step_status = _noop
     set_build_property = _noop
     trigger = _noop
 
   def new_step_stream(self, step_config):
     return self.StepStream()
 
 
 class StreamEngineInvariants(StreamEngine):
   """Checks that the users are using a StreamEngine hygenically.
 
   Multiply with actually functional StreamEngines so you don't have to check
   these all over the place.
   """
   def __init__(self):
     self._streams = set()
 
+  @classmethod
+  def wrap(cls, other):
+    """Returns (ProductStreamEngine): A product applying invariants to "other".
+    """
+    return ProductStreamEngine(cls(), other)
+
   class StepStream(StreamEngine.StepStream):
     def __init__(self, engine, step_name):
+      super(StreamEngineInvariants.StepStream, self).__init__()
       self._engine = engine
       self._step_name = step_name
       self._open = True
       self._logs = {}
       self._status = 'SUCCESS'
 
     def write_line(self, line):
       assert '\n' not in line
       assert self._open
 
@@ skipping 100 matching lines @@
     if self._current_step != step_name:
       self.output_root_annotation('STEP_CURSOR', step_name)
       self._current_step = step_name
     if step_name not in self._opened:
       self.output_current_time()
       self.output_root_annotation('STEP_STARTED')
       self._opened.add(step_name)
 
   class StepStream(StreamEngine.StepStream):
     def __init__(self, engine, outstream, step_name):
-      super(StreamEngine.StepStream, self).__init__()
+      super(AnnotatorStreamEngine.StepStream, self).__init__()
 
       self._engine = engine
       self._outstream = outstream
       self._step_name = step_name
 
     def basic_write(self, line):
       self._engine._step_cursor(self._step_name)
       self._outstream.write(line)
 
     def close(self):
@@ skipping 89 matching lines @@
   string (which means it might not be valid ascii), we decode the string with
   the 'replace' error mode, which replaces invalid characters with a suitable
   replacement character.
   """
   try:
     return str(s)
   except UnicodeEncodeError:
     return s.encode('utf-8', 'replace')
   except UnicodeDecodeError:
     return s.decode('utf-8', 'replace')