utils/concurrent: Fix FutureSetWatcher.wait implementation

Explicitly track done callback invocations.

Author: ales-erjavec

Orange/widgets/utils/concurrent.py

@@ -21,7 +21,7 @@
     pass
 
 from AnyQt.QtCore import (
-    Qt, QObject, QMetaObject, QThreadPool, QThread, QRunnable,
+    Qt, QObject, QMetaObject, QThreadPool, QThread, QRunnable, QSemaphore,
     QEventLoop, QCoreApplication, QEvent, Q_ARG
 )
 
@@ -573,6 +573,7 @@ def __init__(self, futures=None, *args, **kwargs):
         # type: (List[Future], ...) -> None
         super().__init__(*args, **kwargs)
         self.__futures = None
+        self.__semaphore = None
         self.__countdone = 0
         if futures is not None:
             self.setFutures(futures)
@@ -594,21 +595,29 @@ def setFutures(self, futures):
         selfweakref = weakref.ref(self)
         schedule_emit = methodinvoke(self, "__emitpending", (int, Future))
 
+        # Semaphore counting the number of future that have enqueued
+        # done notifications. Used for the `wait` implementation.
+        self.__semaphore = semaphore = QSemaphore(0)
+
         for i, future in enumerate(futures):
             self.__futures.append(future)
 
             def on_done(index, f):
-                selfref = selfweakref()  # not safe really
-                if selfref is None:
-                    return
                 try:
-                    schedule_emit(index, f)
-                except RuntimeError:
-                    # Ignore RuntimeErrors (when C++ side of QObject is deleted)
-                    # (? Use QObject.destroyed and remove the done callback ?)
-                    pass
+                    selfref = selfweakref()  # not safe really
+                    if selfref is None:  # pragma: no cover
+                        return
+                    try:
+                        schedule_emit(index, f)
+                    except RuntimeError:  # pragma: no cover
+                        # Ignore RuntimeErrors (when C++ side of QObject is deleted)
+                        # (? Use QObject.destroyed and remove the done callback ?)
+                        pass
+                finally:
+                    semaphore.release()
 
             future.add_done_callback(partial(on_done, i))
+
         if not self.__futures:
             # `futures` was an empty sequence.
             methodinvoke(self, "doneAll", ())()
@@ -644,13 +653,28 @@ def __emitpending(self, index, future):
     def flush(self):
         """
         Flush all pending signal emits currently enqueued.
+
+        Must only ever be called from the thread this object lives in
+        (:func:`QObject.thread()`).
         """
-        assert QThread.currentThread() is self.thread()
+        if QThread.currentThread() is not self.thread():
+            raise RuntimeError("`flush()` called from a wrong thread.")
+        # NOTE: QEvent.MetaCall is the event implementing the
+        # `Qt.QueuedConnection` method invocation.
         QCoreApplication.sendPostedEvents(self, QEvent.MetaCall)
 
     def wait(self):
-        assert self.__futures is not None
-        concurrent.futures.wait(self.__futures)
+        """
+        Wait for for all the futures to complete and *enqueue* notifications
+        to this object, but do not emit any signals.
+
+        Use `flush()` to emit all signals after a `wait()`
+        """
+        if self.__futures is None:
+            raise RuntimeError("Futures were not set.")
+
+        self.__semaphore.acquire(len(self.__futures))
+        self.__semaphore.release(len(self.__futures))
 
 
 class methodinvoke(object):

Orange/widgets/utils/tests/test_concurrent.py

@@ -1,4 +1,5 @@
 import unittest
+import unittest.mock
 import threading
 import random
 
@@ -226,6 +227,18 @@ def as_set(seq):
         spy = spies(watcher)
         self.assertTrue(spy.doneAll.wait())
 
+        watcher = FutureSetWatcher()
+        watcher.setFutures([])
+        watcher.wait()
+
+        watcher = FutureSetWatcher()
+        with self.assertRaises(RuntimeError):
+            watcher.wait()
+
+        with unittest.mock.patch.object(watcher, "thread", lambda: 42), \
+                self.assertRaises(RuntimeError):
+            watcher.flush()
+
 
 class TestTask(CoreAppTestCase):