import itertools import json import os import platform import sys import threading import time from .collector import Collector, filter_internal_frames from .opcode_utils import get_opcode_info, format_opcode try: from _remote_debugging import THREAD_STATUS_HAS_GIL, THREAD_STATUS_ON_CPU, THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION, THREAD_STATUS_MAIN_THREAD except ImportError: # Fallback if module not available (shouldn't happen in normal use) THREAD_STATUS_HAS_GIL = (1 << 0) THREAD_STATUS_ON_CPU = (1 << 1) THREAD_STATUS_UNKNOWN = (1 << 2) THREAD_STATUS_GIL_REQUESTED = (1 << 3) THREAD_STATUS_HAS_EXCEPTION = (1 << 4) THREAD_STATUS_MAIN_THREAD = (1 << 5) # Categories matching Firefox Profiler expectations GECKO_CATEGORIES = [ {"name": "Other", "color": "grey", "subcategories": ["Other"]}, {"name": "Python", "color": "yellow", "subcategories": ["Other"]}, {"name": "Native", "color": "blue", "subcategories": ["Other"]}, {"name": "GC", "color": "orange", "subcategories": ["Other"]}, {"name": "GIL", "color": "green", "subcategories": ["Other"]}, {"name": "CPU", "color": "purple", "subcategories": ["Other"]}, {"name": "Code Type", "color": "red", "subcategories": ["Other"]}, {"name": "Opcodes", "color": "magenta", "subcategories": ["Other"]}, {"name": "Exception", "color": "lightblue", "subcategories": ["Other"]}, ] # Category indices CATEGORY_OTHER = 0 CATEGORY_PYTHON = 1 CATEGORY_NATIVE = 2 CATEGORY_GC = 3 CATEGORY_GIL = 4 CATEGORY_CPU = 5 CATEGORY_CODE_TYPE = 6 CATEGORY_OPCODES = 7 CATEGORY_EXCEPTION = 8 # Subcategory indices DEFAULT_SUBCATEGORY = 0 GECKO_FORMAT_VERSION = 32 GECKO_PREPROCESSED_VERSION = 57 # Resource type constants RESOURCE_TYPE_LIBRARY = 1 # Frame constants FRAME_ADDRESS_NONE = -1 FRAME_INLINE_DEPTH_ROOT = 0 # Process constants PROCESS_TYPE_MAIN = 0 STACKWALK_DISABLED = 0 class GeckoCollector(Collector): def __init__(self, sample_interval_usec, *, skip_idle=False, opcodes=False): self.sample_interval_usec = sample_interval_usec self.skip_idle = skip_idle self.opcodes_enabled = opcodes self.start_time = time.monotonic() * 1000 # milliseconds since start # Global string table (shared across all threads) self.global_strings = ["(root)"] # Start with root self.global_string_map = {"(root)": 0} # Per-thread data structures self.threads = {} # tid -> thread data # Global tables self.libs = [] # Sampling interval tracking self.sample_count = 0 self.last_sample_time = 0 self.interval = 1.0 # Will be calculated from actual sampling # State tracking for interval markers (tid -> start_time) self.has_gil_start = {} # Thread has the GIL self.no_gil_start = {} # Thread doesn't have the GIL self.on_cpu_start = {} # Thread is running on CPU self.off_cpu_start = {} # Thread is off CPU self.python_code_start = {} # Thread running Python code (has GIL) self.native_code_start = {} # Thread running native code (on CPU without GIL) self.gil_wait_start = {} # Thread waiting for GIL self.exception_start = {} # Thread has an exception set self.no_exception_start = {} # Thread has no exception set # GC event tracking: track GC start time per thread self.gc_start_per_thread = {} # tid -> start_time # Track which threads have been initialized for state tracking self.initialized_threads = set() # Opcode state tracking per thread: tid -> (opcode, lineno, col_offset, funcname, filename, start_time) self.opcode_state = {} # For binary replay: track base timestamp (first sample's timestamp) self._replay_base_timestamp_us = None def _track_state_transition(self, tid, condition, active_dict, inactive_dict, active_name, inactive_name, category, current_time): """Track binary state transitions and emit markers. Args: tid: Thread ID condition: Whether the active state is true active_dict: Dict tracking start time of active state inactive_dict: Dict tracking start time of inactive state active_name: Name for active state marker inactive_name: Name for inactive state marker category: Gecko category for the markers current_time: Current timestamp """ # On first observation of a thread, just record the current state # without creating a marker (we don't know what the previous state was) if tid not in self.initialized_threads: if condition: active_dict[tid] = current_time else: inactive_dict[tid] = current_time return # For already-initialized threads, track transitions if condition: active_dict.setdefault(tid, current_time) if tid in inactive_dict: self._add_marker(tid, inactive_name, inactive_dict.pop(tid), current_time, category) else: inactive_dict.setdefault(tid, current_time) if tid in active_dict: self._add_marker(tid, active_name, active_dict.pop(tid), current_time, category) def collect(self, stack_frames, timestamps_us=None): """Collect samples from stack frames. Args: stack_frames: List of interpreter/thread frame info timestamps_us: List of timestamps in microseconds (None for live sampling) """ # Handle live sampling (no timestamps provided) if timestamps_us is None: current_time = (time.monotonic() * 1000) - self.start_time times = [current_time] else: if not timestamps_us: return # Initialize base timestamp if needed if self._replay_base_timestamp_us is None: self._replay_base_timestamp_us = timestamps_us[0] # Convert all timestamps to times (ms relative to first sample) base = self._replay_base_timestamp_us times = [(ts - base) / 1000 for ts in timestamps_us] first_time = times[0] # Update interval calculation if self.sample_count > 0 and self.last_sample_time > 0: self.interval = (times[-1] - self.last_sample_time) / self.sample_count self.last_sample_time = times[-1] # Process threads for interpreter_info in stack_frames: for thread_info in interpreter_info.threads: frames = filter_internal_frames(thread_info.frame_info) tid = thread_info.thread_id status_flags = thread_info.status is_main_thread = bool(status_flags & THREAD_STATUS_MAIN_THREAD) # Initialize thread if needed if tid not in self.threads: self.threads[tid] = self._create_thread(tid, is_main_thread) thread_data = self.threads[tid] # Decode status flags has_gil = bool(status_flags & THREAD_STATUS_HAS_GIL) on_cpu = bool(status_flags & THREAD_STATUS_ON_CPU) gil_requested = bool(status_flags & THREAD_STATUS_GIL_REQUESTED) # Track state transitions using first timestamp self._track_state_transition( tid, has_gil, self.has_gil_start, self.no_gil_start, "Has GIL", "No GIL", CATEGORY_GIL, first_time ) self._track_state_transition( tid, on_cpu, self.on_cpu_start, self.off_cpu_start, "On CPU", "Off CPU", CATEGORY_CPU, first_time ) # Track code type if has_gil: self._track_state_transition( tid, True, self.python_code_start, self.native_code_start, "Python Code", "Native Code", CATEGORY_CODE_TYPE, first_time ) elif on_cpu: self._track_state_transition( tid, True, self.native_code_start, self.python_code_start, "Native Code", "Python Code", CATEGORY_CODE_TYPE, first_time ) else: if tid in self.initialized_threads: if tid in self.python_code_start: self._add_marker(tid, "Python Code", self.python_code_start.pop(tid), first_time, CATEGORY_CODE_TYPE) if tid in self.native_code_start: self._add_marker(tid, "Native Code", self.native_code_start.pop(tid), first_time, CATEGORY_CODE_TYPE) # Track GIL wait if gil_requested: self.gil_wait_start.setdefault(tid, first_time) elif tid in self.gil_wait_start: self._add_marker(tid, "Waiting for GIL", self.gil_wait_start.pop(tid), first_time, CATEGORY_GIL) # Track exception state has_exception = bool(status_flags & THREAD_STATUS_HAS_EXCEPTION) self._track_state_transition( tid, has_exception, self.exception_start, self.no_exception_start, "Has Exception", "No Exception", CATEGORY_EXCEPTION, first_time ) # Track GC events has_gc_frame = any(frame[2] == "" for frame in frames) if has_gc_frame: if tid not in self.gc_start_per_thread: self.gc_start_per_thread[tid] = first_time elif tid in self.gc_start_per_thread: self._add_marker(tid, "GC Collecting", self.gc_start_per_thread.pop(tid), first_time, CATEGORY_GC) # Mark thread as initialized self.initialized_threads.add(tid) # Skip idle threads if requested is_idle = not has_gil and not on_cpu if self.skip_idle and is_idle: continue if not frames: continue # Process stack once to get stack_index stack_index = self._process_stack(thread_data, frames) # Add samples with timestamps samples = thread_data["samples"] samples_stack = samples["stack"] samples_time = samples["time"] samples_delay = samples["eventDelay"] for t in times: samples_stack.append(stack_index) samples_time.append(t) samples_delay.append(None) # Handle opcodes if self.opcodes_enabled and frames: leaf_frame = frames[0] filename, location, funcname, opcode = leaf_frame if isinstance(location, tuple): lineno, _, col_offset, _ = location else: lineno = location col_offset = -1 current_state = (opcode, lineno, col_offset, funcname, filename) if tid not in self.opcode_state: self.opcode_state[tid] = (*current_state, first_time) elif self.opcode_state[tid][:5] != current_state: prev_opcode, prev_lineno, prev_col, prev_funcname, prev_filename, prev_start = self.opcode_state[tid] self._add_opcode_interval_marker( tid, prev_opcode, prev_lineno, prev_col, prev_funcname, prev_start, first_time ) self.opcode_state[tid] = (*current_state, first_time) self.sample_count += len(times) def _create_thread(self, tid, is_main_thread): """Create a new thread structure with processed profile format.""" thread = { "name": f"Thread-{tid}", "isMainThread": is_main_thread, "processStartupTime": 0, "processShutdownTime": None, "registerTime": 0, "unregisterTime": None, "pausedRanges": [], "pid": str(os.getpid()), "tid": tid, "processType": "default", "processName": "Python Process", # Sample data - processed format with direct arrays "samples": { "stack": [], "time": [], "eventDelay": [], "weight": None, "weightType": "samples", "length": 0, # Will be updated on export }, # Stack table - processed format "stackTable": { "frame": [], "category": [], "subcategory": [], "prefix": [], "length": 0, # Will be updated on export }, # Frame table - processed format "frameTable": { "address": [], "category": [], "subcategory": [], "func": [], "innerWindowID": [], "implementation": [], "optimizations": [], "line": [], "column": [], "inlineDepth": [], "nativeSymbol": [], "length": 0, # Will be updated on export }, # Function table - processed format "funcTable": { "name": [], "isJS": [], "relevantForJS": [], "resource": [], "fileName": [], "lineNumber": [], "columnNumber": [], "length": 0, # Will be updated on export }, # Resource table - processed format "resourceTable": { "lib": [], "name": [], "host": [], "type": [], "length": 0, # Will be updated on export }, # Native symbols table (empty for Python) "nativeSymbols": { "libIndex": [], "address": [], "name": [], "functionSize": [], "length": 0, }, # Markers - processed format (arrays) "markers": { "data": [], "name": [], "startTime": [], "endTime": [], "phase": [], "category": [], "length": 0, }, # Caches for deduplication "_stackCache": {}, "_frameCache": {}, "_funcCache": {}, "_resourceCache": {}, } return thread def _is_python(self, filename: str) -> bool: return not filename.startswith("<") or filename in ["", ""] def _get_category(self, filename: str) -> int: return CATEGORY_PYTHON if self._is_python(filename) else CATEGORY_NATIVE def _intern_string(self, s): """Intern a string in the global string table.""" if s in self.global_string_map: return self.global_string_map[s] idx = len(self.global_strings) self.global_strings.append(s) self.global_string_map[s] = idx return idx def _add_marker(self, tid, name, start_time, end_time, category): """Add an interval marker for a specific thread.""" if tid not in self.threads: return thread_data = self.threads[tid] duration = end_time - start_time name_idx = self._intern_string(name) markers = thread_data["markers"] markers["name"].append(name_idx) markers["startTime"].append(start_time) markers["endTime"].append(end_time) markers["phase"].append(1) # 1 = interval marker markers["category"].append(category) markers["data"].append({ "type": name.replace(" ", ""), "duration": duration, "tid": tid }) def _add_opcode_interval_marker(self, tid, opcode, lineno, col_offset, funcname, start_time, end_time): """Add an interval marker for opcode execution span.""" if tid not in self.threads or opcode is None: return thread_data = self.threads[tid] opcode_info = get_opcode_info(opcode) # Use formatted opcode name (with base opcode for specialized ones) formatted_opname = format_opcode(opcode) name_idx = self._intern_string(formatted_opname) markers = thread_data["markers"] markers["name"].append(name_idx) markers["startTime"].append(start_time) markers["endTime"].append(end_time) markers["phase"].append(1) # 1 = interval marker markers["category"].append(CATEGORY_OPCODES) markers["data"].append({ "type": "Opcode", "opcode": opcode, "opname": formatted_opname, "base_opname": opcode_info["base_opname"], "is_specialized": opcode_info["is_specialized"], "line": lineno, "column": col_offset if col_offset >= 0 else None, "function": funcname, "duration": end_time - start_time, }) def _process_stack(self, thread_data, frames): """Process a stack and return the stack index.""" if not frames: return None # Cache references to avoid repeated dictionary lookups stack_cache = thread_data["_stackCache"] stack_table = thread_data["stackTable"] stack_frames = stack_table["frame"] stack_prefix = stack_table["prefix"] stack_category = stack_table["category"] stack_subcategory = stack_table["subcategory"] # Build stack bottom-up (from root to leaf) prefix_stack_idx = None for frame_tuple in reversed(frames): # frame_tuple is (filename, location, funcname, opcode) # location is (lineno, end_lineno, col_offset, end_col_offset) or just lineno filename, location, funcname, opcode = frame_tuple if isinstance(location, tuple): lineno, end_lineno, col_offset, end_col_offset = location else: # Legacy format: location is just lineno lineno = location col_offset = -1 end_col_offset = -1 # Get or create function func_idx = self._get_or_create_func( thread_data, filename, funcname, lineno ) # Get or create frame (include column for precise source location) frame_idx = self._get_or_create_frame( thread_data, func_idx, lineno, col_offset ) # Check stack cache stack_key = (frame_idx, prefix_stack_idx) if stack_key in stack_cache: prefix_stack_idx = stack_cache[stack_key] else: # Create new stack entry stack_idx = len(stack_frames) stack_frames.append(frame_idx) stack_prefix.append(prefix_stack_idx) # Determine category category = self._get_category(filename) stack_category.append(category) stack_subcategory.append(DEFAULT_SUBCATEGORY) stack_cache[stack_key] = stack_idx prefix_stack_idx = stack_idx return prefix_stack_idx def _get_or_create_func(self, thread_data, filename, funcname, lineno): """Get or create a function entry.""" func_cache = thread_data["_funcCache"] func_key = (filename, funcname) if func_key in func_cache: return func_cache[func_key] # Cache references for func table func_table = thread_data["funcTable"] func_names = func_table["name"] func_is_js = func_table["isJS"] func_relevant = func_table["relevantForJS"] func_resources = func_table["resource"] func_filenames = func_table["fileName"] func_line_numbers = func_table["lineNumber"] func_column_numbers = func_table["columnNumber"] func_idx = len(func_names) # Intern strings in global table name_idx = self._intern_string(funcname) # Determine if Python is_python = self._is_python(filename) # Create resource resource_idx = self._get_or_create_resource(thread_data, filename) # Add function func_names.append(name_idx) func_is_js.append(is_python) func_relevant.append(is_python) func_resources.append(resource_idx) if is_python: filename_idx = self._intern_string(os.path.basename(filename)) func_filenames.append(filename_idx) func_line_numbers.append(lineno) else: func_filenames.append(None) func_line_numbers.append(None) func_column_numbers.append(None) func_cache[func_key] = func_idx return func_idx def _get_or_create_resource(self, thread_data, filename): """Get or create a resource entry.""" resource_cache = thread_data["_resourceCache"] if filename in resource_cache: return resource_cache[filename] # Cache references for resource table resource_table = thread_data["resourceTable"] resource_libs = resource_table["lib"] resource_names = resource_table["name"] resource_hosts = resource_table["host"] resource_types = resource_table["type"] resource_idx = len(resource_names) resource_name = ( os.path.basename(filename) if "/" in filename else filename ) name_idx = self._intern_string(resource_name) resource_libs.append(None) resource_names.append(name_idx) resource_hosts.append(None) resource_types.append(RESOURCE_TYPE_LIBRARY) resource_cache[filename] = resource_idx return resource_idx def _get_or_create_frame(self, thread_data, func_idx, lineno, col_offset=-1): """Get or create a frame entry.""" frame_cache = thread_data["_frameCache"] # Include column in cache key for precise frame identification frame_key = (func_idx, lineno, col_offset if col_offset >= 0 else None) if frame_key in frame_cache: return frame_cache[frame_key] # Cache references for frame table frame_table = thread_data["frameTable"] frame_addresses = frame_table["address"] frame_inline_depths = frame_table["inlineDepth"] frame_categories = frame_table["category"] frame_subcategories = frame_table["subcategory"] frame_funcs = frame_table["func"] frame_native_symbols = frame_table["nativeSymbol"] frame_inner_window_ids = frame_table["innerWindowID"] frame_implementations = frame_table["implementation"] frame_lines = frame_table["line"] frame_columns = frame_table["column"] frame_optimizations = frame_table["optimizations"] frame_idx = len(frame_funcs) # Determine category based on function - use cached func table reference is_python = thread_data["funcTable"]["isJS"][func_idx] category = CATEGORY_PYTHON if is_python else CATEGORY_NATIVE frame_addresses.append(FRAME_ADDRESS_NONE) frame_inline_depths.append(FRAME_INLINE_DEPTH_ROOT) frame_categories.append(category) frame_subcategories.append(DEFAULT_SUBCATEGORY) frame_funcs.append(func_idx) frame_native_symbols.append(None) frame_inner_window_ids.append(None) frame_implementations.append(None) frame_lines.append(lineno if lineno else None) # Store column offset if available (>= 0), otherwise None frame_columns.append(col_offset if col_offset >= 0 else None) frame_optimizations.append(None) frame_cache[frame_key] = frame_idx return frame_idx def _finalize_markers(self): """Close any open markers at the end of profiling.""" end_time = self.last_sample_time # Close all open markers for each thread using a generic approach marker_states = [ (self.has_gil_start, "Has GIL", CATEGORY_GIL), (self.no_gil_start, "No GIL", CATEGORY_GIL), (self.on_cpu_start, "On CPU", CATEGORY_CPU), (self.off_cpu_start, "Off CPU", CATEGORY_CPU), (self.python_code_start, "Python Code", CATEGORY_CODE_TYPE), (self.native_code_start, "Native Code", CATEGORY_CODE_TYPE), (self.gil_wait_start, "Waiting for GIL", CATEGORY_GIL), (self.gc_start_per_thread, "GC Collecting", CATEGORY_GC), (self.exception_start, "Has Exception", CATEGORY_EXCEPTION), (self.no_exception_start, "No Exception", CATEGORY_EXCEPTION), ] for state_dict, marker_name, category in marker_states: for tid in list(state_dict.keys()): self._add_marker(tid, marker_name, state_dict[tid], end_time, category) del state_dict[tid] # Close any open opcode markers for tid, state in list(self.opcode_state.items()): opcode, lineno, col_offset, funcname, filename, start_time = state self._add_opcode_interval_marker(tid, opcode, lineno, col_offset, funcname, start_time, end_time) self.opcode_state.clear() def export(self, filename): """Export the profile to a Gecko JSON file.""" if self.sample_count > 0 and self.last_sample_time > 0: self.interval = self.last_sample_time / self.sample_count # Spinner for progress indication spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']) stop_spinner = threading.Event() def spin(): message = 'Building Gecko profile...' while not stop_spinner.is_set(): sys.stderr.write(f'\r{next(spinner)} {message}') sys.stderr.flush() time.sleep(0.1) # Clear the spinner line sys.stderr.write('\r' + ' ' * (len(message) + 3) + '\r') sys.stderr.flush() spinner_thread = threading.Thread(target=spin, daemon=True) spinner_thread.start() try: # Finalize any open markers before building profile self._finalize_markers() profile = self._build_profile() with open(filename, "w") as f: json.dump(profile, f, separators=(",", ":")) finally: stop_spinner.set() spinner_thread.join(timeout=1.0) # Small delay to ensure the clear happens time.sleep(0.01) print(f"Gecko profile written to {filename}") print( f"Open in Firefox Profiler: https://profiler.firefox.com/" ) def _build_marker_schema(self): """Build marker schema definitions for Firefox Profiler.""" schema = [] # Opcode marker schema (only if opcodes enabled) if self.opcodes_enabled: schema.append({ "name": "Opcode", "display": ["marker-table", "marker-chart"], "tooltipLabel": "{marker.data.opname}", "tableLabel": "{marker.data.opname} at line {marker.data.line}", "chartLabel": "{marker.data.opname}", "fields": [ {"key": "opname", "label": "Opcode", "format": "string", "searchable": True}, {"key": "base_opname", "label": "Base Opcode", "format": "string"}, {"key": "is_specialized", "label": "Specialized", "format": "string"}, {"key": "line", "label": "Line", "format": "integer"}, {"key": "column", "label": "Column", "format": "integer"}, {"key": "function", "label": "Function", "format": "string"}, {"key": "duration", "label": "Duration", "format": "duration"}, ], }) return schema def _build_profile(self): """Build the complete profile structure in processed format.""" # Convert thread data to final format threads = [] for tid, thread_data in self.threads.items(): # Update lengths samples = thread_data["samples"] stack_table = thread_data["stackTable"] frame_table = thread_data["frameTable"] func_table = thread_data["funcTable"] resource_table = thread_data["resourceTable"] samples["length"] = len(samples["stack"]) stack_table["length"] = len(stack_table["frame"]) frame_table["length"] = len(frame_table["func"]) func_table["length"] = len(func_table["name"]) resource_table["length"] = len(resource_table["name"]) thread_data["markers"]["length"] = len(thread_data["markers"]["name"]) # Clean up internal caches del thread_data["_stackCache"] del thread_data["_frameCache"] del thread_data["_funcCache"] del thread_data["_resourceCache"] threads.append(thread_data) # Main profile structure in processed format profile = { "meta": { "interval": self.interval, "startTime": self.start_time, "abi": platform.machine(), "misc": "Python profiler", "oscpu": platform.machine(), "platform": platform.system(), "processType": PROCESS_TYPE_MAIN, "categories": GECKO_CATEGORIES, "stackwalk": STACKWALK_DISABLED, "toolkit": "", "version": GECKO_FORMAT_VERSION, "preprocessedProfileVersion": GECKO_PREPROCESSED_VERSION, "appBuildID": "", "physicalCPUs": os.cpu_count() or 0, "logicalCPUs": os.cpu_count() or 0, "CPUName": "", "product": "Python", "symbolicated": True, "markerSchema": self._build_marker_schema(), "importedFrom": "Tachyon Sampling Profiler", "extensions": { "id": [], "name": [], "baseURL": [], "length": 0, }, }, "libs": self.libs, "threads": threads, "pages": [], "shared": { "stringArray": self.global_strings, "sources": {"length": 0, "uuid": [], "filename": []}, }, } return profile