python_future_task.h

#pragma once
 
#include <chrono>
#include <future>
#include <memory>
#include <mutex>
#include <queue>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
 
#include <iostream>
 
#include <Python.h>
 
#include <ucxx/log.h>
#include <ucxx/python/future.h>
#include <ucxx/python/python_future_task_collector.h>
 
namespace ucxx {
 
namespace python {
 
namespace detail {

template <typename ReturnType, typename... TaskArgs>
struct PythonFutureTask {
 public:
  std::packaged_task<ReturnType(TaskArgs...)> _task{};  
  std::function<PyObject*(ReturnType)>
    _pythonConvert{};                 
  PyObject* _asyncioEventLoop{};      
  PyObject* _handle{};                
  std::future<ReturnType> _future{};  
 
 private:
  void setResult(const ReturnType result)
  {
    // PyLong_FromSize_t requires the GIL
    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
 
    PyGILState_STATE state = PyGILState_Ensure();
    ucxx::python::future_set_result_with_event_loop(
      _asyncioEventLoop, _handle, _pythonConvert(result));
    PyGILState_Release(state);
  }

  void setPythonException(PyObject* pythonException, const std::string& message)
  {
    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
 
    PyGILState_STATE state = PyGILState_Ensure();
    ucxx::python::future_set_exception_with_event_loop(
      _asyncioEventLoop, _handle, pythonException, message.c_str());
    PyGILState_Release(state);
  }

  void setException(const std::exception& exception)
  {
    try {
      throw exception;
    } catch (const std::bad_alloc& e) {
      setPythonException(PyExc_MemoryError, e.what());
    } catch (const std::bad_cast& e) {
      setPythonException(PyExc_TypeError, e.what());
    } catch (const std::bad_typeid& e) {
      setPythonException(PyExc_TypeError, e.what());
    } catch (const std::domain_error& e) {
      setPythonException(PyExc_ValueError, e.what());
    } catch (const std::invalid_argument& e) {
      setPythonException(PyExc_ValueError, e.what());
    } catch (const std::ios_base::failure& e) {
      setPythonException(PyExc_IOError, e.what());
    } catch (const std::out_of_range& e) {
      setPythonException(PyExc_IndexError, e.what());
    } catch (const std::overflow_error& e) {
      setPythonException(PyExc_OverflowError, e.what());
    } catch (const std::range_error& e) {
      setPythonException(PyExc_ArithmeticError, e.what());
    } catch (const std::underflow_error& e) {
      setPythonException(PyExc_ArithmeticError, e.what());
    } catch (const std::exception& e) {
      setPythonException(PyExc_RuntimeError, e.what());
    } catch (...) {
      setPythonException(PyExc_RuntimeError, "Unknown exception");
    }
  }
 
 public:
  explicit PythonFutureTask(std::packaged_task<ReturnType(TaskArgs...)> task,
                            std::function<PyObject*(ReturnType)> pythonConvert,
                            PyObject* asyncioEventLoop,
                            std::launch launchPolicy = std::launch::async)
    : _task{std::move(task)},
      _pythonConvert(std::move(pythonConvert)),
      _asyncioEventLoop(asyncioEventLoop),
      _handle{ucxx::python::create_python_future_with_event_loop(asyncioEventLoop)},
      _future{std::async(launchPolicy, [this]() {
        std::future<ReturnType> result = this->_task.get_future();
        this->_task();
        try {
          const ReturnType r = result.get();
          this->setResult(r);
          return r;
        } catch (std::exception& e) {
          this->setException(e);
        }
      })}
  {
  }
 
  PythonFutureTask(const PythonFutureTask&)            = delete;
  PythonFutureTask& operator=(PythonFutureTask const&) = delete;

  PythonFutureTask(PythonFutureTask&& o) = default;

  PythonFutureTask& operator=(PythonFutureTask&& o) = default;

  ~PythonFutureTask() { ucxx::python::PythonFutureTaskCollector::get().push(_handle); }

  [[nodiscard]] std::future<ReturnType>& getFuture()
  {
    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
 
    return _future;
  }

  [[nodiscard]] PyObject* getHandle()
  {
    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
 
    return _handle;
  }

  [[nodiscard]] PyObject* release()
  {
    if (_handle == nullptr) throw std::runtime_error("Invalid object or already released");
 
    return std::exchange(_handle, nullptr);
  }
};
 
}  // namespace detail

template <typename ReturnType, typename... TaskArgs>
class PythonFutureTask : public std::enable_shared_from_this<PythonFutureTask<ReturnType>> {
 private:
  std::unique_ptr<detail::PythonFutureTask<ReturnType, TaskArgs...>> _detail{
    nullptr};  
 
 public:
  explicit PythonFutureTask(std::packaged_task<ReturnType(TaskArgs...)> task,
                            std::function<PyObject*(ReturnType)> pythonConvert,
                            PyObject* asyncioEventLoop,
                            std::launch launchPolicy = std::launch::async)
    : _detail(std::make_unique<detail::PythonFutureTask<ReturnType, TaskArgs...>>(
        std::move(task), std::move(pythonConvert), asyncioEventLoop, launchPolicy))
  {
  }
  PythonFutureTask(const PythonFutureTask&)            = delete;
  PythonFutureTask& operator=(PythonFutureTask const&) = delete;

  PythonFutureTask(PythonFutureTask&& o) = default;

  PythonFutureTask& operator=(PythonFutureTask&& o) = default;

  [[nodiscard]] std::future<ReturnType>& getFuture() { return _detail->getFuture(); }

  [[nodiscard]] PyObject* getHandle() { return _detail->getHandle(); }

  [[nodiscard]] PyObject* release() { return _detail->release(); }
};
 
}  // namespace python
 
}  // namespace ucxx