From bf0739c19446adeb982a167b0509da29f9860104 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Tue, 20 Feb 2024 08:59:21 -0500
Subject: [PATCH] Add `stackscope` tree pprinter triggered by SIGUSR1

Can be optionally enabled via a new `enable_stack_on_sig()` which will
swap in the SIGUSR1 handler. Much thanks to @oremanj for writing this
amazing project, it's thus far helped me fix some very subtle hangs
inside our new IPC-context cancellation machinery that would have
otherwise taken much more manual pdb-ing and hair pulling XD

Full credit for `dump_task_tree()` goes to the original project author
with some minor tweaks as was handed to me via the trio-general matrix
room B)

Slight changes from orig version:
- use a `log.pdb()` emission to pprint to console
- toss in an ex sh CLI cmd to trigger the dump from another terminal
  using `kill` + `pgrep`.
---
 tractor/devx/__init__.py    |  3 ++
 tractor/devx/_stackscope.py | 84 +++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 tractor/devx/_stackscope.py

diff --git a/tractor/devx/__init__.py b/tractor/devx/__init__.py
index 89b9a33..5f83261 100644
--- a/tractor/devx/__init__.py
+++ b/tractor/devx/__init__.py
@@ -32,6 +32,9 @@ from ._debug import (
     maybe_open_crash_handler,
     post_mortem,
 )
+from ._stackscope import (
+    enable_stack_on_sig as enable_stack_on_sig,
+)
 
 __all__ = [
     'maybe_wait_for_debugger',
diff --git a/tractor/devx/_stackscope.py b/tractor/devx/_stackscope.py
new file mode 100644
index 0000000..706b85d
--- /dev/null
+++ b/tractor/devx/_stackscope.py
@@ -0,0 +1,84 @@
+# tractor: structured concurrent "actors".
+# Copyright eternity Tyler Goodlet.
+
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+'''
+The fundamental cross process SC abstraction: an inter-actor,
+cancel-scope linked task "context".
+
+A ``Context`` is very similar to the ``trio.Nursery.cancel_scope`` built
+into each ``trio.Nursery`` except it links the lifetimes of memory space
+disjoint, parallel executing tasks in separate actors.
+
+'''
+from signal import (
+    signal,
+    SIGUSR1,
+)
+
+import trio
+
+@trio.lowlevel.disable_ki_protection
+def dump_task_tree() -> None:
+    import stackscope
+    from tractor.log import get_console_log
+
+    tree_str: str = str(
+        stackscope.extract(
+            trio.lowlevel.current_root_task(),
+            recurse_child_tasks=True
+        )
+    )
+    log = get_console_log('cancel')
+    log.pdb(
+        f'Dumping `stackscope` tree:\n\n'
+        f'{tree_str}\n'
+    )
+    # import logging
+    # try:
+    #     with open("/dev/tty", "w") as tty:
+    #         tty.write(tree_str)
+    # except BaseException:
+    #     logging.getLogger(
+    #         "task_tree"
+    #     ).exception("Error printing task tree")
+
+
+def signal_handler(sig: int, frame: object) -> None:
+    import traceback
+    try:
+        trio.lowlevel.current_trio_token(
+        ).run_sync_soon(dump_task_tree)
+    except RuntimeError:
+        # not in async context -- print a normal traceback
+        traceback.print_stack()
+
+
+
+def enable_stack_on_sig(
+    sig: int = SIGUSR1
+) -> None:
+    '''
+    Enable `stackscope` tracing on reception of a signal; by
+    default this is SIGUSR1.
+
+    '''
+    signal(
+        sig,
+        signal_handler,
+    )
+    # NOTE: not the above can be triggered from
+    # a (xonsh) shell using:
+    # kill -SIGUSR1 @$(pgrep -f '<cmd>')