retracesoftware-proxy 0.1.5__py3-none-any.whl → 0.2.4__py3-none-any.whl

retracesoftware/proxy/thread.py

@@ -1,6 +1,9 @@
 import retracesoftware.functional as functional
 import retracesoftware_utils as utils
 
+import os
+import _thread
+
 # def thread_aware_writer(writer):
 #     on_thread_switch = functional.sequence(utils.thread_id(), writer.handle('THREAD_SWITCH'))
 #     return utils.threadawareproxy(on_thread_switch = on_thread_switch, target = writer)
@@ -11,23 +14,35 @@ class ThreadSwitch:
     def __init__(self, id):
         self.id = id
 
-def set_thread_id(writer, id):
-    utils.set_thread_id(writer.handle(ThreadSwitch(id)))
+    def __repr__(self):
+        return f'ThreadSwitch<{self.id}>'
+
+    def __str__(self):
+        return f'ThreadSwitch<{self.id}>'
+
+# def set_thread_id(writer, id):
+#     utils.sigtrap(id)
+#     utils.set_thread_id(writer.handle(ThreadSwitch(id)))
 
 def write_thread_switch(writer):
     on_thread_switch = functional.repeatedly(functional.sequence(utils.thread_id, writer))
 
     return lambda f: utils.thread_aware_proxy(target = f, on_thread_switch = on_thread_switch, sticky = False)
 
-def prefix_with_thread_id(f, current):
+def prefix_with_thread_id(f, thread_id):
+    current = None
+
     def next():
         nonlocal current, f
+        if current is None: current = thread_id()
+
         obj = f()
 
         while issubclass(type(obj), ThreadSwitch):
             current = obj.id
             obj = f()
 
+        # print(f'prefix_with_thread_id: {(current, obj)}')
         return (current, obj)
 
     return next
@@ -36,7 +51,15 @@ def per_thread_messages(messages):
     thread_id = utils.thread_id
     # thread_id = lambda: 'FOOOOO!!!'
 
-    demux = utils.demux(source = prefix_with_thread_id(messages, thread_id()), key_function = lambda obj: obj[0])
+    def on_timeout(demux, key):
+        print(f'ON TIMEOUT!!!! {key} pending: {demux.pending} {demux.pending_keys}')
+        utils.sigtrap(demux)
+        os._exit(1)
+
+    demux = utils.demux(source = prefix_with_thread_id(messages, thread_id),
+                        key_function = lambda obj: obj[0],
+                        timeout_seconds = 60,
+                        on_timeout = on_timeout)
 
     # def next():
     #     thread,message = demux(thread_id())
@@ -44,3 +67,40 @@ def per_thread_messages(messages):
     
     # return next
     return functional.repeatedly(lambda: demux(thread_id())[1])
+
+
+# _thread.start_new_thread(function, args[, kwargs])
+counters = _thread._local()
+counters.id = ()
+counters.counter = 0
+
+def with_thread_id(thread_id, on_exit, function):
+    def on_call(*args, **kwargs):
+        counters.id = thread_id
+        counters.counter = 0
+
+    def on_result(res):
+        on_exit(thread_id)
+    
+    def on_error(*args):
+        on_exit(thread_id)
+
+    return utils.observer(on_call = on_call, on_result = on_result, on_error = on_error, function = function)
+
+thread_id = functional.lazy(getattr, counters, 'id')
+
+def start_new_thread_wrapper(thread_state, on_exit, start_new_thread):
+
+    def wrapper(function, *args):
+        
+        next_id = counters.id + (counters.counter,)
+        counters.counter += 1
+
+        wrapped_function = with_thread_id(thread_id = next_id,
+                                          on_exit = on_exit, 
+                                          function = thread_state.wrap('internal', function))
+
+        return start_new_thread(wrapped_function, *args)
+
+    return thread_state.dispatch(start_new_thread, internal = wrapper)
+

retracesoftware/replay.py

