Annotation protobuf code can now dump final.

recipe_engine/stream_logdog.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.
 
 """stream.StreamEngine implementation for LogDog, using Milo annotation
 protobuf.
 """
 
 import collections
 import contextlib
@@ skipping 250 matching lines @@
 
     def trigger(self, trigger_spec):
       if self._engine._ignore_triggers:
         return
       raise NotImplementedError(
           'Stream-based triggering is not supported for LogDog. Please use '
           'a recipe module (e.g., buildbucket) directly for build scheduling.')
 
 
   def __init__(self, client=None, streamserver_uri=None, name_base=None,
-               ignore_triggers=False, environment=None):
+               dump_path=None, ignore_triggers=False, environment=None):
     """Initializes a new LogDog/Annotation StreamEngine.
 
     Args:
       client (libs.logdog.stream.StreamClient or None): the LogDog stream client
           to use. If this is None, a new StreamClient will be instantiated when
           this StreamEngine is opened.
       streamserver_uri (str or None): The LogDog Butler stream server URI. See
           LogDog client library docs for details on supported protocols and
           format. This will only be used when "client" is None. If this is also
           None, a StreamClient will be created through probing.
       name_base (str or None): The default stream name prefix that will be added
           to generated LogDog stream names. If None, no prefix will be applied.
+      dump_path (str or None): If provided, a filesystem path where the final
+          recipe annotation protobuf binary will be dumped.
       ignore_triggers (bool): Triggers are not supported in LogDog annotation
           streams. If True, attempts to trigger will be silently ignored. If
           False, they will cause a NotImplementedError to be raised.
       environment (_Environment or None): The _Environment instance to use for
           operations. This will be None at production, but can be overridden
           here for testing.
     """
 
     self._client = client
     self._streamserver_uri = streamserver_uri
     self._name_base = _StreamName(name_base)
+    self._dump_path = dump_path
     self._ignore_triggers = ignore_triggers
     self._env = environment or _Environment.probe()
 
     self._astate = None
 
     self._annotation_stream = None
     self._annotation_monitor = None
     self._streams = collections.OrderedDict()
 
   def new_step_stream(self, step_config):
@@ skipping 51 matching lines @@
     # Shut down any outstanding streams that may not have been closed for
     # whatever reason.
     for s in reversed(self._streams.values()):
       s.close()
 
     # Close out our root Step. Manually check annotation state afterwards.
     self._close_step(self._astate.base)
     self._check()
 
     # Shut down our annotation monitor and close our annotation stream.
-    self._annotation_monitor.flush_and_join()
+    last_step_data = self._annotation_monitor.flush_and_join()
     self._annotation_stream.close()
 
     # Clear our client and state. We are now closed.
     self._streams.clear()
     self._client = None
     self._astate = None
 
+    # If requested, write out the last step data.
+    if self._dump_path and last_step_data:
+      with open(self._dump_path, 'wb') as fd:

[nodir, 2016/09/28 23:02:22]

please swap these two lines, so in the worst case we have an empty file (which
is still a valid protobuf)

[dnj, 2016/09/28 23:04:11]

Done.

+        fd.write(last_step_data)
+
   def _check(self):
     if self._astate is None:
       return
 
     step_data = self._astate.check()
     if step_data is not None:
       self._annotation_monitor.signal_update(step_data)
 
   def _set_timestamp(self, dst, dt=None):
     """Populates a timestamp_pb2.Timestamp, dst, with a datetime.
@@ skipping 120 matching lines @@
       # If not, make sure our timer is running so it will eventually be flushed.
       # Note that the timer may also suggest that we flush immediately.
       now = datetime.datetime.now()
       self._current_data = step_data
       if structural or self._set_flush_timer_locked(now):
         # We should flush immediately.
         self._flush_now_locked(now)
 
   def flush_and_join(self):
     """Flushes any remaining updates and blocks until the monitor is complete.
+
+    Returns (pb.Step): The final Step protobuf, or None if no step data was
+        sent.
     """
     # Mark that we're finished and signal our event.
     with self._lock:
       self._flush_now_locked(datetime.datetime.now())
+      return self._last_flush_data
 
   @property
   def latest(self):
     with self._lock:
       return self._last_flush_data
 
   def _flush_now_locked(self, now):
     # Clear any current flush timer.
     self._clear_flush_timer_locked()
 
@@ skipping 206 matching lines @@
     return v.replace('/', '_')
 
   @staticmethod
   def _normalize(v):
     return libs.logdog.streamname.normalize(v, prefix='s_')
 
   def __str__(self):
     if not self._base:
       raise ValueError('Cannot generate string from empty StreamName.')
     return self._base