Merge pull request #294 from Distributive-Network/Xmader/feat/timers-unref

Feat: Node.js-style timers

Author: philippedistributive

include/PyEventLoop.hh

@@ -32,15 +32,21 @@ public:
    * @see https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.Handle
    */
   struct AsyncHandle {
+    using id_t = uint32_t;
+    using id_ptr_pair = std::pair<id_t, AsyncHandle *>;
   public:
     explicit AsyncHandle(PyObject *handle) : _handle(handle) {};
     AsyncHandle(const AsyncHandle &old) = delete; // forbid copy-initialization
-    AsyncHandle(AsyncHandle &&old) : _handle(std::exchange(old._handle, nullptr)) {}; // clear the moved-from object
+    AsyncHandle(AsyncHandle &&old) : _handle(std::exchange(old._handle, nullptr)), _refed(old._refed.exchange(false)) {}; // clear the moved-from object
     ~AsyncHandle() {
       if (Py_IsInitialized()) { // the Python runtime has already been finalized when `_timeoutIdMap` is cleared at exit
         Py_XDECREF(_handle);
       }
     }
+    static inline id_ptr_pair newEmpty() {
+      auto handle = AsyncHandle(Py_None);
+      return AsyncHandle::getUniqueIdAndPtr(std::move(handle));
+    }
 
     /**
      * @brief Cancel the scheduled event-loop job.
@@ -52,18 +58,23 @@ public:
      * @brief Get the unique `timeoutID` for JS `setTimeout`/`clearTimeout` methods
      * @see https://developer.mozilla.org/en-US/docs/Web/API/setTimeout#return_value
      */
-    static inline uint32_t getUniqueId(AsyncHandle &&handle) {
+    static inline id_t getUniqueId(AsyncHandle &&handle) {
       // TODO (Tom Tang): mutex lock
       _timeoutIdMap.push_back(std::move(handle));
       return _timeoutIdMap.size() - 1; // the index in `_timeoutIdMap`
     }
-    static inline AsyncHandle *fromId(uint32_t timeoutID) {
+    static inline AsyncHandle *fromId(id_t timeoutID) {
       try {
         return &_timeoutIdMap.at(timeoutID);
       } catch (...) { // std::out_of_range&
         return nullptr; // invalid timeoutID
       }
     }
+    static inline id_ptr_pair getUniqueIdAndPtr(AsyncHandle &&handle) {
+      auto timeoutID = getUniqueId(std::move(handle));
+      auto ptr = fromId(timeoutID);
+      return std::make_pair(timeoutID, ptr);
+    }
 
     /**
      * @brief Get the underlying `asyncio.Handle` Python object
@@ -72,8 +83,44 @@ public:
       Py_INCREF(_handle); // otherwise the object would be GC-ed as the AsyncHandle destructor decreases the reference count
       return _handle;
     }
+
+    /**
+     * @brief Replace the underlying `asyncio.Handle` Python object with the provided value
+     * @return the old `asyncio.Handle` object
+     */
+    inline PyObject *swap(PyObject *newHandleObject) {
+      return std::exchange(_handle, newHandleObject);
+    }
+
+    /**
+     * @brief Getter for if the timer has been ref'ed
+     */
+    inline bool hasRef() {
+      return _refed;
+    }
+
+    /**
+     * @brief Ref the timer so that the event-loop won't exit as long as the timer is active
+     */
+    inline void addRef() {
+      if (!_refed) {
+        _refed = true;
+        PyEventLoop::_locker->incCounter();
+      }
+    }
+
+    /**
+     * @brief Unref the timer so that the event-loop can exit
+     */
+    inline void removeRef() {
+      if (_refed) {
+        _refed = false;
+        PyEventLoop::_locker->decCounter();
+      }
+    }
   protected:
     PyObject *_handle;
+    std::atomic_bool _refed = false;
   };
 
   /**
@@ -86,9 +133,9 @@ public:
    * @brief Schedule a job to the Python event-loop, with the given delay
    * @param jobFn - The JS event-loop job converted to a Python function
    * @param delaySeconds - The job function will be called after the given number of seconds
-   * @return a AsyncHandle, the value can be safely ignored
+   * @return the timeoutId and a pointer to the AsyncHandle
    */
-  AsyncHandle enqueueWithDelay(PyObject *jobFn, double delaySeconds);
+  [[nodiscard]] AsyncHandle::id_ptr_pair enqueueWithDelay(PyObject *jobFn, double delaySeconds);
 
   /**
    * @brief C++ wrapper for Python `asyncio.Future` class

python/pythonmonkey/builtin_modules/internal-binding.d.ts

@@ -43,6 +43,21 @@ declare function internalBinding(namespace: "timers"): {
    * internal binding helper for the `clearTimeout` global function
    */
   cancelByTimeoutId(timeoutId: number): void;
+
+  /**
+   * internal binding helper for if a timer object has been ref'ed
+   */
+  timerHasRef(timeoutId: number): boolean;
+
+  /**
+   * internal binding helper for ref'ing the timer
+   */
+  timerAddRef(timeoutId: number): void;
+
+  /**
+   * internal binding helper for unref'ing the timer
+   */
+  timerRemoveRef(timeoutId: number): void;
 };
 
 export = internalBinding;

python/pythonmonkey/builtin_modules/timers.js

@@ -7,17 +7,85 @@ const internalBinding = require('internal-binding');
 
 const { 
   enqueueWithDelay,
-  cancelByTimeoutId
+  cancelByTimeoutId,
+  timerHasRef,
+  timerAddRef,
+  timerRemoveRef,
 } = internalBinding('timers');
 
+/**
+ * Implement Node.js-style `timeoutId` class returned from setTimeout() and setInterval()
+ * @see https://nodejs.org/api/timers.html#class-timeout
+ */
+class Timeout
+{
+  /** @type {number} an integer */
+  #numericId;
+
+  /**
+   * @param {number} numericId 
+   */
+  constructor(numericId)
+  {
+    this.#numericId = numericId;
+  }
+
+  /**
+   * If `true`, the `Timeout` object will keep the event-loop active.
+   * @returns {boolean}
+   */
+  hasRef()
+  {
+    return timerHasRef(this.#numericId);
+  }
+
+  /**
+   * When called, requests that the event-loop **not exit** so long as the `Timeout` is active.
+   *   
+   * By default, all `Timeout` objects are "ref'ed", making it normally unnecessary to call `timeout.ref()` unless `timeout.unref()` had been called previously.
+   */
+  ref()
+  {
+    timerAddRef(this.#numericId);
+    return this; // allow chaining
+  }
+
+  /**
+   * When called, the active `Timeout` object will not require the event-loop to remain active.  
+   * If there is no other activity keeping the event-loop running, the process may exit before the `Timeout` object's callback is invoked.
+   */
+  unref()
+  {
+    timerRemoveRef(this.#numericId);
+    return this; // allow chaining
+  }
+
+  /**
+   * @returns a number that can be used to reference this timeout
+   */
+  [Symbol.toPrimitive]()
+  {
+    return this.#numericId;
+  }
+
+  /**
+   * Cancels the timeout.
+   * @deprecated legacy Node.js API. Use `clearTimeout()` instead
+   */
+  close()
+  {
+    return clearTimeout(this);
+  }
+}
+
 /**
  * Implement the `setTimeout` global function
  * @see https://developer.mozilla.org/en-US/docs/Web/API/setTimeout and
  * @see https://html.spec.whatwg.org/multipage/timers-and-user-prompts.html#dom-settimeout
  * @param {Function | string} handler
  * @param {number} delayMs timeout milliseconds, use value of 0 if this is omitted
  * @param {any[]} args additional arguments to be passed to the `handler`
- * @return {number} timeoutId
+ * @return {Timeout} timeoutId
  */
 function setTimeout(handler, delayMs = 0, ...args) 
 {
@@ -41,25 +109,28 @@ function setTimeout(handler, delayMs = 0, ...args)
     delayMs = 0; // as spec-ed
   const delaySeconds = delayMs / 1000; // convert ms to s
 
-  return enqueueWithDelay(boundHandler, delaySeconds);
+  return new Timeout(enqueueWithDelay(boundHandler, delaySeconds));
 }
 
 /**
  * Implement the `clearTimeout` global function
  * @see https://developer.mozilla.org/en-US/docs/Web/API/clearTimeout and
  * @see https://html.spec.whatwg.org/multipage/timers-and-user-prompts.html#dom-cleartimeout
- * @param {number} timeoutId
+ * @param {Timeout | number} timeoutId
  * @return {void}
  */
 function clearTimeout(timeoutId) 
 {
-  // silently does nothing when an invalid timeoutId (should be an int32 value) is passed in
-  if (!Number.isInteger(timeoutId))
+  // silently does nothing when an invalid timeoutId (should be a Timeout instance or an int32 value) is passed in
+  if (!(timeoutId instanceof Timeout) && !Number.isInteger(timeoutId))
     return;
 
-  return cancelByTimeoutId(timeoutId);
+  return cancelByTimeoutId(Number(timeoutId));
 }
 
+// expose the `Timeout` class
+setTimeout.Timeout = Timeout;
+
 if (!globalThis.setTimeout)
   globalThis.setTimeout = setTimeout;
 if (!globalThis.clearTimeout)

src/PyEventLoop.cc

@@ -15,27 +15,47 @@ static PyObject *eventLoopJobWrapper(PyObject *jobFn, PyObject *Py_UNUSED(_)) {
     Py_RETURN_NONE;
   }
 }
-static PyMethodDef jobWrapperDef = {"eventLoopJobWrapper", eventLoopJobWrapper, METH_NOARGS, NULL};
+static PyMethodDef loopJobWrapperDef = {"eventLoopJobWrapper", eventLoopJobWrapper, METH_NOARGS, NULL};
+
+/**
+ * @brief Wrapper to remove the reference of the timer after the job finishes
+ */
+static PyObject *timerJobWrapper(PyObject *jobFn, PyObject *handlerPtr) {
+  auto handle = (PyEventLoop::AsyncHandle *)PyLong_AsVoidPtr(handlerPtr);
+  PyObject *ret = PyObject_CallObject(jobFn, NULL); // jobFn()
+  Py_XDECREF(ret); // don't care about its return value
+  handle->removeRef();
+  if (PyErr_Occurred()) {
+    return NULL;
+  } else {
+    Py_RETURN_NONE;
+  }
+}
+static PyMethodDef timerJobWrapperDef = {"timerJobWrapper", timerJobWrapper, METH_O, NULL};
 
 PyEventLoop::AsyncHandle PyEventLoop::enqueue(PyObject *jobFn) {
   PyEventLoop::_locker->incCounter();
-  PyObject *wrapper = PyCFunction_New(&jobWrapperDef, jobFn);
+  PyObject *wrapper = PyCFunction_New(&loopJobWrapperDef, 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", 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);
+PyEventLoop::AsyncHandle::id_ptr_pair PyEventLoop::enqueueWithDelay(PyObject *jobFn, double delaySeconds) {
+  auto handler = PyEventLoop::AsyncHandle::newEmpty();
+  PyObject *wrapper = PyCFunction_New(&timerJobWrapperDef, jobFn);
+  PyObject *handlerPtr = PyLong_FromVoidPtr(handler.second);
   // 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, wrapper); // https://docs.python.org/3/c-api/arg.html#c.Py_BuildValue
+  PyObject *asyncHandle = PyObject_CallMethod(_loop, "call_later", "dOO", delaySeconds, wrapper, handlerPtr); // 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
+    return handler;
   }
-  return PyEventLoop::AsyncHandle(asyncHandle);
+  handler.second->swap(asyncHandle);
+  handler.second->addRef();
+  return handler;
 }
 
 PyEventLoop::Future PyEventLoop::createFuture() {
@@ -145,7 +165,7 @@ void PyEventLoop::AsyncHandle::cancel() {
                                                                        // 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();
+    removeRef(); // automatically unref at finish
   }
   Py_XDECREF(scheduled);
 

src/internalBinding/timers.cc

@@ -10,6 +10,8 @@
 
 #include <jsapi.h>
 
+using AsyncHandle = PyEventLoop::AsyncHandle;
+
 /**
  * See function declarations in python/pythonmonkey/builtin_modules/internal-binding.d.ts :
  *    `declare function internalBinding(namespace: "timers")`
@@ -29,15 +31,14 @@ static bool enqueueWithDelay(JSContext *cx, unsigned argc, JS::Value *vp) {
   // Schedule job to the running Python event-loop
   PyEventLoop loop = PyEventLoop::getRunningLoop();
   if (!loop.initialized()) return false;
-  PyEventLoop::AsyncHandle handle = loop.enqueueWithDelay(job, delaySeconds);
+  PyEventLoop::AsyncHandle::id_ptr_pair handler = loop.enqueueWithDelay(job, delaySeconds);
 
   // Return the `timeoutID` to use in `clearTimeout`
-  args.rval().setNumber(PyEventLoop::AsyncHandle::getUniqueId(std::move(handle)));
+  args.rval().setNumber(handler.first);
   return true;
 }
 
 static bool cancelByTimeoutId(JSContext *cx, unsigned argc, JS::Value *vp) {
-  using AsyncHandle = PyEventLoop::AsyncHandle;
   JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
   double timeoutID = args.get(0).toNumber();
 
@@ -53,8 +54,51 @@ static bool cancelByTimeoutId(JSContext *cx, unsigned argc, JS::Value *vp) {
   return true;
 }
 
+static bool timerHasRef(JSContext *cx, unsigned argc, JS::Value *vp) {
+  JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
+  double timeoutID = args.get(0).toNumber();
+
+  // Retrieve the AsyncHandle by `timeoutID`
+  AsyncHandle *handle = AsyncHandle::fromId((uint32_t)timeoutID);
+  if (!handle) return false; // error no such timeoutID
+
+  args.rval().setBoolean(handle->hasRef());
+  return true;
+}
+
+static bool timerAddRef(JSContext *cx, unsigned argc, JS::Value *vp) {
+  JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
+  double timeoutID = args.get(0).toNumber();
+
+  // Retrieve the AsyncHandle by `timeoutID`
+  AsyncHandle *handle = AsyncHandle::fromId((uint32_t)timeoutID);
+  if (!handle) return false; // error no such timeoutID
+
+  handle->addRef();
+
+  args.rval().setUndefined();
+  return true;
+}
+
+static bool timerRemoveRef(JSContext *cx, unsigned argc, JS::Value *vp) {
+  JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
+  double timeoutID = args.get(0).toNumber();
+
+  // Retrieve the AsyncHandle by `timeoutID`
+  AsyncHandle *handle = AsyncHandle::fromId((uint32_t)timeoutID);
+  if (!handle) return false; // error no such timeoutID
+
+  handle->removeRef();
+
+  args.rval().setUndefined();
+  return true;
+}
+
 JSFunctionSpec InternalBinding::timers[] = {
   JS_FN("enqueueWithDelay", enqueueWithDelay, /* nargs */ 2, 0),
   JS_FN("cancelByTimeoutId", cancelByTimeoutId, 1, 0),
+  JS_FN("timerHasRef", timerHasRef, 1, 0),
+  JS_FN("timerAddRef", timerAddRef, 1, 0),
+  JS_FN("timerRemoveRef", timerRemoveRef, 1, 0),
   JS_FS_END
 };