Make element iterators observable

Mix class ObservableMixin into various values iterators, so that we
have a hook to estimate the size of the elements we iterate over.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=119993968

Author: gildea

google/cloud/dataflow/transforms/trigger.py

@@ -22,6 +22,7 @@
 import collections
 import copy
 
+from google.cloud.dataflow.coders import observable
 from google.cloud.dataflow.transforms import combiners
 from google.cloud.dataflow.transforms import core
 from google.cloud.dataflow.transforms.timeutil import MAX_TIMESTAMP
@@ -721,9 +722,12 @@ def process_elements(self, state, windowed_values, unused_output_watermark):
     if isinstance(windowed_values, list):
       unwindowed = [wv.value for wv in windowed_values]
     else:
-      class UnwindowedValues(object):
+      class UnwindowedValues(observable.ObservableMixin):
         def __iter__(self):
-          return (wv.value for wv in windowed_values)
+          for wv in windowed_values:
+            unwindowed_value = wv.value
+            self.notify_observers(unwindowed_value)
+            yield unwindowed_value
         def __repr__(self):
           return '<UnwindowedValues of %s>' % windowed_values
       unwindowed = UnwindowedValues()

google/cloud/dataflow/worker/shuffle.py

@@ -38,6 +38,7 @@
 import logging
 import struct
 
+from google.cloud.dataflow.coders import observable
 from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.io import range_trackers
 
@@ -246,7 +247,7 @@ def clone(self, start_position, end_position, key):
             self.iterable.reader, start_position, end_position, key))
 
 
-class ShuffleKeyValuesIterable(object):
+class ShuffleKeyValuesIterable(observable.ObservableMixin):
   """An iterable over all values associated with a key.
 
   The class supports reiteration over the values by cloning the underlying
@@ -257,6 +258,7 @@ class ShuffleKeyValuesIterable(object):
 
   def __init__(self, entries_iterator, key, value_coder,
                start_position, end_position=''):
+    super(ShuffleKeyValuesIterable, self).__init__()
     self.key = key
     self.value_coder = value_coder
     self.start_position = start_position
@@ -290,7 +292,9 @@ def values_iterator(self):
         self.end_position = entry.position
         self.entries_iterator.push_back(entry)
         break
-      yield self.value_coder.decode(entry.value)
+      decoded_value = self.value_coder.decode(entry.value)
+      self.notify_observers(entry.value, is_encoded=True)
+      yield decoded_value
 
   def __str__(self):
     return '<%s>' % self._str_internal()

google/cloud/dataflow/worker/windmillio.py

@@ -18,6 +18,7 @@
 
 from __future__ import absolute_import
 
+from google.cloud.dataflow.coders import observable
 from google.cloud.dataflow.io import coders
 from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.io import pubsub
@@ -216,10 +217,11 @@ def __repr__(self):
                                                  self.state_family)
 
 
-class KeyedWorkItem(object):
+class KeyedWorkItem(observable.ObservableMixin):
   """Keyed work item used by a StreamingGroupAlsoByWindowsOperation."""
 
   def __init__(self, work_item, coder):
+    super(KeyedWorkItem, self).__init__()
     self.work_item = work_item
     self.coder = coder
     self.key_coder = coder.key_coder()
@@ -236,7 +238,9 @@ def __init__(self, work_item, coder):
   def elements(self):
     for bundle in self.work_item.message_bundles:
       for message in bundle.messages:
-        yield self.wv_coder.decode(message.data)
+        element = self.wv_coder.decode(message.data)
+        self.notify_observers(message.data, is_encoded=True)
+        yield element
 
   def timers(self):
     if self.work_item.timers:
@@ -252,7 +256,7 @@ def timers(self):
             state_family=timer_item.state_family)
 
   def __repr__(self):
-    return 'KeyedWorkItem(%r)' % self.key
+    return '<%s %s>' % (self.__class__.__name__, self.key)
 
 
 class WindowingWindmillSource(iobase.Source):