Repository: google/turbinia

Issue:
Adding errors to TurbiniaTaskResult causes it to be unpickleable
[INFO]  Took 0.27 sec
Traceback (most recent call last):
  File "/home/turbinia/src/turbinia/turbiniactl", line 175, in <module>
    worker.listen()
  File "/home/turbinia/turbinia-env/local/lib/python2.7/site-packages/psq/worker.py", line 60, in listen
    self.run_task(task)
  File "/home/turbinia/turbinia-env/local/lib/python2.7/site-packages/psq/worker.py", line 69, in run_task
    task.execute(self.queue)
  File "/home/turbinia/turbinia-env/local/lib/python2.7/site-packages/psq/task.py", line 96, in execute
    queue.storage.put_task(self)
  File "/home/turbinia/turbinia-env/local/lib/python2.7/site-packages/psq/datastore_storage.py", line 61, in put_task
    entity['data'] = dumps(task)
  File "/home/turbinia/turbinia-env/local/lib/python2.7/site-packages/google/cloud/client.py", line 134, in __getstate__
    'Clients have non-trivial state that is local and unpickleable.',
pickle.PicklingError: Pickling client objects is explicitly not supported.
Clients have non-trivial state that is local and unpickleable.


Discussion:

Resolving patch:
diff --git a/turbinia/task_manager.py b/turbinia/task_manager.py
--- a/turbinia/task_manager.py
+++ b/turbinia/task_manager.py
@@ -56,22 +56,8 @@ def task_runner(obj, *args, **kwargs):
 
   Returns:
     Output from TurbiniaTask (should be TurbiniaTaskResult).
-
-  Raises:
-    Re-raises exceptions that are thrown from the task.
   """
-  # TODO(aarontp): Add proper error checks/handling
-  res = None
-  try:
-    res = obj.run(*args, **kwargs)
-  except Exception as e:
-    # TODO(aarontp): Create synthetic TurbiniaTaskResult upon failure to
-    # propogate errors to the Task Manager
-    logging.warning(
-        u'Exception thrown from Task: {0:s}'.format(traceback.format_exc()))
-    raise e
-
-  return res
+  return obj.run_wrapper(*args, **kwargs)
 
 
 class BaseTaskManager(object):
@@ -233,8 +219,8 @@ def _finalize_result(self, task_result):
       task_result: The TurbiniaTaskResult object
     """
     if not task_result.successful:
-      log.error(u'Task {0:s} from {1:s} was not successful'.format(
-          task_result.task_name, task_result.worker_name))
+      log.error(u'Task {0:s} from {1:s} failed: [{2:s}]'.format(
+          task_result.task_name, task_result.worker_name, task_result.status))
     else:
       log.info(
           u'Task {0:s} from {1:s} executed with status [{2:s}]'.format(
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -22,6 +22,7 @@
 import os
 import platform
 import time
+import traceback
 import uuid
 
 from turbinia import config
@@ -160,7 +161,7 @@ def save_local_file(self, file_):
       if writer.name != 'LocalOutputWriter':
         writer.write(file_)
 
-  def set_error(self, error, traceback):
+  def set_error(self, error, traceback_):
     """Add error and traceback.
 
     Args:
@@ -168,7 +169,7 @@ def set_error(self, error, traceback):
         traceback: Traceback of the error.
     """
     self.error['error'] = error
-    self.error['traceback'] = traceback
+    self.error['traceback'] = traceback_
 
 
 class TurbiniaTask(object):
@@ -220,12 +221,37 @@ def setup(self, evidence):
     evidence.preprocess()
     return self.result
 
+  def run_wrapper(self, evidence):
+    """Wrapper to manage TurbiniaTaskResults and exception handling.
 
-  def run(self, evidence):
+    This wrapper should be called to invoke the run() methods so it can handle
+    the management of TurbiniaTaskResults and the exception handling.  Otherwise
+    details from exceptions in the worker cannot be propogated back to the
+    Turbinia TaskManager.
+
+    Args:
+      evidence: Evidence object
+
+    Returns:
+      A TurbiniaTaskResult object
+    """
+    result = self.setup(evidence)
+    try:
+      result = self.run(evidence, result)
+    # pylint: disable=broad-except
+    except Exception as e:
+      msg = 'Task failed with exception: [{0!s}]'.format(e)
+      result.close(success=False, status=msg)
+      result.set_error(e.message, traceback.format_exc())
+
+    return result
+
+  def run(self, evidence, result):
     """Entry point to execute the task.
 
     Args:
       evidence: Evidence object.
+      result: A TurbiniaTaskResult object to place task results into.
 
     Returns:
         TurbiniaTaskResult object.
diff --git a/turbinia/workers/plaso.py b/turbinia/workers/plaso.py
--- a/turbinia/workers/plaso.py
+++ b/turbinia/workers/plaso.py
@@ -23,16 +23,16 @@
 class PlasoTask(TurbiniaTask):
   """Task to run Plaso (log2timeline)."""
 
-  def run(self, evidence):
+  def run(self, evidence, result):
     """Task that process data with Plaso.
 
     Args:
         evidence: Path to data to process.
+        result: A TurbiniaTaskResult object to place task results into.
 
     Returns:
         TurbiniaTaskResult object.
     """
-    result = self.setup(evidence)
     plaso_result = PlasoFile()
 
     plaso_file = os.path.join(self.output_dir, u'{0:s}.plaso'.format(self.id))


Test patch: