Updates BatchWorker to report failure to shutdown progress reporter to service properly.

A progress reporter shutdown failure may indicate a failure to send the last
dynamic split response to the service. Hence work item must be marked as failed when such failures occur.

Added tests to cover this.

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

Author: chamikaramj

google/cloud/dataflow/worker/batchworker.py

@@ -435,24 +435,35 @@ def do_work(self, work_item, deferred_exception_details=None):
                       'was found: %s', work_item, deferred_exception_details)
         return
 
+    exception_details = None
     try:
       progress_reporter.start_reporting_progress()
       work_executor.execute(work_item.map_task)
     except Exception:  # pylint: disable=broad-except
-      progress_reporter.stop_reporting_progress()
       exception_details = traceback.format_exc()
-      logging.error('Exception: %s', exception_details, exc_info=True)
-      # Completed with errors means failed.
+      logging.error('An exception was raised when trying to execute the '
+                    'work item %s : %s',
+                    work_item,
+                    exception_details, exc_info=True)
+    finally:
+      try:
+        progress_reporter.stop_reporting_progress()
+      except Exception:  # pylint: disable=broad-except
+        logging.error('An exception was raised when trying to stop the '
+                      'progress reporter : %s',
+                      traceback.format_exc(), exc_info=True)
+        # If 'exception_details' was already set, we were already going to
+        # mark this work item as failed. Hence only logging this error and
+        # reporting the original error.
+        if exception_details is None:
+          # This will be reported to the service and work item will be marked as
+          # failed.
+          exception_details = traceback.format_exc()
+
       with work_item.lock:
-        self.report_completion_status(work_item,
-                                      progress_reporter,
+        self.report_completion_status(work_item, progress_reporter,
                                       exception_details=exception_details)
         work_item.done = True
-    else:
-      progress_reporter.stop_reporting_progress()
-      with work_item.lock:
-        self.report_completion_status(work_item, progress_reporter)
-        work_item.done = True
 
   def status_server(self):
     """Executes the serving loop for the status server."""

google/cloud/dataflow/worker/batchworker_test.py

@@ -108,15 +108,10 @@ def test_worker_fails_for_deferred_exceptions(
     assert not mock_start.called
     assert not mock_execute.called
 
-  @patch.object(executor.MapTaskExecutor, 'execute')
-  @patch.object(batchworker.ProgressReporter, 'start_reporting_progress')
-  @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress')
-  @patch.object(batchworker.ProgressReporter, 'report_status')
-  def test_worker_sends_completion_in_case_of_a_failure(
-      self, mock_report_status, mock_stop, mock_start, mock_execute):
+  def _run_send_completion_test(self, mock_report_status, mock_stop, mock_start,
+                                mock_execute, expected_exception):
     worker = batchworker.BatchWorker(self.dummy_properties(), {})
     mock_work_item = mock.MagicMock()
-    mock_execute.side_effect = Exception('test_exception')
     worker.do_work(mock_work_item)
 
     class AnyStringWith(str):
@@ -125,12 +120,46 @@ def __eq__(self, other):
         return self in other
 
     mock_report_status.assert_called_with(
-        completed=True, exception_details=AnyStringWith('test_exception'))
+        completed=True,
+        exception_details=AnyStringWith(expected_exception))
 
     mock_start.assert_called_once_with()
     mock_execute.assert_called_once_with(mock.ANY)
     mock_stop.assert_called_once_with()
 
+  @patch.object(executor.MapTaskExecutor, 'execute')
+  @patch.object(batchworker.ProgressReporter, 'start_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'report_status')
+  def test_send_completion_execute_failure(self, mock_report_status, mock_stop,
+                                           mock_start, mock_execute):
+    mock_execute.side_effect = Exception('test_exception')
+    self._run_send_completion_test(mock_report_status, mock_stop, mock_start,
+                                   mock_execute, 'test_exception')
+
+  @patch.object(executor.MapTaskExecutor, 'execute')
+  @patch.object(batchworker.ProgressReporter, 'start_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'report_status')
+  def test_send_completion_stop_progress_reporter_failure(self,
+                                                          mock_report_status,
+                                                          mock_stop, mock_start,
+                                                          mock_execute):
+    mock_stop.side_effect = Exception('test_exception')
+    self._run_send_completion_test(mock_report_status, mock_stop, mock_start,
+                                   mock_execute, 'test_exception')
+
+  @patch.object(executor.MapTaskExecutor, 'execute')
+  @patch.object(batchworker.ProgressReporter, 'start_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'stop_reporting_progress')
+  @patch.object(batchworker.ProgressReporter, 'report_status')
+  def test_send_completion_execute_and_stop_progress_reporter_failure(
+      self, mock_report_status, mock_stop, mock_start, mock_execute):
+    mock_execute.side_effect = Exception('test_exception_1')
+    mock_stop.side_effect = Exception('test_exception_2')
+    self._run_send_completion_test(mock_report_status, mock_stop, mock_start,
+                                   mock_execute, 'test_exception_1')
+
 
 class ProgressReporterTest(unittest.TestCase):