diff --git a/src/uringcore_liburing.c b/src/uringcore_liburing.c index 79c1e3e..380a01e 100644 --- a/src/uringcore_liburing.c +++ b/src/uringcore_liburing.c @@ -8,25 +8,235 @@ #include #include -typedef struct { +typedef enum { + URINGCORE_REQUEST_PREPARED = 0, + URINGCORE_REQUEST_SUBMITTED, + URINGCORE_REQUEST_COMPLETED, + URINGCORE_REQUEST_CANCELLED, +} UringCoreRequestState; + +typedef struct UringCoreLiburingRing UringCoreLiburingRing; +typedef struct UringCoreLiburingRequest UringCoreLiburingRequest; + +struct UringCoreLiburingRequest { + PyObject_HEAD + UringCoreLiburingRing *ring; + UringCoreLiburingRequest *previous; + UringCoreLiburingRequest *next; + UringCoreRequestState state; + int result; + unsigned int flags; +}; + +struct UringCoreLiburingRing { PyObject_HEAD struct io_uring ring; + UringCoreLiburingRequest *request_head; + UringCoreLiburingRequest *request_tail; int initialized; unsigned int sq_entries; unsigned int cq_entries; unsigned int features; -} UringCoreLiburingRing; + unsigned int pending; + unsigned int prepared; +}; + +static PyTypeObject UringCoreLiburingRequestType; + +static void +uringcore_liburing_request_link( + UringCoreLiburingRing *ring, + UringCoreLiburingRequest *request) +{ + request->ring = ring; + request->previous = ring->request_tail; + request->next = NULL; + if (ring->request_tail != NULL) { + ring->request_tail->next = request; + } + else { + ring->request_head = request; + } + ring->request_tail = request; + ring->pending++; + ring->prepared++; + + /* Keep the request alive while an SQE or CQE contains its address. */ + Py_INCREF(request); +} + +static void +uringcore_liburing_request_unlink(UringCoreLiburingRequest *request) +{ + UringCoreLiburingRing *ring = request->ring; + + if (ring == NULL) { + return; + } + if (request->previous != NULL) { + request->previous->next = request->next; + } + else { + ring->request_head = request->next; + } + if (request->next != NULL) { + request->next->previous = request->previous; + } + else { + ring->request_tail = request->previous; + } + ring->pending--; + if (request->state == URINGCORE_REQUEST_PREPARED) { + ring->prepared--; + } + request->ring = NULL; + request->previous = NULL; + request->next = NULL; + + Py_DECREF(request); +} static void uringcore_liburing_ring_close_resources(UringCoreLiburingRing *self) { + UringCoreLiburingRequest *request; + if (self->initialized) { io_uring_queue_exit(&self->ring); memset(&self->ring, 0, sizeof(self->ring)); self->initialized = 0; } + while ((request = self->request_head) != NULL) { + request->state = URINGCORE_REQUEST_CANCELLED; + request->result = -ECANCELED; + request->flags = 0; + uringcore_liburing_request_unlink(request); + } + self->pending = 0; + self->prepared = 0; +} + +static void +uringcore_liburing_request_dealloc(UringCoreLiburingRequest *self) +{ + Py_TYPE(self)->tp_free((PyObject *)self); } +static PyObject * +uringcore_liburing_request_get_done( + UringCoreLiburingRequest *self, + void *Py_UNUSED(context)) +{ + return PyBool_FromLong( + self->state == URINGCORE_REQUEST_COMPLETED || + self->state == URINGCORE_REQUEST_CANCELLED); +} + +static PyObject * +uringcore_liburing_request_get_cancelled( + UringCoreLiburingRequest *self, + void *Py_UNUSED(context)) +{ + return PyBool_FromLong(self->state == URINGCORE_REQUEST_CANCELLED); +} + +static PyObject * +uringcore_liburing_request_get_result( + UringCoreLiburingRequest *self, + void *Py_UNUSED(context)) +{ + if (self->state != URINGCORE_REQUEST_COMPLETED && + self->state != URINGCORE_REQUEST_CANCELLED) { + Py_RETURN_NONE; + } + return PyLong_FromLong(self->result); +} + +static PyObject * +uringcore_liburing_request_get_state( + UringCoreLiburingRequest *self, + void *Py_UNUSED(context)) +{ + const char *state; + + switch (self->state) { + case URINGCORE_REQUEST_PREPARED: + state = "prepared"; + break; + case URINGCORE_REQUEST_SUBMITTED: + state = "submitted"; + break; + case URINGCORE_REQUEST_COMPLETED: + state = "completed"; + break; + case URINGCORE_REQUEST_CANCELLED: + state = "cancelled"; + break; + default: + PyErr_SetString(PyExc_SystemError, "request has an invalid state"); + return NULL; + } + return PyUnicode_FromString(state); +} + +static PyGetSetDef uringcore_liburing_request_getset[] = { + { + "done", + (getter)uringcore_liburing_request_get_done, + NULL, + PyDoc_STR("Whether the request reached a terminal state."), + NULL, + }, + { + "cancelled", + (getter)uringcore_liburing_request_get_cancelled, + NULL, + PyDoc_STR("Whether ring teardown cancelled the request."), + NULL, + }, + { + "result", + (getter)uringcore_liburing_request_get_result, + NULL, + PyDoc_STR("Raw CQE result, or None before completion."), + NULL, + }, + { + "state", + (getter)uringcore_liburing_request_get_state, + NULL, + PyDoc_STR("Current native request lifecycle state."), + NULL, + }, + {NULL}, +}; + +static PyMemberDef uringcore_liburing_request_members[] = { + { + "flags", + T_UINT, + offsetof(UringCoreLiburingRequest, flags), + READONLY, + PyDoc_STR("Flags from the completion queue entry."), + }, + {NULL}, +}; + +PyDoc_STRVAR( + uringcore_liburing_request_doc, + "A native request whose lifetime is owned by its ring until completion."); + +static PyTypeObject UringCoreLiburingRequestType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "uringloop._uringcore_liburing.Request", + .tp_basicsize = sizeof(UringCoreLiburingRequest), + .tp_dealloc = (destructor)uringcore_liburing_request_dealloc, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_doc = uringcore_liburing_request_doc, + .tp_members = uringcore_liburing_request_members, + .tp_getset = uringcore_liburing_request_getset, +}; + static PyObject * uringcore_liburing_ring_new( PyTypeObject *type, @@ -112,6 +322,187 @@ uringcore_liburing_ring_close( Py_RETURN_NONE; } +static PyObject * +uringcore_liburing_ring_prepare_nop( + UringCoreLiburingRing *self, + PyObject *Py_UNUSED(ignored)) +{ + UringCoreLiburingRequest *request; + struct io_uring_sqe *sqe; + + if (!self->initialized) { + PyErr_SetString(PyExc_RuntimeError, "ring is closed"); + return NULL; + } + + request = (UringCoreLiburingRequest *) + UringCoreLiburingRequestType.tp_alloc( + &UringCoreLiburingRequestType, 0); + if (request == NULL) { + return NULL; + } + request->ring = NULL; + request->previous = NULL; + request->next = NULL; + request->state = URINGCORE_REQUEST_PREPARED; + request->result = 0; + request->flags = 0; + + sqe = io_uring_get_sqe(&self->ring); + if (sqe == NULL) { + Py_DECREF(request); + PyErr_SetString( + PyExc_BufferError, + "submission queue is full; call submit() before preparing more requests"); + return NULL; + } + io_uring_prep_nop(sqe); + io_uring_sqe_set_data(sqe, request); + uringcore_liburing_request_link(self, request); + return (PyObject *)request; +} + +static PyObject * +uringcore_liburing_ring_submit( + UringCoreLiburingRing *self, + PyObject *Py_UNUSED(ignored)) +{ + UringCoreLiburingRequest *request; + unsigned int remaining; + int result; + + if (!self->initialized) { + PyErr_SetString(PyExc_RuntimeError, "ring is closed"); + return NULL; + } + + do { + result = io_uring_submit(&self->ring); + } while (result == -EINTR); + if (result < 0) { + errno = -result; + PyErr_SetFromErrno(PyExc_OSError); + return NULL; + } + + remaining = (unsigned int)result; + request = self->request_head; + while (request != NULL && remaining > 0) { + if (request->state == URINGCORE_REQUEST_PREPARED) { + request->state = URINGCORE_REQUEST_SUBMITTED; + self->prepared--; + remaining--; + } + request = request->next; + } + if (remaining != 0) { + PyErr_SetString( + PyExc_SystemError, + "liburing submitted more requests than the ring had prepared"); + return NULL; + } + return PyLong_FromLong(result); +} + +static PyObject * +uringcore_liburing_ring_reap( + UringCoreLiburingRing *self, + PyObject *args, + PyObject *kwargs) +{ + static char *keyword_names[] = {"max_completions", NULL}; + PyObject *max_completions_object = NULL; + PyObject *max_completions_index = NULL; + unsigned long parsed_max_completions; + unsigned int max_completions = 64; + PyObject *completed; + unsigned int reaped = 0; + + if (!PyArg_ParseTupleAndKeywords( + args, + kwargs, + "|O:reap", + keyword_names, + &max_completions_object)) { + return NULL; + } + if (max_completions_object != NULL) { + max_completions_index = PyNumber_Index(max_completions_object); + if (max_completions_index == NULL) { + return NULL; + } + parsed_max_completions = + PyLong_AsUnsignedLong(max_completions_index); + Py_DECREF(max_completions_index); + if (parsed_max_completions == (unsigned long)-1 && + PyErr_Occurred()) { + PyErr_Clear(); + PyErr_Format( + PyExc_ValueError, + "max_completions must be between 1 and %u", + UINT_MAX); + return NULL; + } + if (parsed_max_completions == 0 || + parsed_max_completions > UINT_MAX) { + PyErr_Format( + PyExc_ValueError, + "max_completions must be between 1 and %u", + UINT_MAX); + return NULL; + } + max_completions = (unsigned int)parsed_max_completions; + } + if (!self->initialized) { + PyErr_SetString(PyExc_RuntimeError, "ring is closed"); + return NULL; + } + + completed = PyList_New(0); + if (completed == NULL) { + return NULL; + } + while (reaped < max_completions) { + UringCoreLiburingRequest *request; + struct io_uring_cqe *cqe; + int result = io_uring_peek_cqe(&self->ring, &cqe); + + if (result == -EAGAIN) { + break; + } + if (result < 0) { + errno = -result; + PyErr_SetFromErrno(PyExc_OSError); + Py_DECREF(completed); + return NULL; + } + + request = (UringCoreLiburingRequest *)io_uring_cqe_get_data(cqe); + if (request == NULL || request->ring != self || + request->state != URINGCORE_REQUEST_SUBMITTED) { + io_uring_cqe_seen(&self->ring, cqe); + PyErr_SetString( + PyExc_SystemError, + "completion queue entry does not reference a submitted request"); + Py_DECREF(completed); + return NULL; + } + + request->state = URINGCORE_REQUEST_COMPLETED; + request->result = cqe->res; + request->flags = cqe->flags; + io_uring_cqe_seen(&self->ring, cqe); + if (PyList_Append(completed, (PyObject *)request) < 0) { + uringcore_liburing_request_unlink(request); + Py_DECREF(completed); + return NULL; + } + uringcore_liburing_request_unlink(request); + reaped++; + } + return completed; +} + static PyObject * uringcore_liburing_ring_enter( UringCoreLiburingRing *self, @@ -142,6 +533,24 @@ uringcore_liburing_ring_get_closed( } static PyMethodDef uringcore_liburing_ring_methods[] = { + { + "prepare_nop", + (PyCFunction)uringcore_liburing_ring_prepare_nop, + METH_NOARGS, + PyDoc_STR("Prepare one native no-op request without submitting it."), + }, + { + "submit", + (PyCFunction)uringcore_liburing_ring_submit, + METH_NOARGS, + PyDoc_STR("Submit the requests currently prepared on the ring."), + }, + { + "reap", + _PyCFunction_CAST(uringcore_liburing_ring_reap), + METH_VARARGS | METH_KEYWORDS, + PyDoc_STR("Return up to max_completions completed native requests."), + }, { "close", (PyCFunction)uringcore_liburing_ring_close, @@ -164,6 +573,13 @@ static PyMethodDef uringcore_liburing_ring_methods[] = { }; static PyMemberDef uringcore_liburing_ring_members[] = { + { + "pending", + T_UINT, + offsetof(UringCoreLiburingRing, pending), + READONLY, + PyDoc_STR("Number of prepared or submitted requests owned by the ring."), + }, { "sq_entries", T_UINT, @@ -232,6 +648,9 @@ PyInit__uringcore_liburing(void) { PyObject *module; + if (PyType_Ready(&UringCoreLiburingRequestType) < 0) { + return NULL; + } if (PyType_Ready(&UringCoreLiburingRingType) < 0) { return NULL; } @@ -248,7 +667,14 @@ PyInit__uringcore_liburing(void) Py_DECREF(module); return NULL; } - if (PyModule_AddIntConstant(module, "ABI_VERSION", 1) < 0) { + if (PyModule_AddObjectRef( + module, + "Request", + (PyObject *)&UringCoreLiburingRequestType) < 0) { + Py_DECREF(module); + return NULL; + } + if (PyModule_AddIntConstant(module, "ABI_VERSION", 2) < 0) { Py_DECREF(module); return NULL; } diff --git a/tests/unit/test_uringcore_liburing.py b/tests/unit/test_uringcore_liburing.py index 22b7f9a..c01e3cf 100644 --- a/tests/unit/test_uringcore_liburing.py +++ b/tests/unit/test_uringcore_liburing.py @@ -1,10 +1,24 @@ +import errno +import time + import pytest from uringloop import _uringcore_liburing def test_static_liburing_core_has_versioned_abi(): - assert _uringcore_liburing.ABI_VERSION == 1 + assert _uringcore_liburing.ABI_VERSION == 2 + + +def reap_requests(ring, count): + completed = [] + deadline = time.monotonic() + 1 + while len(completed) < count and time.monotonic() < deadline: + completed.extend(ring.reap(count - len(completed))) + if len(completed) < count: + time.sleep(0.001) + assert len(completed) == count + return completed @pytest.mark.parametrize("entries", [-(2**32), -1, 0, 2**32, 2**64]) @@ -34,3 +48,58 @@ def test_static_liburing_ring_context_manager_closes_resources(): assert ring.closed is True with pytest.raises(RuntimeError, match="ring is closed"): ring.__enter__() + + +def test_static_liburing_ring_submits_and_reaps_native_nop_requests(): + ring = _uringcore_liburing.Ring(8) + requests = [ring.prepare_nop(), ring.prepare_nop()] + + assert ring.pending == 2 + assert [request.state for request in requests] == ["prepared", "prepared"] + assert [request.result for request in requests] == [None, None] + assert ring.submit() == 2 + assert [request.state for request in requests] == ["submitted", "submitted"] + + completed = reap_requests(ring, 2) + + assert set(completed) == set(requests) + assert ring.pending == 0 + assert [request.state for request in requests] == ["completed", "completed"] + assert [request.done for request in requests] == [True, True] + assert [request.cancelled for request in requests] == [False, False] + assert [request.result for request in requests] == [0, 0] + assert [request.flags for request in requests] == [0, 0] + + +def test_static_liburing_ring_keeps_request_alive_until_completion(): + ring = _uringcore_liburing.Ring(8) + + ring.prepare_nop() + assert ring.submit() == 1 + + [request] = reap_requests(ring, 1) + assert request.state == "completed" + assert request.result == 0 + + +def test_static_liburing_ring_close_releases_prepared_request_ownership(): + ring = _uringcore_liburing.Ring(8) + request = ring.prepare_nop() + + ring.close() + + assert ring.pending == 0 + assert request.state == "cancelled" + assert request.done is True + assert request.cancelled is True + assert request.result == -errno.ECANCELED + + +@pytest.mark.parametrize("max_completions", [-1, 0, 2**32, 2**64]) +def test_static_liburing_ring_reap_rejects_out_of_range_batch_size( + max_completions, +): + ring = _uringcore_liburing.Ring(8) + + with pytest.raises(ValueError, match="max_completions must be between"): + ring.reap(max_completions) diff --git a/uringloop/_uringcore_liburing.pyi b/uringloop/_uringcore_liburing.pyi index 959461d..56078cd 100644 --- a/uringloop/_uringcore_liburing.pyi +++ b/uringloop/_uringcore_liburing.pyi @@ -2,6 +2,20 @@ from typing import Final, Literal, Self ABI_VERSION: Final[int] +class Request: + @property + def state( + self, + ) -> Literal["prepared", "submitted", "completed", "cancelled"]: ... + @property + def done(self) -> bool: ... + @property + def cancelled(self) -> bool: ... + @property + def result(self) -> int | None: ... + @property + def flags(self) -> int: ... + class Ring: def __init__(self, entries: int = 256) -> None: ... @property @@ -12,6 +26,11 @@ class Ring: def features(self) -> int: ... @property def closed(self) -> bool: ... + @property + def pending(self) -> int: ... + def prepare_nop(self) -> Request: ... + def submit(self) -> int: ... + def reap(self, max_completions: int = 64) -> list[Request]: ... def close(self) -> None: ... def __enter__(self) -> Self: ... def __exit__(