Merge pull request #155 from Distributive-Network/Xmader/feat/pmjs-event-loop

pmjs event loop

Author: Xmader

include/PyEventLoop.hh

@@ -15,6 +15,7 @@
 #include <Python.h>
 #include <vector>
 #include <utility>
+#include <atomic>
 
 struct PyEventLoop {
 public:
@@ -179,9 +180,54 @@ public:
    */
   static PyEventLoop getMainLoop();
 
-protected:
-  PyObject *_loop;
+  struct Lock {
+  public:
+    explicit Lock() {
+      PyObject *asyncio = PyImport_ImportModule("asyncio");
+      _queueIsEmpty = PyObject_CallMethod(asyncio, "Event", NULL); // _queueIsEmpty = asyncio.Event()
+      Py_DECREF(asyncio);
+
+      // The flag should initially be set as the queue is initially empty
+      Py_XDECREF(PyObject_CallMethod(_queueIsEmpty, "set", NULL)); // _queueIsEmpty.set()
+    };
+    ~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;
+
+  PyObject *_loop;
+protected:
   PyEventLoop() = delete;
   PyEventLoop(PyObject *loop) : _loop(loop) {};
 private:

peter-jr

@@ -21,6 +21,7 @@ cd "$runDir"
 set -u
 set -o pipefail
 peter_exit_code=2
+defaultTimeout="${defaultTimeout:-15}"
 
 if [ "${1:-}" = "-h" ] || [ "${1:-}" = "--help" ]; then
 cat <<EOF
@@ -160,13 +161,14 @@ findTests \
       FAIL="$(red FAIL)"
     fi
     (
+      timeout=$(egrep '^( *[*#] *timeout: )' "$file" | head -1 | sed -e 's/^\( *[*#] *timeout: \)//')
       case "$testType" in
         "simple")
-          "${PMJS}" "$file"
+          eval timeout -s9 "${timeout:-${defaultTimeout}}" "${PMJS}" \"$file\"
           exitCode="$?"
           ;;
         "bash")
-          bash "$file"
+          eval timeout -s9 "${timeout:-${defaultTimeout}}" bash \"$file\"
           exitCode="$?"
           ;;
         *)

python/pythonmonkey/cli/pmjs.py

@@ -5,6 +5,7 @@
 
 import sys, os, signal, getopt
 import readline
+import asyncio
 import pythonmonkey as pm
 from pythonmonkey.lib import pmdb
 
@@ -356,8 +357,11 @@ def main():
             assert False, "unhandled option"
 
     if (len(args) > 0):
-        globalInitModule.patchGlobalRequire()
-        pm.runProgramModule(args[0], args, requirePath)
+        async def runJS():
+            globalInitModule.patchGlobalRequire()
+            pm.runProgramModule(args[0], args, requirePath)
+            await pm.wait() # blocks until all asynchronous calls finish
+        asyncio.run(runJS())
     elif (enterRepl or forceRepl):
         globalInitModule.initReplLibs()
         repl()

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,17 +2,36 @@
 
 #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();
+  if (PyErr_Occurred()) {
+    return NULL;
+  } else {
+    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
-  PyObject *asyncHandle = PyObject_CallMethod(_loop, "call_soon_threadsafe", "O", jobFn); // https://docs.python.org/3/c-api/arg.html#c.Py_BuildValue
+  PyObject *asyncHandle = PyObject_CallMethod(_loop, "call_soon_threadsafe", "O", wrapper); // https://docs.python.org/3/c-api/arg.html#c.Py_BuildValue
   return PyEventLoop::AsyncHandle(asyncHandle);
 }
 
 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
-  PyObject *asyncHandle = PyObject_CallMethod(_loop, "call_later", "dO", delaySeconds, jobFn); // https://docs.python.org/3/c-api/arg.html#c.Py_BuildValue
+  PyObject *asyncHandle = PyObject_CallMethod(_loop, "call_later", "dO", delaySeconds, wrapper); // https://docs.python.org/3/c-api/arg.html#c.Py_BuildValue
   if (asyncHandle == nullptr) {
     PyErr_Print(); // RuntimeError: Non-thread-safe operation invoked on an event loop other than the current one
   }
@@ -122,6 +141,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

@@ -268,6 +268,17 @@ static PyObject *eval(PyObject *self, PyObject *args) {
   }
 }
 
+static PyObject *waitForEventLoop(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(_)) {
+  PyObject *waiter = PyEventLoop::_locker->_queueIsEmpty; // instance of asyncio.Event
+
+  // Making sure it's attached to the current event-loop
+  PyEventLoop loop = PyEventLoop::getRunningLoop();
+  if (!loop.initialized()) return NULL;
+  PyObject_SetAttrString(waiter, "_loop", loop._loop);
+
+  return PyObject_CallMethod(waiter, "wait", NULL);
+}
+
 static PyObject *isCompilableUnit(PyObject *self, PyObject *args) {
   StrType *buffer = new StrType(PyTuple_GetItem(args, 0));
   const char *bufferUtf8;
@@ -289,6 +300,7 @@ static PyObject *isCompilableUnit(PyObject *self, PyObject *args) {
 
 PyMethodDef PythonMonkeyMethods[] = {
   {"eval", eval, METH_VARARGS, "Javascript evaluator in Python"},
+  {"wait", waitForEventLoop, METH_NOARGS, "The event-loop shield. Blocks until all asynchronous jobs finish."},
   {"isCompilableUnit", isCompilableUnit, METH_VARARGS, "Hint if a string might be compilable Javascript"},
   {"collect", collect, METH_VARARGS, "Calls the spidermonkey garbage collector"},
   {NULL, NULL, 0, NULL}
@@ -404,6 +416,9 @@ PyMODINIT_FUNC PyInit_pythonmonkey(void)
     return NULL;
   }
 
+  // Initialize event-loop shield
+  PyEventLoop::_locker = new PyEventLoop::Lock();
+
   PyObject *internalBindingPy = getInternalBindingPyFn(GLOBAL_CX);
   if (PyModule_AddObject(pyModule, "internalBinding", internalBindingPy) < 0) {
     Py_DECREF(internalBindingPy);