From 8a4a293ab4f58cd5885ce3067d6629df374090c3 Mon Sep 17 00:00:00 2001
From: englehardt <englehardt@gmail.com>
Date: Fri, 9 Aug 2019 16:55:53 -0700
Subject: [PATCH] Pass all log information across process boundary.

The current inter-process log handler discards `exc_info` and adds the
arguments to the exception message. This prevents Sentry from properly
grouping log messages by their exception and call information. We send
the information if possible (or fall back to the previous approach if
not).
---
 automation/MPLogger.py | 41 +++++++++++++++++++++++++++++------------
 test/test_mp_logger.py | 23 ++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/automation/MPLogger.py b/automation/MPLogger.py
index 779cc05f..336ff216 100644
--- a/automation/MPLogger.py
+++ b/automation/MPLogger.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, print_function
 
+import copy
 import json
 import logging
 import logging.handlers
@@ -10,14 +11,19 @@ import sys
 import threading
 import time
 
+import dill
 import sentry_sdk
+import six
 from multiprocess import JoinableQueue
 from sentry_sdk.integrations.logging import BreadcrumbHandler, EventHandler
 from six.moves.queue import Empty as EmptyQueue
 from six.moves.urllib import parse as urlparse
+from tblib import pickling_support
 
 from .SocketInterface import serversocket
 
+pickling_support.install()
+
 BROWSER_PREFIX = re.compile(r"^BROWSER (-)?\d+:\s*")
 NETERROR_RE = re.compile(
     r"WebDriverException: Message: Reached error page: about:neterror\?(.*)\."
@@ -34,18 +40,25 @@ class ClientSocketHandler(logging.handlers.SocketHandler):
         Serializes the record via json and prepends a length/serialization
         flag. Returns it ready for transmission across the socket.
         """
-        ei = record.exc_info
-        if ei:
-            # just to get traceback text into record.exc_text ...
-            dummy = self.format(record)  # noqa
-            record.exc_info = None  # to avoid Unpickleable error
-        d = dict(record.__dict__)
-        d['msg'] = record.getMessage()
-        d['args'] = None
-        s = json.dumps(d).encode('utf-8')
-        if ei:
-            record.exc_info = ei  # for next handler
-        return struct.pack('>Lc', len(s), b'j') + s
+        d = copy.deepcopy(record.__dict__)
+
+        # Pickle fields to so record is safe to send across socket
+        if 'exc_info' in d and d['exc_info']:
+            try:
+                d['exc_info'] = dill.dumps(d['exc_info'])
+            except dill.PicklingError:
+                d['exc_info'] = None
+        if 'args' in d and d['args']:
+            try:
+                d['args'] = dill.dumps(d['args'])
+            except dill.PicklingError:
+                d['msg'] = record.getMessage()
+                d['args'] = None
+
+        # Serialize logging record so it can be sent to MPLogger
+        # s = json.dumps(d).encode('utf-8')
+        s = dill.dumps(d)
+        return struct.pack('>Lc', len(s), b'd') + s
 
 
 class MPLogger(object):
@@ -207,6 +220,10 @@ class MPLogger(object):
         This is currently records that are written to a file on disk
         and those sent to Sentry.
         """
+        if obj['exc_info']:
+            obj['exc_info'] = dill.loads(six.ensure_str(obj['exc_info']))
+        if obj['args']:
+            obj['args'] = dill.loads(six.ensure_str(obj['args']))
         record = logging.makeLogRecord(obj)
         self._file_handler.emit(record)
         if self._sentry_dsn:
diff --git a/test/test_mp_logger.py b/test/test_mp_logger.py
index 16f52373..e0808d18 100644
--- a/test/test_mp_logger.py
+++ b/test/test_mp_logger.py
@@ -40,10 +40,31 @@ def child_proc(index):
 
 
 def child_proc_with_exception(index):
+    def test_func():
+        print("test")
+
+    class TestClass(object):
+        def __init__(self):
+            return
+
+        def test_method(self):
+            return
+
+    class TestSubClass(TestClass):
+        def test_method(self):
+            print("test method")
+            return
+
     logger = logging.getLogger('openwpm')
     logger.info(CHILD_INFO_STR_1 % index)
     logger.info(CHILD_INFO_STR_2 % index)
-    raise RuntimeError(CHILD_EXCEPTION_STR % index)
+    test_class = TestClass()
+    test_subclass = TestSubClass()
+    raise IOError(
+        CHILD_EXCEPTION_STR % index,
+        ('blah', 1, test_func, test_class, test_subclass,
+         TestClass, TestSubClass)
+    )
 
 
 class TestMPLogger(OpenWPMTest):