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 136 matching lines @@
   """
 
   # The name of the annotation stream.
   ANNOTATION_NAME = 'annotations'
 
   # The default amount of time in between anotation pushes.
   DEFAULT_UPDATE_INTERVAL = datetime.timedelta(seconds=30)
 
 
   def __init__(self, client=None, streamserver_uri=None, name_base=None,
-               ignore_triggers=False, environment=None, update_interval=None):
+               dump_path=None, ignore_triggers=False, environment=None,
+               update_interval=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.
       update_interval (datetime.timedelta or None): The interval of time between
           annotation data pushes. If None, DEFAULT_UPDATE_INTERVAL will be
           used.
     """
 
     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.real()
     self._update_interval = update_interval or self.DEFAULT_UPDATE_INTERVAL
 
     self._astate = None
 
     self._annotation_stream = None
     self._annotation_monitor = None
     self._streams = collections.OrderedDict()
 
@@ skipping 221 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._notify_annotation_changed()
 
     # 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 there is no last step data, this will write an empty file, which is
+    # still a valid protobuf.
+    if self._dump_path and last_step_data:
+      with open(self._dump_path, 'wb') as fd:
+          fd.write(last_step_data)
+
   def _notify_annotation_changed(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 122 matching lines @@
       # Note that the timer may also suggest that we flush immediately if we're
       # already past our last flush interval.
       now = self._env.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(self._env.now)
+      return self._last_flush_data

[martiniss, 2016/10/12 23:01:25]

this is the data we last flushed. We always flush the whole annotation data?

[dnj, 2016/10/12 23:15:02]

Yep.

 
   @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