Merge pull request #7 from P403n1x87/chore/patch-on-import

chore: patch modules on import

.flake8

@@ -1,6 +1,6 @@
 [flake8]
 select = ANN,B,B9,C,D,E,F,W,I
-ignore = ANN001,ANN101,ANN102,ANN201,ANN401,B950,BLK100,D100,D101,D102,D103,D104,D107,E203,E501,I001,I005,I201,I202,W503,W606
+ignore = ANN001,ANN101,ANN102,ANN201,ANN401,B950,BLK100,D100,D101,D102,D103,D104,D105,D107,E203,E501,I001,I005,I201,I202,W503,W606
 exclude =
   austin/format/pprof/profile_pb2.py
 max-line-length = 88

echion/bootstrap/__init__.py

@@ -5,9 +5,10 @@
 import atexit
 import os
 import sys
-from threading import Thread
+from types import ModuleType
 
 import echion.core as ec
+from echion.module import ModuleWatchdog
 
 
 # We cannot unregister the fork hook, so we use this flag instead
@@ -34,10 +35,26 @@ def start():
     ec.set_native(bool(int(os.getenv("ECHION_NATIVE", 0))))
     ec.set_where(bool(int(os.getenv("ECHION_WHERE", 0) or 0)))
 
-    # Monkey-patch the standard library
-    for module in ("echion.monkey.asyncio", "echion.monkey.threading"):
-        __import__(module)
-        sys.modules[module].patch()
+    # Monkey-patch the standard library on import
+    try:
+        ModuleWatchdog.install()
+        atexit.register(ModuleWatchdog.uninstall)
+    except RuntimeError:
+        # If ModuleWatchdog we have already registered the import hooks.
+        pass
+    else:
+        for module in ("asyncio", "threading"):
+
+            @ModuleWatchdog.after_module_imported(module)
+            def _(module: ModuleType) -> None:
+                echion_module = f"echion.monkey.{module.__name__}"
+                __import__(echion_module)
+                sys.modules[echion_module].patch()
+
+            try:
+                sys.modules[f"echion.monkey.{module}"].track()
+            except KeyError:
+                pass
 
     do_on_fork = True
     os.register_at_fork(after_in_child=restart_on_fork)
@@ -46,6 +63,8 @@ def start():
     if int(os.getenv("ECHION_STEALTH", 0)):
         ec.start_async()
     else:
+        from threading import Thread
+
         Thread(target=ec.start, name="echion.core.sampler", daemon=True).start()
 
 
@@ -54,8 +73,16 @@ def stop():
 
     ec.stop()
 
-    for module in ("echion.monkey.asyncio", "echion.monkey.threading"):
-        sys.modules[module].unpatch()
+    for module in ("asyncio", "threading"):
+        echion_module = f"echion.monkey.{module}"
+        try:
+            sys.modules[echion_module].unpatch()
+        except KeyError:
+            pass
+
+    ModuleWatchdog.uninstall()
 
     atexit.unregister(stop)
+    atexit.unregister(ModuleWatchdog.uninstall)
+
     do_on_fork = False

echion/coremodule.cc

@@ -99,7 +99,6 @@ static void for_each_thread(std::function<void(PyThreadState *, ThreadInfo *)> c
 
             if (thread_info_map.find(tstate.thread_id) == thread_info_map.end())
             {
-#if defined PL_DARWIN
                 // If the threading module was not imported in the target then
                 // we mistakenly take the hypno thread as the main thread. We
                 // assume that any missing thread is the actual main thread.
@@ -110,15 +109,14 @@ static void for_each_thread(std::function<void(PyThreadState *, ThreadInfo *)> c
                 new_info->name = new char[name.length() + 1];
                 std::strcpy((char *)new_info->name, name.c_str());
 
+#if defined PL_DARWIN
                 pthread_threadid_np((pthread_t)tstate.thread_id, (__uint64_t *)&new_info->native_id);
                 new_info->mach_port = pthread_mach_thread_np((pthread_t)tstate.thread_id);
+#endif
 
                 new_info->update_cpu_time();
 
                 thread_info_map[tstate.thread_id] = new_info;
-#else
-                return;
-#endif
             }
 
             ThreadInfo *info = thread_info_map[tstate.thread_id];

echion/frame.h

@@ -304,7 +304,7 @@ bool Frame::is_valid()
 
 // ----------------------------------------------------------------------------
 
-static Frame *INVALID_FRAME = new Frame("INVALID");
+static Frame *INVALID_FRAME = nullptr;
 static Frame *UNKNOWN_FRAME = new Frame("<unknown>");
 
 static LRUCache<uintptr_t, Frame> *frame_cache = nullptr;
@@ -369,7 +369,8 @@ Frame *Frame::read(PyObject *frame_addr, PyObject **prev_addr)
     const int lasti = ((int)(iframe.prev_instr - (_Py_CODEUNIT *)(iframe.f_code))) - offsetof(PyCodeObject, co_code_adaptive) / sizeof(_Py_CODEUNIT);
     Frame *frame = Frame::get(iframe.f_code, lasti);
 
-    frame->is_entry = iframe.is_entry;
+    if (frame != INVALID_FRAME)
+        frame->is_entry = iframe.is_entry;
 
     *prev_addr = frame == INVALID_FRAME ? NULL : (PyObject *)iframe.previous;