Fix non-deterministic caller resolution in call graph tracking

src/fever/call_tracker.py

@@ -114,6 +114,7 @@ def __init__(
         self.resume_event = threading.Event()
         self.stop_event = threading.Event()
         self._propagate_trace_on_cache_hit = propagate_trace_on_cache_hit
+        self._execution_stack: List[TraceNode] = []
 
     def track_calls(
         self,
@@ -191,16 +192,24 @@ def fever_wrapper(*args, **kwargs):
                 if key == params.hash:
                     params = data["params"]
                     break
-            # INFO: We don't know the caller's parameters, but they are in the graph
-            # somewhere. For now, assuming a single thread, we can connect the caller's
-            # node to the previously registered node for that function. This is britle
-            # though.
-            # FIXME: This will break if we go multithreaded:
-            # FIXME: Include the module too!!
+            # INFO: Use the execution stack to find the correct caller node.
+            # The stack contains the current call chain, so the top of the stack
+            # is the immediate caller.
             caller_params_hash = -1
-            for n in self._call_graph.nodes:
-                if n.func == k:
-                    caller_params_hash = n.params_hash
+            if self._execution_stack:
+                # The top of the stack is the current caller
+                caller_node = self._execution_stack[-1]
+                if caller_node.func == k and caller_node.module == caller_module:
+                    caller_params_hash = caller_node.params_hash
+
+            # If we couldn't find it in the stack, fall back to searching the graph
+            # (this handles cases where the caller isn't tracked)
+            if caller_params_hash == -1:
+                for n in self._call_graph.nodes:
+                    if n.func == k and n.module == caller_module:
+                        caller_params_hash = n.params_hash
+                        break
+
             k, v = (
                 TraceNode(caller_module, k, caller_params_hash),
                 TraceNode(module.name, v, params.hash),
@@ -276,41 +285,57 @@ def fever_wrapper(*args, **kwargs):
                                         )
                                     )
                 return cached_result
+
+            # Push the current callee node onto the execution stack
+            self._execution_stack.append(v)
+
             start = timeit.default_timer()
+            exception_occurred = False
             try:
-                result = func_ptr(*args, **kwargs)
-            except Exception as e:
-                self._on_exception(e)
-                # Wait for resume, but check stop_event periodically
-                while not self.resume_event.is_set():
-                    if self.stop_event.is_set():
-                        log.debug(
-                            "Stop event detected while waiting on exception, terminating thread"
+                try:
+                    result = func_ptr(*args, **kwargs)
+                except Exception as e:
+                    exception_occurred = True
+                    self._on_exception(e)
+                    # Wait for resume, but check stop_event periodically
+                    while not self.resume_event.is_set():
+                        if self.stop_event.is_set():
+                            log.debug(
+                                "Stop event detected while waiting on exception, terminating thread"
+                            )
+                            raise SystemExit("Thread termination requested")
+                        self.resume_event.wait(timeout=0.1)
+                    # After resuming, return None since the function didn't complete successfully
+                    result = None
+                end = timeit.default_timer()
+                log.debug(f"Call to '{callable_full_name}' took {end - start:.6f} seconds")
+                self._on_new_call(k, v)
+                # WARN: The caller object will change as the caller function is recompiled!
+                # Because we look for it in the call stack. This is normal, but we might
+                # want the caller to be the function name instead of the pointer, so we
+                # parameterize the strategy with self._tracking_mode.
+                edge_data = self._call_graph.edges[k, v, params.hash]
+                if "weight" not in edge_data:
+                    edge_data["cum_time"] = 0.0
+                    edge_data["calls"] = 0
+                edge_data["cum_time"] += end - start
+                edge_data["calls"] += 1
+                edge_data["weight"] = edge_data["cum_time"] / edge_data["calls"]
+                edge_data["last_timestamp"] = start
+
+                # Only cache the result if the function executed successfully
+                if not exception_occurred:
+                    try:
+                        self._cache.set(func_ptr, params, edge_data, result)
+                    except Exception as e:
+                        log.error(
+                            f"Error setting cache for {func_ptr} with params {params}: {e}"
                         )
-                        raise SystemExit("Thread termination requested")
-                    self.resume_event.wait(timeout=0.1)
-            end = timeit.default_timer()
-            log.debug(f"Call to '{callable_full_name}' took {end - start:.6f} seconds")
-            self._on_new_call(k, v)
-            # WARN: The caller object will change as the caller function is recompiled!
-            # Because we look for it in the call stack. This is normal, but we might
-            # want the caller to be the function name instead of the pointer, so we
-            # parameterize the strategy with self._tracking_mode.
-            edge_data = self._call_graph.edges[k, v, params.hash]
-            if "weight" not in edge_data:
-                edge_data["cum_time"] = 0.0
-                edge_data["calls"] = 0
-            edge_data["cum_time"] += end - start
-            edge_data["calls"] += 1
-            edge_data["weight"] = edge_data["cum_time"] / edge_data["calls"]
-            edge_data["last_timestamp"] = start
-            try:
-                self._cache.set(func_ptr, params, edge_data, result)
-            except Exception as e:
-                log.error(
-                    f"Error setting cache for {func_ptr} with params {params}: {e}"
-                )
-            return result
+
+                return result
+            finally:
+                # Always pop the callee node from the execution stack
+                self._execution_stack.pop()
 
         return fever_wrapper
 

tests/test_imports/module_a.py

@@ -6,24 +6,52 @@
 def function() -> str:
     print("test.function()")
     test_module = TestCase()
-    return test_module.elaborate_function("world", [False, False, False])
+    return test_module.elaborate_function("earth", [False, False, False])
 
 
 def function_deep_nested() -> str:
     print("test.function_deep_nested()")
     test_module = TestCase()
     test_module()
-    return test_module.elaborate_function("world", [False, False, False])
+    return test_module.elaborate_function("earth", [False, False, False])
 
 
 def function_with_lambda_call() -> str:
     print("test.function_with_lambda_call()")
     test_module = TestCase()
-    return test_module.elaborate_function("world", [True, False, True])
+    return test_module.elaborate_function("earth", [True, False, True])
 
 
 def second_function():
     print("im an other function")
     for _ in range(random.randint(3, 7)):
         module_level_func("second_function")
     other_module_level_func("ya")
+
+
+global_var_a = 10
+global_var_b = 20
+
+def function_with_globals() -> int:
+    return global_var_a + global_var_b
+
+
+import numpy as np
+
+def function_with_numpy() -> np.ndarray:
+    a = np.array([4, 5, 6])
+    return a + 1
+
+
+from numpy import array
+
+def function_with_numpy_array() -> array:
+    a = array([7, 8, 9])
+    return a + 2
+
+
+def new_fn_a():
+    return 123
+
+def new_fn_b():
+    return new_fn_a()+1

tests/test_imports/module_c.py

@@ -6,3 +6,10 @@ def other_function(from_module: str) -> bool:
     module_a.second_function()
     print("new statement!")
     return True
+
+
+class TestClass:
+    def call_me_baby(self):
+        return "Hey, I just met you!"
+def new_function() -> str:
+    return "I think therefore I am"

tests/test_imports/module_d.py

@@ -3,7 +3,7 @@
 
 def function_d(name: str) -> int:
     print(f"Nothing to show here, {name}")
-    return 123
+    return 456
 
 
 def other_function_d(name: str) -> int:
@@ -28,7 +28,7 @@ def hello(self, name: str):
         print(f"hello {a(name)}")
 
     def return_string(self) -> str:
-        return "just a string"
+        return "not just a string"
 
     class NestedTestClass:
         def __init__(self, owner):
@@ -41,13 +41,13 @@ def nested_test(self, name: str) -> int:
 
 class MiniTestClass:
     def __init__(self, name: str):
-        self._name = name
+        self._name = 'testy girl'
 
     def __call__(self) -> str:
-        return self._name
+        return 'oh la la!'
 
     def __len__(self) -> int:
-        return 10
+        return 20
 
     def __str__(self) -> str:
         return f"My name is {self._name}"

tests/test_imports/submodules/module_e.py

@@ -5,15 +5,15 @@ def function_e():
 
 def function_foreign_imports() -> np.ndarray:
     x = np.array([1, 2, 3])
-    return x * 2
+    return x * 3
 
 
 def nested_functions() -> int:
     def nested_a() -> str:
         def nested_b() -> int:
             return 123
 
-        return f"nested_a calls nested_b: {nested_b()}"
+        return f"nested_a calls modified nested_b: {nested_b()}"
 
     return len(nested_a())