From faa5accaa284fe12cb66726233594770f15664b3 Mon Sep 17 00:00:00 2001 From: Krzysztof Godlewski Date: Mon, 13 Jan 2025 18:06:15 +0100 Subject: [PATCH] Run._wait() also returns during closing when called from the main thread --- src/neptune_scale/api/run.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/neptune_scale/api/run.py b/src/neptune_scale/api/run.py index 5ee82a5..506674c 100644 --- a/src/neptune_scale/api/run.py +++ b/src/neptune_scale/api/run.py @@ -614,13 +614,15 @@ def _wait( # Handle the case where we get notified on `wait_seq` before we actually wait. # Otherwise, we would unnecessarily block, waiting on a notify_all() that never happens. if wait_seq.value >= self._operations_queue.last_sequence_id: - break + return True - if is_closing and threading.current_thread() != self._closing_thread: - if verbose: - logger.warning("Waiting interrupted by run termination") + if is_closing: + if threading.current_thread() != self._closing_thread: + if verbose: + logger.warning("Waiting interrupted by run termination") + + self._close_completed.wait(wait_time) - self._close_completed.wait(wait_time) return False with wait_seq: