Fix erros in `check_thread_leak`
@@ -1417,34 +1417,35 @@ def save_sys_modules():
@contextmanager
def check_thread_leak():
- active_threads_start = set(threading._active)
+ """ Context manager to ensure we haven't leaked any threads """
+ active_threads_start = threading.enumerate()
yield
start = time()
while True:
- bad = [
- t
- for t, v in threading._active.items()
- if t not in active_threads_start
- and "Threaded" not in v.name
- and "watch message" not in v.name
- and "TCP-Executor" not in v.name
+ bad_threads = [
+ thread
+ for thread in threading.enumerate()
+ if thread not in active_threads_start
+ and "Threaded" not in thread.name
+ and "watch message" not in thread.name
+ and "TCP-Executor" not in thread.name
# TODO: Make sure profile thread is cleaned up
# and remove the line below
- and "Profile" not in v.name
+ and "Profile" not in thread.name
]
- if not bad:
+ if not bad_threads:
break
else:
sleep(0.01)
if time() > start + 5:
+ # Raise an error with information about leaked threads
from distributed import profile
- tid = bad[0]
- thread = threading._active[tid]
- call_stacks = profile.call_stack(sys._current_frames()[tid])
- assert False, (thread, call_stacks)
+ bad_thread = bad_threads[0]
+ call_stacks = profile.call_stack(sys._current_frames()[bad_thread.ident])
+ assert False, (bad_thread, call_stacks)
@contextmanager
