Emit CURRENT_TIMESTAMP annotation

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
 
 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."""
       for actual_line in string.splitlines() or ['']: # preserve empty lines
@@ skipping 30 matching lines @@
     def set_nest_level(self, nest_level):
       raise NotImplementedError()
 
     def set_build_property(self, key, value):
       raise NotImplementedError()
 
     def trigger(self, trigger_spec):
       raise NotImplementedError()
 
   def new_step_stream(self, step_name, allow_subannotations=False):
-    """Craete a new StepStream in this engine.
+    """Creates a new StepStream in this engine.
 
     The step will be considered started at the moment this method is called.
 
     TODO(luqui): allow_subannotations is a bit of a hack, whether to allow
     annotations that this step emits through to the annotator (True), or
     guard them by prefixing them with ! (False).  The proper way to do this
     is to implement an annotations parser that converts to StreamEngine calls;
     i.e. parse -> re-emit.
     """
 
     raise NotImplementedError()
 
+  def open(self):
+    pass
+
+  def close(self):
+    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):
   def __init__(self, engine_a, engine_b):
     self._engine_a = engine_a
     self._engine_b = engine_b
 
   class Stream(StreamEngine.Stream):
@@ skipping 28 matching lines @@
     set_step_status = _void_product('set_step_status')
     set_nest_level = _void_product('set_nest_level')
     set_build_property = _void_product('set_build_property')
     trigger = _void_product('trigger')
 
   def new_step_stream(self, step_name, allow_subannotations=False):
     return self.StepStream(
         self._engine_a.new_step_stream(step_name, allow_subannotations),
         self._engine_b.new_step_stream(step_name, allow_subannotations))
 
+  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):
   pass
 
 class NoopStreamEngine(StreamEngine):
   class Stream(StreamEngine.Stream):
     write_line = _noop
     close = _noop
 
   class StepStream(Stream):
@@ skipping 92 matching lines @@
       assert self._open
       self._open = False
 
   def new_step_stream(self, step_name, allow_subannotations=False):
     assert step_name not in self._streams, 'Step %s already exists' % step_name
     self._streams.add(step_name)
     return self.StepStream(self, step_name)
 
 
 class AnnotationStepStream(StreamEngine.StepStream):
+  def __init__(self, emit_timestamps=False, time_fn=None):
+    self.emit_timestamps = emit_timestamps
+    self.time_fn = time_fn or time.time
+
   def basic_write(self, line):
     raise NotImplementedError()
 
   def output_annotation(self, *args):
-    self.basic_write('@@@' + '@'.join(args) + '@@@\n')
+    self.basic_write('@@@' + '@'.join(map(str, args)) + '@@@\n')
 
   def write_line(self, line):
     if line.startswith('@@@'):
       self.basic_write('!' + line + '\n')
     else:
       self.basic_write(line + '\n')
 
   def close(self):
+    if self.emit_timestamps:
+      self.output_annotation('CURRENT_TIMESTAMP', self.time_fn())
     self.output_annotation('STEP_CLOSED')
 
   def new_log_stream(self, log_name):
     return self.StepLogStream(self, log_name)
 
   def add_step_text(self, text):
     self.output_annotation('STEP_TEXT', text)
 
   def add_step_summary_text(self, text):
     self.output_annotation('STEP_SUMMARY_TEXT', text)
@@ skipping 28 matching lines @@
       self._log_name = log_name.replace('/', '/')
 
     def write_line(self, line):
       self._step_stream.output_annotation('STEP_LOG_LINE', self._log_name, line)
 
     def close(self):
       self._step_stream.output_annotation('STEP_LOG_END', self._log_name)
 
 
 class AnnotatorStreamEngine(StreamEngine):
-  def __init__(self, outstream):
+  def __init__(self, outstream, emit_timestamps=False, time_fn=None):
     self._current_step = None
     self._opened = set()
     self._outstream = outstream
+    self.emit_timestamps = emit_timestamps
+    self.time_fn = time_fn or time.time
+
+  def open(self):
+    super(AnnotatorStreamEngine, self).open()
+    if self.emit_timestamps:
+      self.output_current_time()
     self.output_annotation('HONOR_ZERO_RETURN_CODE')
 
+  def close(self):
+    super(AnnotatorStreamEngine, self).close()
+    if self.emit_timestamps:
+      self.output_current_time()
+
+  def output_current_time(self):
+    """Prints CURRENT_TIMESTAMP annotation with current time."""
+    self.output_annotation('CURRENT_TIMESTAMP', self.time_fn())
+
   def output_annotation(self, *args):
     # Flush the stream before & after engine annotations, because they can
     # change which step we are talking about and this matters to buildbot.
     self._outstream.flush()
-    self._outstream.write('@@@' + '@'.join(args) + '@@@\n')
+    self._outstream.write('@@@' + '@'.join(map(str, args)) + '@@@\n')
     self._outstream.flush()
 
   def _step_cursor(self, name):
     if self._current_step != name:
       self.output_annotation('STEP_CURSOR', name)
       self._current_step = name
     if name not in self._opened:
+      if self.emit_timestamps:
+        self.output_current_time()
       self.output_annotation('STEP_STARTED')
       self._opened.add(name)
 
   class StepStream(AnnotationStepStream):
     def __init__(self, engine, step_name):
+      AnnotationStepStream.__init__(
+          self, emit_timestamps=engine.emit_timestamps, time_fn=engine.time_fn)
       self._engine = engine
       self._step_name = step_name
 
     def basic_write(self, line):
       self._engine._step_cursor(self._step_name)
       self._engine._outstream.write(line)
 
   class AllowSubannotationsStepStream(StepStream):
     def write_line(self, line):
       self.basic_write(line + '\n')
@@ skipping 17 matching lines @@
   """A StepStream that is not tied to any engine, and emits assuming it has the
   cursor.
 
   This is used for capturing the annotations in the engine.
   """
   def __init__(self, outstream):
     self._outstream = outstream
 
   def basic_write(self, line):
     self._outstream.write(line)
-