Updates and simplifies logic related to progress reporting.

Adds a new class ProgressReporter that encapsulates logic related to progress reporting. Each progress reporter object will start a new thread. A progress reporter should be created for each work item processed by a worker. This simplifies logic since we can now cleanly start and stop progress reporting of a work item without having to worry about unnecessary progress reports being send by a progress reporting thread.

Updates ProgressReporter to shutdown cleanly and send any unsent split requests.

Significantly improves test coverage of BatchWorker by adding tests to a batchworker_test.py.

Updates apiclient so that its 'report_status' and 'lease_work' methods accept an object of type WorkerInfo that contains necessary information about the worker instead of the worker object preventing apiclient from accessing internal state of the worker.

Fixes a bug in CounterFactory. Updated it to be thread safe so that progress reporting thread and main thread can update the counters map simultaneously.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=119423075

Author: chamikaramj

google/cloud/dataflow/internal/apiclient.py

@@ -632,20 +632,20 @@ def __init__(self, worker, skip_get_credentials=False):
             get_credentials=(not skip_get_credentials)))
 
   @retry.with_exponential_backoff()  # Using retry defaults from utils/retry.py
-  def lease_work(self, worker):
+  def lease_work(self, worker_info, desired_lease_duration):
     """Leases a work item from the service."""
     work_request = dataflow.LeaseWorkItemRequest()
-    work_request.workerId = worker.worker_id
-    work_request.requestedLeaseDuration = worker.desired_lease_duration()
-    work_request.currentWorkerTime = worker.current_time
-    work_request.workerCapabilities.append(worker.worker_id)
-    for value in worker.capabilities:
+    work_request.workerId = worker_info.worker_id
+    work_request.requestedLeaseDuration = desired_lease_duration
+    work_request.currentWorkerTime = worker_info.formatted_current_time
+    work_request.workerCapabilities.append(worker_info.worker_id)
+    for value in worker_info.capabilities:
       work_request.workerCapabilities.append(value)
-    for value in worker.work_types:
+    for value in worker_info.work_types:
       work_request.workItemTypes.append(value)
     request = dataflow.DataflowProjectsJobsWorkItemsLeaseRequest()
-    request.jobId = worker.job_id
-    request.projectId = worker.project_id
+    request.jobId = worker_info.job_id
+    request.projectId = worker_info.project_id
     try:
       request.leaseWorkItemRequest = work_request
     except AttributeError:
@@ -656,7 +656,8 @@ def lease_work(self, worker):
     return response
 
   def report_status(self,
-                    worker,
+                    worker_info,
+                    desired_lease_duration,
                     work_item,
                     completed,
                     progress,
@@ -670,7 +671,12 @@ def report_status(self,
     work item.
 
     Args:
-      worker: The Worker instance executing the work item.
+      worker_info: A batchworker.BatchWorkerInfo that contains
+        information about the Worker instance executing the work
+        item.
+      desired_lease_duration: The duration for which the worker would like to
+        extend the lease of the work item. Should be in seconds formatted as a
+        string.
       work_item: The work item for which to report status.
       completed: True if there is no further work to be done on this work item
         either because it succeeded or because it failed. False if this is a
@@ -695,7 +701,7 @@ def report_status(self,
     work_item_status.completed = completed
 
     if not completed:
-      work_item_status.requestedLeaseDuration = worker.desired_lease_duration()
+      work_item_status.requestedLeaseDuration = desired_lease_duration
 
     if progress is not None:
       work_item_progress = dataflow.ApproximateProgress()
@@ -742,13 +748,13 @@ def report_status(self,
       append_counter(work_item_status, counter, tentative=not completed)
 
     report_request = dataflow.ReportWorkItemStatusRequest()
-    report_request.currentWorkerTime = worker.current_time
-    report_request.workerId = worker.worker_id
+    report_request.currentWorkerTime = worker_info.formatted_current_time
+    report_request.workerId = worker_info.worker_id
     report_request.workItemStatuses.append(work_item_status)
 
     request = dataflow.DataflowProjectsJobsWorkItemsReportStatusRequest()
-    request.jobId = worker.job_id
-    request.projectId = worker.project_id
+    request.jobId = worker_info.job_id
+    request.projectId = worker_info.project_id
     try:
       request.reportWorkItemStatusRequest = report_request
     except AttributeError:

google/cloud/dataflow/utils/counters.py

@@ -17,6 +17,8 @@
 
 """Counters collect the progress of the Worker for reporting to the service."""
 
+import threading
+
 
 class Counter(object):
   """A counter aggregates a series of values.
@@ -109,6 +111,9 @@ class CounterFactory(object):
   def __init__(self):
     self.counters = {}
 
+    # Lock to be acquired when accessing the counters map.
+    self._lock = threading.Lock()
+
   def get_counter(self, name, aggregation_kind):
     """Returns a counter with the requested name.
 
@@ -122,13 +127,14 @@ def get_counter(self, name, aggregation_kind):
     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
+    with self._lock:
+      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.
@@ -141,16 +147,26 @@ def get_aggregator_counter(self, step_name, aggregator):
     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()
+    with self._lock:
+      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 get_counters(self):
+    """Returns the current set of counters.
+
+    Returns:
+      An iterable that contains the current set of counters. To make sure that
+      multiple threads can iterate over the set of counters, we return a new
+      iterable here. Note that the actual set of counters may get modified after
+      this method returns hence the returned iterable may be stale.
+    """
+    with self._lock:
+      return self.counters.values()