From 8276b81b7d9fc649ded73484f9340c86c7c8a719 Mon Sep 17 00:00:00 2001
From: Alex Willmer <alex@moreati.org.uk>
Date: Fri, 26 Nov 2021 21:51:16 +0000
Subject: [PATCH] ci: Account for pre-existing children in process leak checks

---
 tests/testlib.py | 36 +++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/tests/testlib.py b/tests/testlib.py
index d40ce573..ef35e13e 100644
--- a/tests/testlib.py
+++ b/tests/testlib.py
@@ -1,4 +1,5 @@
 
+import errno
 import logging
 import os
 import random
@@ -11,6 +12,7 @@ import threading
 import time
 import traceback
 
+import psutil
 import unittest2
 
 import mitogen.core
@@ -67,7 +69,6 @@ def get_fd_count():
     """
     Return the number of FDs open by this process.
     """
-    import psutil
     return psutil.Process().num_fds()
 
 
@@ -334,6 +335,10 @@ class TestCase(unittest2.TestCase):
         # Broker() instantiations in setUp() etc.
         mitogen.fork.on_fork()
         cls._fd_count_before = get_fd_count()
+        # Ignore children started by external packages - in particular
+        # multiprocessing.resource_tracker.main()`, started when some Ansible
+        # versions instantiate a `multithreading.Lock()`.
+        cls._children_before = frozenset(psutil.Process().children())
         super(TestCase, cls).setUpClass()
 
     ALLOWED_THREADS = set([
@@ -361,7 +366,10 @@ class TestCase(unittest2.TestCase):
     def _teardown_check_fds(self):
         mitogen.core.Latch._on_fork()
         if get_fd_count() != self._fd_count_before:
-            import os; os.system('lsof +E -w -p %s | grep -vw mem' % (os.getpid(),))
+            if sys.platform == 'linux':
+                os.system('lsof +E -w -p %i | grep -vw mem' % (os.getpid(),))
+            else:
+                os.system('lsof -w -p %i | grep -vw mem' % (os.getpid(),))
             assert 0, "%s leaked FDs. Count before: %s, after: %s" % (
                 self, self._fd_count_before, get_fd_count(),
             )
@@ -374,19 +382,33 @@ class TestCase(unittest2.TestCase):
         if self.no_zombie_check:
             return
 
+        # pid=0: Wait for any child process in the same process group as us.
+        # WNOHANG: Don't block if no processes ready to report status.
         try:
             pid, status = os.waitpid(0, os.WNOHANG)
-        except OSError:
-            return  # ECHILD
+        except OSError as e:
+            # ECHILD: there are no child processes in our group.
+            if e.errno == errno.ECHILD:
+                return
+            raise
 
         if pid:
             assert 0, "%s failed to reap subprocess %d (status %d)." % (
                 self, pid, status
             )
 
-        print('')
-        print('Children of unit test process:')
-        os.system('ps uww --ppid ' + str(os.getpid()))
+        children_after = frozenset(psutil.Process().children())
+        children_leaked = children_after.difference(self._children_before)
+        if not children_leaked:
+            return
+
+        print('Leaked children of unit test process:')
+        os.system('ps -o "user,pid,%%cpu,%%mem,vsz,rss,tty,stat,start,time,command" -ww -p %s'
+                  % (','.join(str(p.pid) for p in children_leaked),))
+        if self._children_before:
+            print('Pre-existing children of unit test process:')
+            os.system('ps -o "user,pid,%%cpu,%%mem,vsz,rss,tty,stat,start,time,command" -ww -p %s'
+                      % (','.join(str(p.pid) for p in self._children_before),))
         assert 0, "%s leaked still-running subprocesses." % (self,)
 
     def tearDown(self):