casync.cc

#include <pybind11/pybind11.h>

#include <deque>
#include <iostream>
#include <mutex>

namespace py = pybind11;

class AsyncTest {
 private:
  struct Input {
    py::object loop;
    py::object future;
    int compute_id;
  };

 public:
  AsyncTest() : compute_id_(0) {}

  py::object compute() {
    py::object loop = py::module::import("asyncio").attr("get_event_loop")();
    py::object future = loop.attr("create_future")();

    std::unique_lock<std::mutex> lock(mu_);
    futures_.push_back(Input{loop, future, compute_id_++});
    return future;
  }

  void finish() {
    py::object result;
    {
      py::gil_scoped_release release;
      // Do lots of work.
    }
    if (futures_.empty()) {
      throw std::exception();
    }
    Input input = futures_.front();
    futures_.pop_front();
    result = py::str("Habemus result: " + std::to_string(input.compute_id));

    input.loop.attr("call_soon_threadsafe")(input.future.attr("set_result"),
                                            result);
  }

 private:
  std::mutex mu_;
  int compute_id_;
  std::deque<Input> futures_;
};

PYBIND11_MODULE(casync, m) {
  py::class_<AsyncTest>(m, "AsyncTest")
      .def(py::init<>())
      .def("compute", &AsyncTest::compute)
      .def("finish", &AsyncTest::finish);
}