From 7dacb68eeb242e74022da9e71d891356ab731cd3 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Mon, 25 Feb 2019 16:33:03 +0000
Subject: [PATCH] issue #552: include process identity in log messages.

---
 ansible_mitogen/logging.py  | 26 +++++++++++++++++++++++++-
 ansible_mitogen/process.py  |  6 ++++--
 ansible_mitogen/strategy.py |  1 +
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/ansible_mitogen/logging.py b/ansible_mitogen/logging.py
index 1c439be8..ce6f1659 100644
--- a/ansible_mitogen/logging.py
+++ b/ansible_mitogen/logging.py
@@ -40,6 +40,25 @@ except ImportError:
     display = Display()
 
 
+#: The process name set via :func:`set_process_name`.
+_process_name = None
+
+#: The PID of the process that last called :func:`set_process_name`, so its
+#: value can be ignored in unknown fork children.
+_process_pid = None
+
+
+def set_process_name(name):
+    """
+    Set a name to adorn log messages with.
+    """
+    global _process_name
+    _process_name = name
+
+    global _process_pid
+    _process_pid = os.getpid()
+
+
 class Handler(logging.Handler):
     """
     Use Mitogen's log format, but send the result to a Display method.
@@ -65,7 +84,12 @@ class Handler(logging.Handler):
         if mitogen_name in self.NOISY_LOGGERS and record.levelno >= logging.WARNING:
             record.levelno = logging.DEBUG
 
-        s = '[pid %d] %s' % (os.getpid(), self.format(record))
+        if _process_pid == os.getpid():
+            process_name = _process_name
+        else:
+            process_name = '?'
+
+        s = '[%-4s %d] %s' % (process_name, os.getpid(), self.format(record))
         if record.levelno >= logging.ERROR:
             display.error(s, wrap_text=False)
         elif record.levelno >= logging.WARNING:
diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py
index d7f36496..e4e61e8b 100644
--- a/ansible_mitogen/process.py
+++ b/ansible_mitogen/process.py
@@ -185,19 +185,21 @@ class MuxProcess(object):
         cls.profiling = os.environ.get('MITOGEN_PROFILING') is not None
         if cls.profiling:
             mitogen.core.enable_profiling()
+        if _init_logging:
+            ansible_mitogen.logging.setup()
 
         cls.original_env = dict(os.environ)
         cls.child_pid = os.fork()
-        if _init_logging:
-            ansible_mitogen.logging.setup()
         if cls.child_pid:
             save_pid('controller')
+            ansible_mitogen.logging.set_process_name('top')
             ansible_mitogen.affinity.policy.assign_controller()
             cls.child_sock.close()
             cls.child_sock = None
             mitogen.core.io_op(cls.worker_sock.recv, 1)
         else:
             save_pid('mux')
+            ansible_mitogen.logging.set_process_name('mux')
             ansible_mitogen.affinity.policy.assign_muxprocess()
             cls.worker_sock.close()
             cls.worker_sock = None
diff --git a/ansible_mitogen/strategy.py b/ansible_mitogen/strategy.py
index 50486841..ba0ff525 100644
--- a/ansible_mitogen/strategy.py
+++ b/ansible_mitogen/strategy.py
@@ -108,6 +108,7 @@ def wrap_worker__run(*args, **kwargs):
     if mitogen.core._profile_hook.__name__ != '_profile_hook':
         signal.signal(signal.SIGTERM, signal.SIG_IGN)
 
+    ansible_mitogen.logging.set_process_name('task')
     ansible_mitogen.affinity.policy.assign_worker()
     return mitogen.core._profile_hook('WorkerProcess',
         lambda: worker__run(*args, **kwargs)