Optimize logging context

As this is called per-element, it's important to minimize logging
overhead here.  Instead, reduce enter/exit to a single thread-local
access, and puts the overwriting into the (already expensive) logging
itself.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117290642

Author: robertwb

google/cloud/dataflow/worker/logger.py

@@ -23,45 +23,32 @@
 # Per-thread worker information. This is used only for logging to set
 # context information that changes while work items get executed:
 # work_item_id, step_name, stage_name.
-per_thread_worker_data = threading.local()
+class _PerThreadWorkerData(threading.local):
+
+  def __init__(self):
+    super(_PerThreadWorkerData, self).__init__()
+    self.stack = []
+
+  def get_data(self):
+    all_data = {}
+    for datum in self.stack:
+      all_data.update(datum)
+    return all_data
+
+per_thread_worker_data = _PerThreadWorkerData()
 
 
 class PerThreadLoggingContext(object):
   """A context manager to add per thread attributes."""
-  _instance = dict()
-
-  def __new__(cls, *args, **kwargs):
-    # TODO(robertwb): make the class generic, this is special-cased to save
-    # time on the DoFn
-    if not args and len(kwargs) == 1 and 'step_name' in kwargs:
-      k = kwargs['step_name']
-      if k not in cls._instance:
-        cls._instance[k] = super(PerThreadLoggingContext, cls).__new__(
-            cls, *args, **kwargs)
-      return cls._instance[k]
-    else:
-      return super(PerThreadLoggingContext, cls).__new__(cls, *args, **kwargs)
 
-  def __init__(self, *args, **kwargs):
-    if args:
-      raise ValueError(
-          'PerThreadLoggingContext expects only keyword arguments.')
+  def __init__(self, **kwargs):
     self.kwargs = kwargs
-    self.previous = {}
 
   def __enter__(self):
-    for key in self.kwargs:
-      if hasattr(per_thread_worker_data, key):
-        self.previous[key] = getattr(per_thread_worker_data, key)
-      setattr(per_thread_worker_data, key, self.kwargs[key])
-    return self
+    per_thread_worker_data.stack.append(self.kwargs)
 
   def __exit__(self, exn_type, exn_value, exn_traceback):
-    for key in self.kwargs:
-      if key in self.previous:
-        setattr(per_thread_worker_data, key, self.previous[key])
-      else:
-        delattr(per_thread_worker_data, key)
+    per_thread_worker_data.stack.pop()
 
 
 class JsonLogFormatter(logging.Formatter):
@@ -138,12 +125,13 @@ def format(self, record):
     # Stage, step and work item ID come from thread local storage since they
     # change with every new work item leased for execution. If there is no
     # work item ID then we make sure the step is undefined too.
-    if hasattr(per_thread_worker_data, 'work_item_id'):
-      output['work'] = getattr(per_thread_worker_data, 'work_item_id')
-    if hasattr(per_thread_worker_data, 'stage_name'):
-      output['stage'] = getattr(per_thread_worker_data, 'stage_name')
-    if hasattr(per_thread_worker_data, 'step_name'):
-      output['step'] = getattr(per_thread_worker_data, 'step_name')
+    data = per_thread_worker_data.get_data()
+    if 'work_item_id' in data:
+      output['work'] = data['work_item_id']
+    if 'stage_name' in data:
+      output['stage'] = data['stage_name']
+    if 'step_name' in data:
+      output['step'] = data['step_name']
     # All logging happens using the root logger. We will add the basename of the
     # file and the function name where the logging happened to make it easier
     # to identify who generated the record.

google/cloud/dataflow/worker/logger_test.py

@@ -26,45 +26,43 @@
 class PerThreadLoggingContextTest(unittest.TestCase):
 
   def thread_check_attribute(self, name):
-    self.assertFalse(hasattr(logger.per_thread_worker_data, name))
-    with logger.PerThreadLoggingContext(xyz='thread-value'):
+    self.assertFalse(name in logger.per_thread_worker_data.get_data())
+    with logger.PerThreadLoggingContext(**{name: 'thread-value'}):
       self.assertEqual(
-          getattr(logger.per_thread_worker_data, name), 'thread-value')
-    self.assertFalse(hasattr(logger.per_thread_worker_data, name))
+          logger.per_thread_worker_data.get_data()[name], 'thread-value')
+    self.assertFalse(name in logger.per_thread_worker_data.get_data())
 
   def test_no_positional_args(self):
-    with self.assertRaises(ValueError) as exn:
+    with self.assertRaises(TypeError):
       with logger.PerThreadLoggingContext('something'):
         pass
-    self.assertEqual(
-        exn.exception.message,
-        'PerThreadLoggingContext expects only keyword arguments.')
 
   def test_per_thread_attribute(self):
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
     with logger.PerThreadLoggingContext(xyz='value'):
-      self.assertEqual(logger.per_thread_worker_data.xyz, 'value')
+      self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value')
       thread = threading.Thread(
           target=self.thread_check_attribute, args=('xyz',))
       thread.start()
       thread.join()
-      self.assertEqual(logger.per_thread_worker_data.xyz, 'value')
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+      self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value')
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
 
   def test_set_when_undefined(self):
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
     with logger.PerThreadLoggingContext(xyz='value'):
-      self.assertEqual(logger.per_thread_worker_data.xyz, 'value')
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+      self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value')
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
 
   def test_set_when_already_defined(self):
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
     with logger.PerThreadLoggingContext(xyz='value'):
-      self.assertEqual(logger.per_thread_worker_data.xyz, 'value')
+      self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value')
       with logger.PerThreadLoggingContext(xyz='value2'):
-        self.assertEqual(logger.per_thread_worker_data.xyz, 'value2')
-      self.assertEqual(logger.per_thread_worker_data.xyz, 'value')
-    self.assertFalse(hasattr(logger.per_thread_worker_data, 'xyz'))
+        self.assertEqual(
+            logger.per_thread_worker_data.get_data()['xyz'], 'value2')
+      self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value')
+    self.assertFalse('xyz' in logger.per_thread_worker_data.get_data())
 
 
 class JsonLogFormatterTest(unittest.TestCase):
@@ -140,6 +138,31 @@ def test_record_with_per_thread_info(self):
         {'work': 'workitem', 'stage': 'stage', 'step': 'step'})
     self.assertEqual(log_output, expected_output)
 
+  def test_nested_with_per_thread_info(self):
+    formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid')
+    with logger.PerThreadLoggingContext(
+        work_item_id='workitem', stage_name='stage', step_name='step1'):
+      record = self.create_log_record(**self.SAMPLE_RECORD)
+      log_output1 = json.loads(formatter.format(record))
+
+      with logger.PerThreadLoggingContext(step_name='step2'):
+        record = self.create_log_record(**self.SAMPLE_RECORD)
+        log_output2 = json.loads(formatter.format(record))
+
+      record = self.create_log_record(**self.SAMPLE_RECORD)
+      log_output3 = json.loads(formatter.format(record))
+
+    record = self.create_log_record(**self.SAMPLE_RECORD)
+    log_output4 = json.loads(formatter.format(record))
+
+    self.assertEqual(log_output1, dict(
+        self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
+    self.assertEqual(log_output2, dict(
+        self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step2'))
+    self.assertEqual(log_output3, dict(
+        self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
+    self.assertEqual(log_output4, self.SAMPLE_OUTPUT)
+
   def test_exception_record(self):
     formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid')
     try: