PyEventLoop.hh

Go to the documentation of this file.

 
#ifndef PythonMonkey_PyEventLoop_
#define PythonMonkey_PyEventLoop_
 
#include <Python.h>
#include <jsapi.h>
#include <vector>
#include <utility>
#include <atomic>
 
struct PyEventLoop {
public:
  ~PyEventLoop() {
    Py_XDECREF(_loop);
  }
 
  bool initialized() const {
    return !!_loop;
  }
 
  struct AsyncHandle {
    using id_t = uint32_t;
  public:
    explicit AsyncHandle(PyObject *handle) : _handle(handle) {};
    AsyncHandle(const AsyncHandle &old) = delete; // forbid copy-initialization
    AsyncHandle(AsyncHandle &&old) : _handle(std::exchange(old._handle, nullptr)), _refed(old._refed.exchange(false)), _debugInfo(std::exchange(old._debugInfo, nullptr)) {}; // 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_t newEmpty() {
      auto handle = AsyncHandle(Py_None);
      return AsyncHandle::getUniqueId(std::move(handle));
    }
 
    void cancel();
    bool cancelled();
    bool _finishedOrCancelled();
 
    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(id_t timeoutID) {
      try {
        return &_timeoutIdMap.at(timeoutID);
      } catch (...) { // std::out_of_range&
        return nullptr; // invalid timeoutID
      }
    }
 
    static bool cancelAll();
 
    inline PyObject *getHandleObject() const {
      Py_INCREF(_handle); // otherwise the object would be GC-ed as the AsyncHandle destructor decreases the reference count
      return _handle;
    }
 
    inline PyObject *swap(PyObject *newHandleObject) {
      return std::exchange(_handle, newHandleObject);
    }
 
    inline bool hasRef() {
      return _refed;
    }
 
    inline void addRef() {
      if (!_refed) {
        _refed = true;
        if (!_finishedOrCancelled()) { // noop if the timer is finished or canceled
          PyEventLoop::_locker->incCounter();
        }
      }
    }
 
    inline void removeRef() {
      if (_refed) {
        _refed = false;
        PyEventLoop::_locker->decCounter();
      }
    }
 
    inline void setDebugInfo(PyObject *obj) {
      _debugInfo = obj;
    }
    inline PyObject *getDebugInfo() {
      return _debugInfo;
    }
 
    static inline auto &getAllTimers() {
      return _timeoutIdMap;
    }
  protected:
    PyObject *_handle;
    std::atomic_bool _refed = false;
    PyObject *_debugInfo = nullptr;
  };
 
  AsyncHandle enqueue(PyObject *jobFn);
  [[nodiscard]] AsyncHandle::id_t enqueueWithDelay(PyObject *jobFn, double delaySeconds, bool repeat);
 
  struct Future {
  public:
    explicit Future(PyObject *future) : _future(future) {};
    Future(const Future &old) = delete; // forbid copy-initialization
    Future(Future &&old) : _future(std::exchange(old._future, nullptr)) {}; // clear the moved-from object
    ~Future() {
      Py_XDECREF(_future);
    }
 
    void setResult(PyObject *result);
 
    void setException(PyObject *exception);
 
    void addDoneCallback(PyObject *cb);
 
    bool isCancelled();
 
    PyObject *getResult();
 
    PyObject *getException();
 
    inline PyObject *getFutureObject() const {
      Py_INCREF(_future); // otherwise the object would be GC-ed as this `PyEventLoop::Future` destructs
      return _future;
    }
  protected:
    PyObject *_future;
  };
 
  Future createFuture();
 
  Future ensureFuture(PyObject *awaitable);
 
  static PyEventLoop getRunningLoop();
 
  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);
 
      // 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);
    }
 
    inline void incCounter() {
      _counter++;
      Py_XDECREF(PyObject_CallMethod(_queueIsEmpty, "clear", NULL)); // _queueIsEmpty.clear()
    }
 
    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.");
      }
    }
 
    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:
  static PyEventLoop _loopNotFound();
 
  static PyEventLoop _getLoopOnThread(PyThreadState *tstate);
 
  static PyThreadState *_getMainThread();
  static inline PyThreadState *_getCurrentThread();
 
  // TODO (Tom Tang): use separate pools of IDs for different global objects
  static inline std::vector<AsyncHandle> _timeoutIdMap;
};
 
#endif