@@ -0,0 +1,104 @@
+import retracesoftware.stream as stream
+from retracesoftware.run import install, run_with_retrace, ImmutableTypes, thread_states
+import argparse
+import sys
+from pathlib import Path
+from retracesoftware.proxy.thread import thread_id
+from retracesoftware.proxy.replay import ReplayProxySystem
+import retracesoftware.utils as utils
+import json
+from retracesoftware.stackdifference import on_stack_difference
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        prog=f"python -m {sys.argv[0]}",
+        description="Run a Python module with debugging, logging, etc."
+    )
+
+    parser.add_argument(
+        '--verbose', 
+        action='store_true', 
+        help='Enable verbose output'
+    )
+
+    parser.add_argument(
+        '--timeout',   # or '-r'
+        type = int,      # ensures it's a string (optional, but safe)
+        default = 60,    # default value if not provided
+        help = 'the directory to place the recording files'
+    )
+
+    parser.add_argument(
+        '--recording',   # or '-r'
+        type = str,      # ensures it's a string (optional, but safe)
+        default = '.',    # default value if not provided
+        help = 'the directory to place the recording files'
+    )
+
+    return parser.parse_args()
+
+def load_json(file):
+    with open(file, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+def main():
+    # import debugpy
+
+    # port = 1977
+    # debugpy.listen(("127.0.0.1", port))
+    # debugpy.wait_for_client()
+
+        # self.reader = stream.reader(path, 
+        #                             thread = thread_id,
+        #                             timeout_seconds = 60,
+        #                             verbose = verbose,
+        #                             on_stack_difference = thread_state.wrap('disabled', on_stack_difference),
+        #                             magic_markers = magic_markers)
+
+    # return ReplayProxySystem(thread_state = thread_state, 
+    #                             immutable_types = immutable_types,
+    #                             tracing_config = tracing_config,
+    #                             mainscript = mainscript,
+    #                             path = recording_path / 'trace.bin',
+    #                             tracecalls = env_truthy('RETRACE_ALL', False),
+    #                             verbose = verbose,
+    #                             magic_markers = env_truthy('RETRACE_MAGIC_MARKERS', False))
+
+    args = parse_args()    
+    path = Path(args.recording)
+
+    if not path.exists():
+        raise Exception(f"Recording path: {path} does not exist")
+
+    settings = load_json(path / "settings.json")
+
+    thread_state = utils.ThreadState(*thread_states)
+
+    with stream.reader(path = path / 'trace.bin',
+                       thread = thread_id, 
+                       timeout_seconds = args.timeout,
+                       verbose = args.verbose,
+                       on_stack_difference = thread_state.wrap('disabled', on_stack_difference),
+                       magic_markers = settings['magic_markers']) as reader:
+
+        tracing_config = {}
+
+        system = ReplayProxySystem(
+            reader = reader,
+            thread_state = thread_state,
+            immutable_types = ImmutableTypes(), 
+            tracing_config = tracing_config,
+            tracecalls = settings['trace_inputs'])
+
+        install(system)
+
+        run_with_retrace(system, settings['argv'])
+
+        # install(system)
+
+        # run_with_retrace(system, args.rest[1:])
+
+        # runpy.run_module('foo', run_name="__main__", alter_sys=False)
+
+if __name__ == "__main__":
+    main()

retracesoftware/run.py

@@ -0,0 +1,378 @@
+import sys
+import os
+import runpy
+from pathlib import Path
+# from retracesoftware.install.phases import *
+import pkgutil
+import tomllib
+import retracesoftware.functional as functional
+import builtins
+import importlib
+import _imp
+import importlib._bootstrap_external as _bootstrap_external
+import atexit
+import threading
+import _signal
+import retracesoftware.utils as utils
+from retracesoftware.install.patcher import patch_module, create_patcher, patch_imported_module
+from retracesoftware.proxy.startthread import patch_thread_start
+from retracesoftware.install.replace import update
+
+thread_states = [
+    "disabled", # Default state when retrace is disabled for a thread
+    "internal", # Default state when retrace is disabled for a thread
+    "external", # When target thread is running outside the python system to be recorded
+    "retrace", # When target thread is running outside the retrace system
+    "importing", # When target thread is running outside the retrace system
+    "gc", # When the target thread is running inside the pyton garbage collector
+]
+
+class ImmutableTypes(set):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __contains__(self, item):
+        assert isinstance(item, type)
+
+        if super().__contains__(item):
+            return True
+
+        for elem in self:
+            if issubclass(item, elem):
+                self.add(item)
+                return True
+
+        return False            
+
+def load_module_config(filename):
+    data = pkgutil.get_data("retracesoftware", filename)
+    assert data is not None
+    return tomllib.loads(data.decode("utf-8"))
+
+def wait_for_non_daemon_threads(timeout=None):
+    """
+    Wait for all non-daemon threads to finish, just like Python does on exit.
+    
+    Args:
+        timeout (float, optional): Max seconds to wait. None = wait forever.
+    
+    Returns:
+        bool: True if all threads finished, False if timeout exceeded.
+    """
+    import threading
+    import time
+
+    start_time = time.time()
+    main_thread = threading.main_thread()
+
+    while True:
+        # Get all active threads
+        active = threading.enumerate()
+
+        # Filter: non-daemon and not the main thread
+        non_daemon_threads = [
+            t for t in active
+            if t is not main_thread and not t.daemon
+        ]
+
+        if not non_daemon_threads:
+            return True  # All done!
+
+        # Check timeout
+        if timeout is not None:
+            elapsed = time.time() - start_time
+            if elapsed >= timeout:
+                print(f"Timeout: {len(non_daemon_threads)} thread(s) still alive")
+                return False
+
+        # Sleep briefly to avoid busy-wait
+        time.sleep(0.1)
+
+def run_python_command(argv):
+    """
+    Run a Python app from a command list using runpy.
+
+    Supports:
+        ['-m', 'module', 'arg1', ...]     → like `python -m module arg1 ...`
+        ['script.py', 'arg1', ...]        → like `python script.py arg1 ...`
+
+    Args:
+        argv: List of command-line arguments (first item is either '-m' or script path)
+
+    Returns:
+        Exit code (0 on success, 1+ on error)
+    """
+    if not argv:
+        print("Error: No command provided", file=sys.stderr)
+        return 1
+
+    original_argv = sys.argv[:]
+    original_cwd = os.getcwd()
+
+    try:
+        if argv[0] == '-m':
+            if len(argv) < 2:
+                print("Error: -m requires a module name", file=sys.stderr)
+                return 1
+            module_name = argv[1]
+            module_args = argv[2:]
+            sys.argv = ['-m', module_name] + module_args
+            runpy.run_module(module_name, run_name="__main__")
+            return 0
+
+        else:
+            script_path = argv[0]
+            script_args = argv[1:]
+            path = Path(script_path)
+
+            if not path.exists():
+                print(f"Error: Script not found: {script_path}", file=sys.stderr)
+                return 1
+
+            if path.suffix != ".py":
+                print(f"Error: Not a Python script: {script_path}", file=sys.stderr)
+                return 1
+
+            full_path = str(path.resolve())
+            sys.argv = [full_path] + script_args
+
+            # Change to script's directory
+            os.chdir(path.parent)
+
+            runpy.run_path(full_path, run_name="__main__")
+            return 0
+
+    except ModuleNotFoundError:
+        print(f"Error: No module named '{argv[1]}'", file=sys.stderr)
+        return 1
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    finally:
+        sys.argv = original_argv
+        os.chdir(original_cwd)
+
+def patch_import(thread_state, patcher, sync, checkpoint):
+    # builtins.__import__ = thread_state.dispatch(builtins.__import__, internal = sync(thread_state.wrap('importing', builtins.__import__)))
+    # bi = builtins.__import__
+    # def foo(*args, **kwargs):
+    #     print(f'in patched __import__: {thread_state.value} {args[0]}')
+    #     if thread_state.value == 'internal':
+    #         with thread_state.select('importing'):
+    #             return bi(*args, **kwargs)
+    #     else:
+    #         return bi(*args, **kwargs)
+        
+    # builtins.__import__ = foo
+
+    builtins.__import__ = thread_state.dispatch(builtins.__import__, internal = thread_state.wrap('importing', builtins.__import__))
+
+    def exec(source, globals = None, locals = None):
+        checkpoint(f"exec module: {globals.get('__name__', 'unknown')}")
+        res = builtins.exec(source, globals, locals)
+        patcher(globals, False)
+        return res
+
+    def patch(module):
+        patcher(module.__dict__, False)
+        return module
+
+    _imp.exec_dynamic = thread_state.dispatch(_imp.exec_dynamic, 
+        importing = functional.juxt(
+                        thread_state.wrap('internal', _imp.exec_dynamic), 
+                        patch))
+    # _imp.exec_dynamic = lambda mod: utils.sigtrap(f'{thread_state.value} - {mod}')
+    # _imp.exec_builtin = lambda mod: utils.sigtrap(f'{thread_state.value} - {mod}')
+
+    _imp.exec_builtin = thread_state.dispatch(_imp.exec_builtin, 
+        importing = functional.juxt(
+                        thread_state.wrap('internal', _imp.exec_builtin), 
+                        patch))
+
+    # def runpy_exec(source, globals = None, locals = None):
+    #     print(f'In runpy exec!!!!!')
+    #     return builtins.exec(source, globals, locals)
+    
+    # utils.update(runpy, "_run_code",
+    #              utils.wrap_func_with_overrides,
+    #              exec = runpy_exec)
+
+    utils.update(_bootstrap_external._LoaderBasics, "exec_module", 
+                 utils.wrap_func_with_overrides,
+                 exec = thread_state.dispatch(builtins.exec, importing = thread_state.wrap('internal', sync(exec))))
+
+preload = [
+    "logging",
+    "pathlib",
+    "_signal",
+    "_posixsubprocess",
+    "socket",
+    "select",
+    "ssl",
+    "random",
+    "email",
+    "email.errors",
+    "http.client",
+
+    "json",
+    "typing",
+    "queue",
+    "mimetypes",
+    "tempfile",
+    "zipfile",
+    "importlib.resources",
+    "importlib.metadata",
+    "encodings.idna"
+
+    # "http.client",
+    # "queue",
+    # "mimetypes",
+    # "encodings.idna",
+    # "hmac",
+    # "ipaddress",
+    # "tempfile",
+    # "zipfile",
+    # "importlib.resources",
+    # "importlib.metadata",
+    # "atexit",
+    # "weakref"
+  ]
+
+# def debugger_is_active():
+#     return sys.gettrace() and 'debugpy' in sys.modules
+
+def init_weakref():
+    import weakref
+    def dummy_callback():
+        pass
+
+    class DummyTarget:
+        pass
+
+    f = weakref.finalize(DummyTarget(), dummy_callback)
+    f.detach()
+
+def wrapped_weakref(ref, thread_state, wrap_callback):
+    orig_new = ref.__new__
+
+    def __new__(cls, ob, callback=None, **kwargs):
+        return orig_new(cls, ob, wrap_callback(callback) if callback else None)
+    
+    return type('ref', (ref, ), {'__new__': thread_state.dispatch(orig_new, internal = __new__)})
+
+def patch_weakref(thread_state, wrap_callback):
+    import _weakref
+
+    update(_weakref.ref, wrapped_weakref(_weakref.ref, thread_state, wrap_callback))
+    # _weakref.ref = wrapped_weakref(_weakref.ref)
+
+# def patch_signal(thread_state, wrap_callback):  
+
+#     def wrap_handler(handler):
+#         return utils.observer(on_call = ..., function = handler)
+
+#     _signal.signal = 
+#     update(_signal.signal, thread_state.dispatch(_signal.signal, internal = thread_state.wrap('internal', _signal.signal)))
+
+
+
+def install(system):
+
+    patch_weakref(thread_state = system.thread_state, wrap_callback = system.wrap_weakref_callback)
+
+    init_weakref()
+    
+    preload = pkgutil.get_data("retracesoftware", "preload.txt")
+
+    for name in preload.decode("utf-8").splitlines():
+        try:
+            importlib.import_module(name.strip())
+        except ModuleNotFoundError:
+            pass
+
+    # if 'pydevd' in sys.modules:
+    #     utils.update(sys.modules['pydevd'].PyDB, 'enable_tracing', system.disable_for)
+    #     utils.update(sys.modules['pydevd'].PyDB, 'set_suspend', system.disable_for)
+    #     utils.update(sys.modules['pydevd'].PyDB, 'do_wait_suspend', system.disable_for)
+
+    # if '_pydevd_bundle.pydevd_trace_dispatch_regular' in sys.modules:
+    #     mod = sys.modules['_pydevd_bundle.pydevd_trace_dispatch_regular']
+    #     utils.update(mod.ThreadTracer, '__call__', system.disable_for)
+        
+    for function in utils.stack_functions():
+        system.exclude_from_stacktrace(function)
+
+    def recursive_disable(func):
+        if not callable(func):
+            return func
+        
+        def wrapped(*args, **kwargs):
+            with system.thread_state.select('disabled'):
+                return recursive_disable(func(*args, **kwargs))
+        
+        return wrapped
+
+    sys.settrace = functional.sequence(recursive_disable, sys.settrace)
+    sys.setprofile = functional.sequence(recursive_disable, sys.setprofile)
+    threading.settrace = functional.sequence(recursive_disable, threading.settrace)
+
+    sys.settrace(sys.gettrace())
+    sys.setprofile(sys.getprofile())
+
+    system.checkpoint('About to install retrace system')
+
+    module_config = load_module_config('modules.toml')
+
+    patch_loaded = functional.partial(patch_module, create_patcher(system), module_config)
+    patch_imported = functional.partial(patch_imported_module, create_patcher(system), system.checkpoint, module_config)
+    
+    system.checkpoint('Started installing system 1')
+    
+    for modname in module_config.keys():
+        if modname in sys.modules:
+            patch_loaded(sys.modules[modname].__dict__, True)
+
+    system.checkpoint('About to patch threading')
+
+    patch_thread_start(system.thread_state)
+    threading.current_thread().__retrace__ = system
+
+    system.checkpoint('About to patch import')
+    patch_import(thread_state = system.thread_state, 
+                 patcher = patch_imported, 
+                 sync = system.sync, 
+                 checkpoint = system.checkpoint)
+
+    # print(f'MODULES: {list(sys.modules.keys())}')
+
+    importlib.import_module = \
+        system.thread_state.dispatch(system.disable_for(importlib.import_module),
+                              internal = system.thread_state.wrap('importing', importlib.import_module))
+
+
+    system.checkpoint('About to patch preload libraries')
+
+    system.checkpoint('system patched...')
+
+def run_with_retrace(system, argv, trace_shutdown = False):
+
+    def runpy_exec(source, globals = None, locals = None):
+        with system.thread_state.select('internal'):
+            return builtins.exec(source, globals, locals)
+    
+    utils.update(runpy, "_run_code",
+                 utils.wrap_func_with_overrides,
+                 exec = runpy_exec)
+
+    try:
+        run_python_command(argv)
+    finally:
+        wait_for_non_daemon_threads()
+        try:
+            if trace_shutdown:
+                with system.thread_state.select('internal'):
+                    atexit._run_exitfuncs()
+            else:
+                atexit._run_exitfuncs()
+        except Exception as e:
+            print(f"Error in atexit hook: {e}", file=sys.stderr)