Skip to content

Commit 29daff6

Browse files
godlygeekpablogsal
authored andcommitted
Lock around Tracker creation and destruction
The `Tracker.__enter__` and `Tracker.__exit` methods may wind up releasing the GIL, which allows another thread to see an intermediate state where the tracker is not fully installed. Add a lock, shared across all trackers, to serialize access to the global state that tracker installation and uninstallation writes to. Signed-off-by: Matt Wozniski <[email protected]>
1 parent f848fc3 commit 29daff6

File tree

2 files changed

+38
-32
lines changed

2 files changed

+38
-32
lines changed

news/667.bugfix.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix a race condition that was able to cause strange exception messages if two different threads tried to initialize Memray tracking at once.

src/memray/_memray.pyx

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,9 @@ cdef class ProfileFunctionGuard:
581581
NativeTracker.forgetPythonStack()
582582

583583

584+
tracker_creation_lock = threading.Lock()
585+
586+
584587
cdef class Tracker:
585588
"""Context manager for tracking memory allocations in a Python script.
586589
@@ -690,46 +693,48 @@ cdef class Tracker:
690693

691694
@cython.profile(False)
692695
def __enter__(self):
693-
if NativeTracker.getTracker() != NULL:
694-
raise RuntimeError("No more than one Tracker instance can be active at the same time")
695-
696696
cdef unique_ptr[RecordWriter] writer
697-
if self._writer == NULL:
698-
raise RuntimeError("Attempting to use stale output handle")
699-
writer = move(self._writer)
700-
701-
for attr in ("_name", "_ident"):
702-
assert not hasattr(threading.Thread, attr)
703-
setattr(
704-
threading.Thread,
705-
attr,
706-
ThreadNameInterceptor(attr, NativeTracker.registerThreadNameById),
707-
)
697+
with tracker_creation_lock:
698+
if NativeTracker.getTracker() != NULL:
699+
raise RuntimeError("No more than one Tracker instance can be active at the same time")
700+
701+
if self._writer == NULL:
702+
raise RuntimeError("Attempting to use stale output handle")
703+
writer = move(self._writer)
704+
705+
for attr in ("_name", "_ident"):
706+
assert not hasattr(threading.Thread, attr)
707+
setattr(
708+
threading.Thread,
709+
attr,
710+
ThreadNameInterceptor(attr, NativeTracker.registerThreadNameById),
711+
)
708712

709-
self._previous_profile_func = sys.getprofile()
710-
self._previous_thread_profile_func = threading._profile_hook
711-
threading.setprofile(start_thread_trace)
713+
self._previous_profile_func = sys.getprofile()
714+
self._previous_thread_profile_func = threading._profile_hook
715+
threading.setprofile(start_thread_trace)
712716

713-
if "greenlet" in sys.modules:
714-
NativeTracker.beginTrackingGreenlets()
717+
if "greenlet" in sys.modules:
718+
NativeTracker.beginTrackingGreenlets()
715719

716-
NativeTracker.createTracker(
717-
move(writer),
718-
self._native_traces,
719-
self._memory_interval_ms,
720-
self._follow_fork,
721-
self._trace_python_allocators,
722-
)
723-
return self
720+
NativeTracker.createTracker(
721+
move(writer),
722+
self._native_traces,
723+
self._memory_interval_ms,
724+
self._follow_fork,
725+
self._trace_python_allocators,
726+
)
727+
return self
724728

725729
@cython.profile(False)
726730
def __exit__(self, exc_type, exc_value, exc_traceback):
727-
NativeTracker.destroyTracker()
728-
sys.setprofile(self._previous_profile_func)
729-
threading.setprofile(self._previous_thread_profile_func)
731+
with tracker_creation_lock:
732+
NativeTracker.destroyTracker()
733+
sys.setprofile(self._previous_profile_func)
734+
threading.setprofile(self._previous_thread_profile_func)
730735

731-
for attr in ("_name", "_ident"):
732-
delattr(threading.Thread, attr)
736+
for attr in ("_name", "_ident"):
737+
delattr(threading.Thread, attr)
733738

734739

735740
def start_thread_trace(frame, event, arg):

0 commit comments

Comments
 (0)