PYTHON-3406 Log traceback when fork() test encounters a deadlock (#1045)

Co-authored-by: Ben Warner <[email protected]>

Author: ShaneHarvey

test/__init__.py

@@ -17,7 +17,9 @@
 
 import base64
 import gc
+import multiprocessing
 import os
+import signal
 import socket
 import sys
 import threading
@@ -43,7 +45,7 @@
 from contextlib import contextmanager
 from functools import wraps
 from test.version import Version
-from typing import Dict, Generator, no_type_check
+from typing import Callable, Dict, Generator, no_type_check
 from unittest import SkipTest
 from urllib.parse import quote_plus
 
@@ -999,31 +1001,33 @@ def fail_point(self, command_args):
             )
 
     @contextmanager
-    def fork(self) -> Generator[int, None, None]:
+    def fork(
+        self, target: Callable, timeout: float = 60
+    ) -> Generator[multiprocessing.Process, None, None]:
         """Helper for tests that use os.fork()
 
         Use in a with statement:
 
-            with self.fork() as pid:
-                if pid == 0:  # Child
-                    pass
-                else:  # Parent
-                    pass
+            with self.fork(target=lambda: print('in child')) as proc:
+                self.assertTrue(proc.pid)  # Child process was started
         """
-        pid = os.fork()
-        in_child = pid == 0
+        ctx = multiprocessing.get_context("fork")
+        proc = ctx.Process(target=target)
+        proc.start()
         try:
-            yield pid
-        except:
-            if in_child:
-                traceback.print_exc()
-                os._exit(1)
-            raise
+            yield proc  # type: ignore
         finally:
-            if in_child:
-                os._exit(0)
-            # In parent, assert child succeeded.
-            self.assertEqual(0, os.waitpid(pid, 0)[1])
+            proc.join(timeout)
+            pid = proc.pid
+            assert pid
+            if proc.exitcode is None:
+                # If it failed, SIGINT to get traceback and wait 10s.
+                os.kill(pid, signal.SIGINT)
+                proc.join(10)
+                proc.kill()
+                proc.join(1)
+                self.fail(f"child timed out after {timeout}s (see traceback in logs): deadlock?")
+            self.assertEqual(proc.exitcode, 0)
 
 
 class IntegrationTest(PyMongoTestCase):

test/test_encryption.py

@@ -341,9 +341,13 @@ def test_use_after_close(self):
     def test_fork(self):
         opts = AutoEncryptionOpts(KMS_PROVIDERS, "keyvault.datakeys")
         client = rs_or_single_client(auto_encryption_opts=opts)
-        with self.fork():
+        self.addCleanup(client.close)
+
+        def target():
             client.admin.command("ping")
-            client.close()
+
+        with self.fork(target):
+            target()
 
 
 class TestEncryptedBulkWrite(BulkTestBase, EncryptionIntegrationTest):

test/test_fork.py

@@ -15,22 +15,26 @@
 """Test that pymongo is fork safe."""
 
 import os
+import sys
+import unittest
 from multiprocessing import Pipe
+
+from bson.objectid import ObjectId
+
+sys.path[0:0] = [""]
+
 from test import IntegrationTest
 from test.utils import (
     ExceptionCatchingThread,
     is_greenthread_patched,
     rs_or_single_client,
 )
-from unittest import skipIf
 
-from bson.objectid import ObjectId
 
-
-@skipIf(
+@unittest.skipIf(
     not hasattr(os, "register_at_fork"), "register_at_fork not available in this version of Python"
 )
-@skipIf(
+@unittest.skipIf(
     is_greenthread_patched(),
     "gevent and eventlet do not support POSIX-style forking.",
 )
@@ -40,20 +44,26 @@ def test_lock_client(self):
         # Parent => All locks should be as before the fork.
         # Child => All locks should be reset.
         with self.client._MongoClient__lock:
-            with self.fork() as pid:
-                if pid == 0:  # Child
-                    self.client.admin.command("ping")
+
+            def target():
+                self.client.admin.command("ping")
+
+            with self.fork(target):
+                pass
         self.client.admin.command("ping")
 
     def test_lock_object_id(self):
         # Forks the client with ObjectId's _inc_lock locked.
         # Parent => _inc_lock should remain locked.
         # Child => _inc_lock should be unlocked.
         with ObjectId._inc_lock:
-            with self.fork() as pid:
-                if pid == 0:  # Child
-                    self.assertFalse(ObjectId._inc_lock.locked())
-                    self.assertTrue(ObjectId())
+
+            def target():
+                self.assertFalse(ObjectId._inc_lock.locked())
+                self.assertTrue(ObjectId())
+
+            with self.fork(target):
+                pass
 
     def test_topology_reset(self):
         # Tests that topologies are different from each other.
@@ -63,22 +73,23 @@ def test_topology_reset(self):
         parent_conn, child_conn = Pipe()
         init_id = self.client._topology._pid
         parent_cursor_exc = self.client._kill_cursors_executor
-        with self.fork() as pid:
-            if pid == 0:  # Child
-                self.client.admin.command("ping")
-                child_conn.send(self.client._topology._pid)
-                child_conn.send(
-                    (
-                        parent_cursor_exc != self.client._kill_cursors_executor,
-                        "client._kill_cursors_executor was not reinitialized",
-                    )
+
+        def target():
+            self.client.admin.command("ping")
+            child_conn.send(self.client._topology._pid)
+            child_conn.send(
+                (
+                    parent_cursor_exc != self.client._kill_cursors_executor,
+                    "client._kill_cursors_executor was not reinitialized",
                 )
-            else:  # Parent
-                self.assertEqual(self.client._topology._pid, init_id)
-                child_id = parent_conn.recv()
-                self.assertNotEqual(child_id, init_id)
-                passed, msg = parent_conn.recv()
-                self.assertTrue(passed, msg)
+            )
+
+        with self.fork(target):
+            self.assertEqual(self.client._topology._pid, init_id)
+            child_id = parent_conn.recv()
+            self.assertNotEqual(child_id, init_id)
+            passed, msg = parent_conn.recv()
+            self.assertTrue(passed, msg)
 
     def test_many_threaded(self):
         # Fork randomly while doing operations.
@@ -106,10 +117,13 @@ def action(client):
                     rc = self.clients[i % len(self.clients)]
                     if i % 50 == 0 and self.fork:
                         # Fork
-                        with self.runner.fork() as pid:
-                            if pid == 0:  # Child
-                                for c in self.clients:
-                                    action(c)
+                        def target():
+                            for c_ in self.clients:
+                                action(c_)
+                                c_.close()
+
+                        with self.runner.fork(target=target) as proc:
+                            self.runner.assertTrue(proc.pid)
                     action(rc)
 
         threads = [ForkThread(self, clients) for _ in range(10)]
@@ -125,3 +139,7 @@ def action(client):
 
         for c in clients:
             c.close()
+
+
+if __name__ == "__main__":
+    unittest.main()