switch to context manager for task timing (#13473)

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__":

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__":

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