Use a CounterFactory to create counters

With this change, steps with the same name use the same counter.
Previously, we would send two counters with the same name to the
Dataflow service, which would reject the update.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=118901263

Author: gildea

google/cloud/dataflow/internal/apiclient.py

@@ -738,9 +738,8 @@ def report_status(self,
       work_item_status.errors.append(status)
 
     # Look through the work item for metrics to send.
-    for op in work_item.map_task.executed_operations:
-      for counter in op.itercounters():
-        append_counter(work_item_status, counter, tentative=not completed)
+    for counter in work_item.map_task.itercounters():
+      append_counter(work_item_status, counter, tentative=not completed)
 
     report_request = dataflow.ReportWorkItemStatusRequest()
     report_request.currentWorkerTime = worker.current_time

google/cloud/dataflow/runners/common.py

@@ -24,7 +24,6 @@
 from google.cloud.dataflow.transforms.window import TimestampedValue
 from google.cloud.dataflow.transforms.window import WindowedValue
 from google.cloud.dataflow.transforms.window import WindowFn
-from google.cloud.dataflow.utils import counters
 
 
 class FakeLogger(object):
@@ -172,17 +171,11 @@ class DoFnState(object):
   """Keeps track of state that DoFns want, currently, user counters.
   """
 
-  def __init__(self):
+  def __init__(self, counter_factory):
     self.step_name = ''
-    self._user_counters = {}
+    self._counter_factory = counter_factory
 
   def counter_for(self, aggregator):
     """Looks up the counter for this aggregator, creating one if necessary."""
-    if aggregator not in self._user_counters:
-      self._user_counters[aggregator] = counters.AggregatorCounter(
-          self.step_name, aggregator)
-    return self._user_counters[aggregator]
-
-  def itercounters(self):
-    """Returns an iterable of Counters (to be sent to the service)."""
-    return self._user_counters.values()
+    return self._counter_factory.get_aggregator_counter(
+        self.step_name, aggregator)

google/cloud/dataflow/runners/direct_runner.py

@@ -39,6 +39,7 @@
 from google.cloud.dataflow.typehints.typecheck import OutputCheckWrapperDoFn
 from google.cloud.dataflow.typehints.typecheck import TypeCheckError
 from google.cloud.dataflow.typehints.typecheck import TypeCheckWrapperDoFn
+from google.cloud.dataflow.utils import counters
 from google.cloud.dataflow.utils.options import TypeOptions
 
 
@@ -52,6 +53,7 @@ class DirectPipelineRunner(PipelineRunner):
   def __init__(self, cache=None):
     # Cache of values computed while the runner executes a pipeline.
     self._cache = cache if cache is not None else PValueCache()
+    self._counter_factory = counters.CounterFactory()
 
   def get_pvalue(self, pvalue):
     """Gets the PValue's computed value from the runner's cache."""
@@ -78,7 +80,8 @@ def func_wrapper(self, pvalue, *args, **kwargs):
   def run_ParDo(self, transform_node):
     transform = transform_node.transform
     # TODO(gildea): what is the appropriate object to attach the state to?
-    context = DoFnProcessContext(label=transform.label, state=DoFnState())
+    context = DoFnProcessContext(label=transform.label,
+                                 state=DoFnState(self._counter_factory))
 
     # Construct the list of values from side-input PCollections that we'll
     # substitute into the arguments for DoFn methods.

google/cloud/dataflow/utils/counters.py

@@ -97,9 +97,60 @@ def _str_internal(self):
 
 
 class AggregatorCounter(Counter):
-  """A Counter that represents a step-specific instance of an Aggregator."""
+  """A Counter that represents a step-specific instance of an Aggregator.
 
-  def __init__(self, step_name, aggregator):
-    super(AggregatorCounter, self).__init__(
-        'user-%s-%s' % (step_name, aggregator.name),
-        aggregator.aggregation_kind)
+  Do not create directly, call CounterFactory.get_aggregator_counter instead.
+  """
+
+
+class CounterFactory(object):
+  """Keeps track of unique counters."""
+
+  def __init__(self):
+    self.counters = {}
+
+  def get_counter(self, name, aggregation_kind):
+    """Returns a counter with the requested name.
+
+    Passing in the same name will return the same counter; the
+    aggregation_kind must agree.
+
+    Args:
+      name: the name of this counter.  Typically has three parts:
+        "step-output-counter".
+      aggregation_kind: one of the kinds defined by this class.
+    Returns:
+      A new or existing counter with the requested name.
+    """
+    counter = self.counters.get(name, None)
+    if counter:
+      assert counter.aggregation_kind == aggregation_kind
+    else:
+      counter = Counter(name, aggregation_kind)
+      self.counters[name] = counter
+    return counter
+
+  def get_aggregator_counter(self, step_name, aggregator):
+    """Returns an AggregationCounter for this step's aggregator.
+
+    Passing in the same values will return the same counter.
+
+    Args:
+      step_name: the name of this step.
+      aggregator: an Aggregator object.
+    Returns:
+      A new or existing counter.
+    """
+    name = 'user-%s-%s' % (step_name, aggregator.name)
+    aggregation_kind = aggregator.aggregation_kind
+    counter = self.counters.get(name, None)
+    if counter:
+      assert isinstance(counter, AggregatorCounter)
+      assert counter.aggregation_kind == aggregation_kind
+    else:
+      counter = AggregatorCounter(name, aggregation_kind)
+      self.counters[name] = counter
+    return counter
+
+  def itercounters(self):
+    return self.counters.itervalues()

google/cloud/dataflow/worker/executor.pxd

@@ -16,6 +16,7 @@ cimport cython
 
 cdef class Operation(object):
   cdef public spec
+  cdef public counter_factory
   cdef public list receivers
   cdef public list counters
   cdef readonly bint debug_logging_enabled