feat(event-loop): event-loop shield that blocks until all async jobs finish

await pm.wait()

Author: Xmader

include/PyEventLoop.hh

@@ -15,6 +15,7 @@
 #include <Python.h>
 #include <vector>
 #include <utility>
+#include <atomic>
 
 struct PyEventLoop {
 public:
@@ -179,6 +180,48 @@ public:
    */
   static PyEventLoop getMainLoop();
 
+  struct Lock {
+  public:
+    explicit Lock() {
+      PyObject *asyncio = PyImport_ImportModule("asyncio");
+      _queueIsEmpty = PyObject_CallMethod(asyncio, "Event", NULL); // _queueIsEmpty = asyncio.Event()
+      Py_DECREF(asyncio);
+    };
+    ~Lock() {
+      Py_DECREF(_queueIsEmpty);
+    }
+
+    /**
+     * @brief Increment the counter for the number of our job functions in the Python event-loop
+     */
+    inline void incCounter() {
+      _counter++;
+      Py_XDECREF(PyObject_CallMethod(_queueIsEmpty, "clear", NULL)); // _queueIsEmpty.clear()
+    }
+
+    /**
+     * @brief Decrement the counter for the number of our job functions in the Python event-loop
+     */
+    inline void decCounter() {
+      _counter--;
+      if (_counter == 0) { // no job queueing
+        // Notify that the queue is empty and awake (unblock) the event-loop shield
+        Py_XDECREF(PyObject_CallMethod(_queueIsEmpty, "set", NULL)); // _queueIsEmpty.set()
+      } else if (_counter < 0) { // something went wrong
+        PyErr_SetString(PyExc_RuntimeError, "Event-loop job counter went below zero.");
+      }
+    }
+
+    /**
+     * @brief An `asyncio.Event` instance to notify that there are no queueing asynchronous jobs
+     * @see https://docs.python.org/3/library/asyncio-sync.html#asyncio.Event
+     */
+    PyObject *_queueIsEmpty = nullptr;
+  protected:
+    std::atomic_int _counter = 0;
+  };
+
+  static inline PyEventLoop::Lock *_locker;
 protected:
   PyObject *_loop;
 

python/pythonmonkey/pythonmonkey.pyi

@@ -22,6 +22,17 @@ def eval(code: str, evalOpts: EvalOptions = {}, /) -> _typing.Any:
     JavaScript evaluator in Python
     """
 
+def wait() -> _typing.Awaitable[None]:
+    """
+    Block until all asynchronous jobs (Promise/setTimeout/etc.) finish.
+    
+    ```py
+    await pm.wait()
+    ```
+
+    This is the event-loop shield that protects the loop from being prematurely terminated.
+    """
+
 def isCompilableUnit(code: str) -> bool:
     """
     Hint if a string might be compilable Javascript without actual evaluation

src/PyEventLoop.cc

@@ -2,14 +2,19 @@
 
 #include <Python.h>
 
+/**
+ * @brief Wrapper to decrement the counter of queueing event-loop jobs after the job finishes
+ */
 static PyObject *eventLoopJobWrapper(PyObject *jobFn, PyObject *Py_UNUSED(_)) {
   PyObject *ret = PyObject_CallObject(jobFn, NULL); // jobFn()
   Py_XDECREF(ret); // don't care about its return value
+  PyEventLoop::_locker->decCounter();
   Py_RETURN_NONE;
 }
 static PyMethodDef jobWrapperDef = {"eventLoopJobWrapper", eventLoopJobWrapper, METH_NOARGS, NULL};
 
 PyEventLoop::AsyncHandle PyEventLoop::enqueue(PyObject *jobFn) {
+  PyEventLoop::_locker->incCounter();
   PyObject *wrapper = PyCFunction_New(&jobWrapperDef, jobFn);
   // Enqueue job to the Python event-loop
   //    https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.call_soon
@@ -18,6 +23,7 @@ PyEventLoop::AsyncHandle PyEventLoop::enqueue(PyObject *jobFn) {
 }
 
 PyEventLoop::AsyncHandle PyEventLoop::enqueueWithDelay(PyObject *jobFn, double delaySeconds) {
+  PyEventLoop::_locker->incCounter();
   PyObject *wrapper = PyCFunction_New(&jobWrapperDef, jobFn);
   // Schedule job to the Python event-loop
   //    https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.call_later
@@ -131,6 +137,14 @@ PyEventLoop PyEventLoop::getRunningLoop() {
 }
 
 void PyEventLoop::AsyncHandle::cancel() {
+  PyObject *scheduled = PyObject_GetAttrString(_handle, "_scheduled"); // this attribute only exists on asyncio.TimerHandle returned by loop.call_later
+                                                                       // NULL if no such attribute (on a strict asyncio.Handle returned by loop.call_soon)
+  bool finishedOrCanceled = scheduled && scheduled == Py_False; // the job function has already been executed or canceled
+  if (!finishedOrCanceled) {
+    PyEventLoop::_locker->decCounter();
+  }
+  Py_XDECREF(scheduled);
+
   // https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.Handle.cancel
   PyObject *ret = PyObject_CallMethod(_handle, "cancel", NULL); // returns None
   Py_XDECREF(ret);

src/modules/pythonmonkey/pythonmonkey.cc

@@ -404,6 +404,15 @@ PyMODINIT_FUNC PyInit_pythonmonkey(void)
     return NULL;
   }
 
+  // Initialize event-loop shield
+  PyEventLoop::_locker = new PyEventLoop::Lock();
+  PyObject *waiter = PyObject_GetAttrString(PyEventLoop::_locker->_queueIsEmpty /* instance of asyncio.Event*/, "wait");
+  if (!waiter || PyModule_AddObject(pyModule, "wait", waiter) < 0) {
+    Py_XDECREF(waiter);
+    Py_DECREF(pyModule);
+    return NULL;
+  }
+
   PyObject *internalBindingPy = getInternalBindingPyFn(GLOBAL_CX);
   if (PyModule_AddObject(pyModule, "internalBinding", internalBindingPy) < 0) {
     Py_DECREF(internalBindingPy);