Update tracer.py

Author: KRRT7

codeflash/tracer.py

@@ -11,6 +11,7 @@
 #
 from __future__ import annotations
 
+import contextlib
 import importlib.machinery
 import io
 import json
@@ -42,7 +43,6 @@
 from codeflash.tracing.replay_test import create_trace_replay_test
 from codeflash.tracing.tracing_utils import FunctionModules
 from codeflash.verification.verification_utils import get_test_file_path
-import contextlib
 
 if TYPE_CHECKING:
     from types import FrameType, TracebackType
@@ -77,7 +77,7 @@ def __init__(
         self,
         output: str = "codeflash.trace",
         functions: list[str] | None = None,
-        disable: bool = False,  # noqa: FBT001, FBT002
+        disable: bool = False,
         config_file_path: Path | None = None,
         max_function_count: int = 256,
         timeout: int | None = None,  # seconds
@@ -100,6 +100,7 @@ def __init__(
             )
             disable = True
         self.disable = disable
+        self._db_lock: threading.Lock | None = None
         if self.disable:
             return
         if sys.getprofile() is not None or sys.gettrace() is not None:
@@ -109,6 +110,9 @@ def __init__(
             )
             self.disable = True
             return
+
+        self._db_lock = threading.Lock()
+
         self.con = None
         self.output_file = Path(output).resolve()
         self.functions = functions
@@ -180,34 +184,55 @@ def __enter__(self) -> None:
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
-        if self.disable:
+        if self.disable or self._db_lock is None:
             return
         sys.setprofile(None)
-        self.con.commit()
-        console.rule("Codeflash: Traced Program Output End", style="bold blue")
-        self.create_stats()
+        threading.setprofile(None)
 
-        cur = self.con.cursor()
-        cur.execute(
-            "CREATE TABLE pstats (filename TEXT, line_number INTEGER, function TEXT, class_name TEXT, "
-            "call_count_nonrecursive INTEGER, num_callers INTEGER, total_time_ns INTEGER, "
-            "cumulative_time_ns INTEGER, callers BLOB)"
-        )
-        for func, (cc, nc, tt, ct, callers) in self.stats.items():
-            remapped_callers = [{"key": k, "value": v} for k, v in callers.items()]
+        with self._db_lock:
+            if self.con is None:
+                return
+
+            self.con.commit()  # Commit any pending from tracer_logic
+            console.rule("Codeflash: Traced Program Output End", style="bold blue")
+            self.create_stats()  # This calls snapshot_stats which uses self.timings
+
+            cur = self.con.cursor()
             cur.execute(
-                "INSERT INTO pstats VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
-                (str(Path(func[0]).resolve()), func[1], func[2], func[3], cc, nc, tt, ct, json.dumps(remapped_callers)),
+                "CREATE TABLE pstats (filename TEXT, line_number INTEGER, function TEXT, class_name TEXT, "
+                "call_count_nonrecursive INTEGER, num_callers INTEGER, total_time_ns INTEGER, "
+                "cumulative_time_ns INTEGER, callers BLOB)"
             )
-        self.con.commit()
+            # self.stats is populated by snapshot_stats() called within create_stats()
+            # Ensure self.stats is accessed after create_stats() and within the lock if it involves DB data
+            # For now, assuming self.stats is primarily in-memory after create_stats()
+            for func, (cc, nc, tt, ct, callers) in self.stats.items():
+                remapped_callers = [{"key": k, "value": v} for k, v in callers.items()]
+                cur.execute(
+                    "INSERT INTO pstats VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                    (
+                        str(Path(func[0]).resolve()),
+                        func[1],
+                        func[2],
+                        func[3],
+                        cc,
+                        nc,
+                        tt,
+                        ct,
+                        json.dumps(remapped_callers),
+                    ),
+                )
+            self.con.commit()
 
