From cb0fbe3208a43857e3062412ae375ea1f6f09fcd Mon Sep 17 00:00:00 2001
From: Kyle Altendorf <sda@fstab.net>
Date: Sat, 5 Nov 2022 18:46:40 -0400
Subject: [PATCH] switch to context manager for task timing (#13473)

---
 chia/server/start_full_node.py | 10 +++-------
 chia/server/start_wallet.py    |  9 +++------
 chia/util/task_timing.py       | 21 ++++++++++++++++++++-
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/chia/server/start_full_node.py b/chia/server/start_full_node.py
index 6a07dd611d54..908754caba9e 100644
--- a/chia/server/start_full_node.py
+++ b/chia/server/start_full_node.py
@@ -18,6 +18,7 @@ from chia.util.chia_logging import initialize_service_logging
 from chia.util.config import load_config, load_config_cli
 from chia.util.default_root import DEFAULT_ROOT_PATH
 from chia.util.ints import uint16
+from chia.util.task_timing import maybe_manage_task_instrumentation
 
 # See: https://bugs.python.org/issue29288
 "".encode("idna")
@@ -84,14 +85,9 @@ async def async_main() -> int:
 
 def main() -> int:
     freeze_support()
-    if os.getenv("CHIA_INSTRUMENT_NODE", 0) != 0:
-        import atexit
 
-        from chia.util.task_timing import start_task_instrumentation, stop_task_instrumentation
-
-        start_task_instrumentation()
-        atexit.register(stop_task_instrumentation)
-    return async_run(async_main())
+    with maybe_manage_task_instrumentation(enable=os.environ.get("CHIA_INSTRUMENT_NODE") is not None):
+        return async_run(async_main())
 
 
 if __name__ == "__main__":
diff --git a/chia/server/start_wallet.py b/chia/server/start_wallet.py
index c5e99bae79f3..ff84917b1783 100644
--- a/chia/server/start_wallet.py
+++ b/chia/server/start_wallet.py
@@ -14,6 +14,7 @@ from chia.util.chia_logging import initialize_service_logging
 from chia.util.config import load_config_cli, load_config
 from chia.util.default_root import DEFAULT_ROOT_PATH
 from chia.util.keychain import Keychain
+from chia.util.task_timing import maybe_manage_task_instrumentation
 from chia.wallet.wallet_node import WalletNode
 
 # See: https://bugs.python.org/issue29288
@@ -108,13 +109,9 @@ async def async_main() -> int:
 
 def main() -> int:
     freeze_support()
-    if os.getenv("CHIA_INSTRUMENT_WALLET", 0) != 0:
-        from chia.util.task_timing import start_task_instrumentation, stop_task_instrumentation
-        import atexit
 
-        start_task_instrumentation()
-        atexit.register(stop_task_instrumentation)
-    return async_run(async_main())
+    with maybe_manage_task_instrumentation(enable=os.environ.get("CHIA_INSTRUMENT_WALLET") is not None):
+        return async_run(async_main())
 
 
 if __name__ == "__main__":
diff --git a/chia/util/task_timing.py b/chia/util/task_timing.py
index 0fbc4c721459..213e489fa07a 100644
--- a/chia/util/task_timing.py
+++ b/chia/util/task_timing.py
@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import inspect
 import os
 import sys
 import time
 from types import FrameType
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
 # This is a development utility that instruments tasks (coroutines) and records
 # wall-clock time they spend in various functions. Since it relies on
@@ -334,6 +335,24 @@ def stop_task_instrumentation(target_dir: str = f"task-profile-{os.getpid()}") -
             f.write("}\n")
 
 
+@contextlib.contextmanager
+def manage_task_instrumentation() -> Iterator[None]:
+    start_task_instrumentation()
+    try:
+        yield
+    finally:
+        stop_task_instrumentation()
+
+
+@contextlib.contextmanager
+def maybe_manage_task_instrumentation(enable: bool) -> Iterator[None]:
+    if enable:
+        with manage_task_instrumentation():
+            yield
+    else:
+        yield
+
+
 def main(args: List[str]) -> int:
     import glob
     import pathlib