fix: handle condition waiting better in thread_task_executor

packages/gooddata-flight-server/src/gooddata_flight_server/server/flight_rpc/server_methods.py

@@ -101,6 +101,8 @@ def do_get_task_result(
                     f"While the result exists, it is of an unexpected type: {type(result).__name__} ",
                 ).to_internal_error()
 
+            finalizer = FlightServerMethods.call_finalizer_middleware(context)
+
             rlock, data = result.acquire_data()
 
             def _on_end(_: pyarrow.ArrowException | None) -> None:
@@ -121,7 +123,6 @@ def _on_end(_: pyarrow.ArrowException | None) -> None:
                         # log and sink these Exceptions - not much to do
                         _LOGGER.error("do_get_close_failed", exc_info=True)
 
-            finalizer = FlightServerMethods.call_finalizer_middleware(context)
             finalizer.register_on_end(_on_end)
 
             if isinstance(data, pyarrow.Table):

packages/gooddata-flight-server/src/gooddata_flight_server/tasks/thread_task_executor.py

@@ -158,6 +158,7 @@ class _TaskExecution:
         "_result_future",
         "_lock",
         "_completed",
+        "_done",
         "_stats",
     )
 
@@ -189,6 +190,8 @@ def __init__(
         # all these are protected using the lock
         self._result_future: Future[Union[TaskResult, TaskError]] | None = None
         self._completed: threading.Condition = threading.Condition(self._lock)
+        # indicates the task actually finished
+        self._done: bool = False
 
     @property
     def task(self) -> Task:
@@ -235,6 +238,7 @@ def on_result_done(self, fut: Future) -> None:
 
         with self._lock:
             execution_result = self._cb.process_task_result(self, self._result_future)
+            self._done = True
             self._completed.notify_all()
 
         self._complete_execution_span(execution_result)
@@ -291,7 +295,7 @@ def cancel(self) -> bool:
 
     def wait_for_completion(self, timeout: float | None = None) -> None:
         with self._lock:
-            completed = self._completed.wait(timeout=timeout)
+            completed = self._completed.wait_for(lambda: self._done, timeout=timeout)
 
         if not completed:
             raise TaskWaitTimeoutError(task_id=self._task.task_id, cmd=self._task.cmd)