-        self.make_pstats_compatible()
-        self.print_stats("tottime")
-        cur = self.con.cursor()
-        cur.execute("CREATE TABLE total_time (time_ns INTEGER)")
-        cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,))
-        self.con.commit()
-        self.con.close()
+            self.make_pstats_compatible()  # Modifies self.stats and self.timings in-memory
+            self.print_stats("tottime")  # Uses self.stats, prints to console
+
+            cur = self.con.cursor()  # New cursor
+            cur.execute("CREATE TABLE total_time (time_ns INTEGER)")
+            cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,))
+            self.con.commit()
+            self.con.close()
+            self.con = None  # Mark connection as closed
 
         # filter any functions where we did not capture the return
         self.function_modules = [
@@ -244,16 +269,24 @@ def __exit__(
             overflow="ignore",
         )
 
-    def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
+    def tracer_logic(self, frame: FrameType, event: str) -> None:
         if event != "call":
             return
         if None is not self.timeout and (time.time() - self.start_time) > self.timeout:
             sys.setprofile(None)
             threading.setprofile(None)
             console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.")
             return
+        if self.disable or self._db_lock is None or self.con is None:
+            return
+
         code = frame.f_code
 
+        # Check function name first before resolving path
+        if code.co_name in self.ignored_functions:
+            return
+
+        # Now resolve file path only if we need it
         co_filename = code.co_filename
         if co_filename in self.path_cache:
             file_name = self.path_cache[co_filename]
@@ -262,8 +295,6 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
             self.path_cache[co_filename] = file_name
         # TODO : It currently doesn't log the last return call from the first function
 
-        if code.co_name in self.ignored_functions:
-            return
         if not file_name.is_relative_to(self.project_root):
             return
         if not file_name.exists():
@@ -290,7 +321,12 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
         except:  # noqa: E722
             # someone can override the getattr method and raise an exception. I'm looking at you wrapt
             return
-        function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}"
+
+        try:
+            function_qualified_name = f"{file_name}:{code.co_qualname}"
+        except AttributeError:
+            function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}"
+
         if function_qualified_name in self.ignored_qualified_functions:
             return
         if function_qualified_name not in self.function_count:
@@ -323,49 +359,59 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
 
         # TODO: Also check if this function arguments are unique from the values logged earlier
 
-        cur = self.con.cursor()
+        with self._db_lock:
+            # Check connection again inside lock, in case __exit__ closed it.
+            if self.con is None:
+                return
 
-        t_ns = time.perf_counter_ns()
-        original_recursion_limit = sys.getrecursionlimit()
-        try:
-            # pickling can be a recursive operator, so we need to increase the recursion limit
-            sys.setrecursionlimit(10000)
-            # We do not pickle self for __init__ to avoid recursion errors, and instead instantiate its class
-            # directly with the rest of the arguments in the replay tests. We copy the arguments to avoid memory
-            # leaks, bad references or side effects when unpickling.
-            arguments = dict(arguments.items())
-            if class_name and code.co_name == "__init__":
-                del arguments["self"]
-            local_vars = pickle.dumps(arguments, protocol=pickle.HIGHEST_PROTOCOL)
-            sys.setrecursionlimit(original_recursion_limit)
-        except (TypeError, pickle.PicklingError, AttributeError, RecursionError, OSError):
-            # we retry with dill if pickle fails. It's slower but more comprehensive
+            cur = self.con.cursor()
+
+            t_ns = time.perf_counter_ns()
+            original_recursion_limit = sys.getrecursionlimit()
             try:
-                local_vars = dill.dumps(arguments, protocol=dill.HIGHEST_PROTOCOL)
+                # pickling can be a recursive operator, so we need to increase the recursion limit
+                sys.setrecursionlimit(10000)
+                # We do not pickle self for __init__ to avoid recursion errors, and instead instantiate its class
+                # directly with the rest of the arguments in the replay tests. We copy the arguments to avoid memory
+                # leaks, bad references or side effects when unpickling.
+                arguments_copy = dict(arguments.items())  # Use the local 'arguments' from frame.f_locals
+                if class_name and code.co_name == "__init__" and "self" in arguments_copy:
+                    del arguments_copy["self"]
+                local_vars = pickle.dumps(arguments_copy, protocol=pickle.HIGHEST_PROTOCOL)
                 sys.setrecursionlimit(original_recursion_limit)
+            except (TypeError, pickle.PicklingError, AttributeError, RecursionError, OSError):
+                # we retry with dill if pickle fails. It's slower but more comprehensive
+                try:
+                    sys.setrecursionlimit(10000)  # Ensure limit is high for dill too
+                    # arguments_copy should be used here as well if defined above
+                    local_vars = dill.dumps(
+                        arguments_copy if "arguments_copy" in locals() else dict(arguments.items()),
+                        protocol=dill.HIGHEST_PROTOCOL,
+                    )
+                    sys.setrecursionlimit(original_recursion_limit)
+
+                except (TypeError, dill.PicklingError, AttributeError, RecursionError, OSError):
+                    self.function_count[function_qualified_name] -= 1
+                    return
 
