From 79622bbeeae86412ce718e5c21fdf198db559893 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Tue, 7 Mar 2023 16:46:14 -0500
Subject: [PATCH] Restore `breakpoint()` hook after runtime exits

Previously we were leaking our (pdb++) override into the Python runtime
which would always result in a runtime error whenever `breakpoint()` is
called outside our runtime; after exit of the root actor . This
explicitly restores any previous hook override (detected during startup)
or deletes the hook and restores the environment if none existed prior.

Also adds a new WIP debugging example script to ensure breakpointing
works as normal after runtime close; this will be added to the test
suite.
---
 .../debugging/restore_builtin_breakpoint.py   | 24 +++++++++++++++++++
 tractor/_root.py                              | 18 +++++++++++---
 2 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 examples/debugging/restore_builtin_breakpoint.py

diff --git a/examples/debugging/restore_builtin_breakpoint.py b/examples/debugging/restore_builtin_breakpoint.py
new file mode 100644
index 0000000..6e141df
--- /dev/null
+++ b/examples/debugging/restore_builtin_breakpoint.py
@@ -0,0 +1,24 @@
+import os
+import sys
+
+import trio
+import tractor
+
+
+async def main() -> None:
+    async with tractor.open_nursery(debug_mode=True) as an:
+
+        assert os.environ['PYTHONBREAKPOINT'] == 'tractor._debug._set_trace'
+
+        # TODO: an assert that verifies the hook has indeed been, hooked
+        # XD
+        assert sys.breakpointhook is not tractor._debug._set_trace
+
+        breakpoint()
+
+    # TODO: an assert that verifies the hook is unhooked..
+    assert sys.breakpointhook
+    breakpoint()
+
+if __name__ == '__main__':
+    trio.run(main)
diff --git a/tractor/_root.py b/tractor/_root.py
index 840b288..64652a1 100644
--- a/tractor/_root.py
+++ b/tractor/_root.py
@@ -22,8 +22,9 @@ from contextlib import asynccontextmanager
 from functools import partial
 import importlib
 import logging
-import os
 import signal
+import sys
+import os
 import typing
 import warnings
 
@@ -84,8 +85,10 @@ async def open_root_actor(
 
     '''
     # Override the global debugger hook to make it play nice with
-    # ``trio``, see:
+    # ``trio``, see much discussion in:
     # https://github.com/python-trio/trio/issues/1155#issuecomment-742964018
+    builtin_bp_handler = sys.breakpointhook
+    orig_bp_path: str | None = os.environ.get('PYTHONBREAKPOINT', None)
     os.environ['PYTHONBREAKPOINT'] = 'tractor._debug._set_trace'
 
     # attempt to retreive ``trio``'s sigint handler and stash it
@@ -254,6 +257,15 @@ async def open_root_actor(
                 await actor.cancel()
     finally:
         _state._current_actor = None
+
+        # restore breakpoint hook state
+        sys.breakpointhook = builtin_bp_handler
+        if orig_bp_path is not None:
+            os.environ['PYTHONBREAKPOINT'] = orig_bp_path
+        else:
+            # clear env back to having no entry
+            os.environ.pop('PYTHONBREAKPOINT')
+
         logger.runtime("Root actor terminated")
 
 
@@ -289,7 +301,7 @@ def run_daemon(
     async def _main():
 
         async with open_root_actor(
-            arbiter_addr=registry_addr,
+            registry_addr=registry_addr,
             name=name,
             start_method=start_method,
             debug_mode=debug_mode,