diff --git a/src/liburing/probe.pxd b/src/liburing/probe.pxd index 906f250..2ca3b10 100644 --- a/src/liburing/probe.pxd +++ b/src/liburing/probe.pxd @@ -1,11 +1,13 @@ -from cpython.mem cimport PyMem_RawCalloc +from cpython.mem cimport PyMem_RawCalloc, PyMem_RawFree from .lib.uring cimport * from .error cimport trap_error, memory_error from .queue cimport io_uring cdef class io_uring_probe: - cdef __io_uring_probe *ptr + cdef: + __io_uring_probe* ptr + unsigned int len cpdef io_uring_probe io_uring_get_probe_ring(io_uring ring) cpdef io_uring_probe io_uring_get_probe() diff --git a/src/liburing/probe.py b/src/liburing/probe.py index 90bdaee..f40bbba 100644 --- a/src/liburing/probe.py +++ b/src/liburing/probe.py @@ -17,8 +17,10 @@ def probe() -> dict: ''' r = {} p = io_uring_get_probe() - for i in io_uring_op: - if i.name != 'IORING_OP_LAST': - r[i.name] = io_uring_opcode_supported(p, i) - io_uring_free_probe(p) + try: + for i in io_uring_op: + if i.name != 'IORING_OP_LAST': + r[i.name] = io_uring_opcode_supported(p, i) + finally: + io_uring_free_probe(p) return r diff --git a/src/liburing/probe.pyx b/src/liburing/probe.pyx index 9934e0a..b59f6d1 100644 --- a/src/liburing/probe.pyx +++ b/src/liburing/probe.pyx @@ -2,25 +2,29 @@ cdef class io_uring_probe: def __cinit__(self, unsigned int num=0): if num: - self.ptr = <__io_uring_probe*>PyMem_RawCalloc(num, sizeof(__io_uring_probe)) + self.ptr = <__io_uring_probe*>PyMem_RawCalloc( + num, sizeof(__io_uring_probe) + num * sizeof(__io_uring_probe_op) + ) if self.ptr is NULL: memory_error(self) + self.len = num def __dealloc__(self): - if self.ptr is not NULL: - # just in case user forgets to call `io_uring_free_probe` or error happened - __io_uring_free_probe(self.ptr) + if self.len and self.ptr is not NULL: + PyMem_RawFree(self.ptr) self.ptr = NULL @property def last_op(self): if self.ptr is not NULL: return self.ptr.last_op + memory_error(self) @property def ops_len(self): if self.ptr is not NULL: return self.ptr.ops_len + memory_error(self) cpdef io_uring_probe io_uring_get_probe_ring(io_uring ring): cdef io_uring_probe probe = io_uring_probe() diff --git a/test/probe_test.py b/test/probe_test.py index 930cc72..e20b21d 100644 --- a/test/probe_test.py +++ b/test/probe_test.py @@ -1,3 +1,4 @@ +from re import escape from errno import EINVAL from pytest import raises from liburing import probe, io_uring_get_probe, io_uring_get_probe_ring, io_uring_free_probe, \ @@ -9,17 +10,17 @@ def test_probe(): assert op['IORING_OP_NOP'] is True assert op.get('IORING_OP_LAST', None) is None - # triggers `__dealloc__` - io_uring_get_probe() + p = io_uring_get_probe() + io_uring_free_probe(p) + with raises(MemoryError, match=escape('`io_uring_probe()` is out of memory!')): + p.ops_len def test_probe_ring(ring): p = io_uring_get_probe_ring(ring) - assert p.ops_len > 31 assert p.last_op > 30 assert p.last_op == p.ops_len - 1 - io_uring_free_probe(p) @@ -28,8 +29,8 @@ def test_probe_register(ring): with raises(OSError) as e: io_uring_register_probe(ring, p, 256) assert e.value.errno == EINVAL - io_uring_free_probe(p) + # free is doen by the `__dealloc__` since `num` is set p = io_uring_probe(4) assert io_uring_register_probe(ring, p, 4) == 0 - io_uring_free_probe(p) + # free is doen by the `__dealloc__` since `num` is set