-            except (TypeError, dill.PicklingError, AttributeError, RecursionError, OSError):
-                # give up
-                self.function_count[function_qualified_name] -= 1
-                return
-        cur.execute(
-            "INSERT INTO function_calls VALUES(?, ?, ?, ?, ?, ?, ?, ?)",
-            (
-                event,
-                code.co_name,
-                class_name,
-                str(file_name),
-                frame.f_lineno,
-                frame.f_back.__hash__(),
-                t_ns,
-                local_vars,
-            ),
-        )
-        self.trace_count += 1
-        self.next_insert -= 1
-        if self.next_insert == 0:
-            self.next_insert = 1000
-            self.con.commit()
+            cur.execute(
+                "INSERT INTO function_calls VALUES(?, ?, ?, ?, ?, ?, ?, ?)",
+                (
+                    event,
+                    code.co_name,
+                    class_name,
+                    str(file_name),
+                    frame.f_lineno,
+                    frame.f_back.__hash__(),
+                    t_ns,
+                    local_vars,
+                ),
+            )
+            self.trace_count += 1
+            self.next_insert -= 1
+            if self.next_insert == 0:
+                self.next_insert = 1000
+                self.con.commit()
 
     def trace_callback(self, frame: FrameType, event: str, arg: str | None) -> None:
         # profiler section
@@ -413,7 +459,7 @@ def trace_dispatch_call(self, frame: FrameType, t: int) -> int:
                     class_name = arguments["self"].__class__.__name__
                 elif "cls" in arguments and hasattr(arguments["cls"], "__name__"):
                     class_name = arguments["cls"].__name__
-            except Exception:  # noqa: S110
+            except Exception:  # noqa: BLE001, S110
                 pass
 
             fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name)
@@ -425,7 +471,7 @@ def trace_dispatch_call(self, frame: FrameType, t: int) -> int:
             else:
                 timings[fn] = 0, 0, 0, 0, {}
             return 1  # noqa: TRY300
-        except Exception:
+        except Exception:  # noqa: BLE001
             # Handle any errors gracefully
             return 0
 
@@ -488,8 +534,9 @@ def trace_dispatch_return(self, frame: FrameType, t: int) -> int:
             cc = cc + 1
 
         if pfn in callers:
-            callers[pfn] = callers[pfn] + 1  # TODO: gather more
-            # stats such as the amount of time added to ct courtesy
+            # Increment call count between these functions
+            callers[pfn] = callers[pfn] + 1
+            # Note: This tracks stats such as the amount of time added to ct
             # of this specific call, and the contribution to cc
             # courtesy of this call.
         else:
@@ -579,7 +626,7 @@ def print_stats(self, sort: str | int | tuple = -1) -> None:
 
                     # Store with new format
                     new_stats[new_func] = (cc, nc, tt, ct, new_callers)
-                except Exception as e:
+                except Exception as e:  # noqa: BLE001
                     console.print(f"Error converting stats for {func}: {e}")
                     continue
 
@@ -616,7 +663,7 @@ def print_stats(self, sort: str | int | tuple = -1) -> None:
                         new_callers[new_caller_func] = count
 
                     new_timings[new_func] = (cc, ns, tt, ct, new_callers)
-                except Exception as e:
+                except Exception as e:  # noqa: BLE001
                     console.print(f"Error converting timings for {func}: {e}")
                     continue
 
@@ -686,7 +733,7 @@ def print_stats(self, sort: str | int | tuple = -1) -> None:
 
             console.print(Align.center(table))
 
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             console.print(f"[bold red]Error in stats processing:[/bold red] {e}")
             console.print(f"Traced {self.trace_count:,} function calls")
             self.total_tt = 0
@@ -716,7 +763,7 @@ def create_stats(self) -> None:
 
     def snapshot_stats(self) -> None:
         self.stats = {}
-        for func, (cc, _ns, tt, ct, caller_dict) in self.timings.items():
+        for func, (cc, _ns, tt, ct, caller_dict) in list(self.timings.items()):
             callers = caller_dict.copy()
             nc = 0
             for callcnt in callers.values():