30852224785

Committed 03 Aug 2026 08:53PM UTC coverage: 87.218%. First build

Build # 30852224785

Build Type

Pull #69

github

Committed by

web-flow

Commit Message

Merge 3237c4d90 into d43cac632

Pull Request Pull Request #69: Fix segfault, hang, and ~30 real bugs across all 4 I/O backends

Coverage Stats

83 of 91 new or added lines in 5 files covered. (91.21%)

348 of 399 relevant lines covered (87.22%)

8.17 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.3

/caio/python_aio.py

import os
import sys
import threading
from collections import defaultdict
from collections.abc import Callable
from enum import IntEnum, unique
from multiprocessing.pool import ThreadPool
from threading import Lock, RLock
from types import MappingProxyType
from typing import Any

from .abstract import AbstractContext, AbstractOperation

fdsync = getattr(os, "fdatasync", os.fsync)
NATIVE_PREAD_PWRITE = hasattr(os, "pread") and hasattr(os, "pwrite")


@unique
class OpCode(IntEnum):
    READ = 0
    WRITE = 1
    FSYNC = 2
    FDSYNC = 3
    NOOP = -1


@unique
class ContextState(IntEnum):
    OPEN = 0
    CLOSING = 1
    CLOSED = 2


class Context(AbstractContext):
    """
    python aio context implementation
    """

    MAX_POOL_SIZE = 128

    def __init__(self, max_requests: int = 32, pool_size: int = 8):
        # Set before any validation that can raise: __del__ runs even on a
        # partially-constructed object, and close() below relies on _state
        # existing and being non-OPEN to safely no-op before ever touching
        # _lock/pool, which aren't created until validation passes.
        self._state = ContextState.CLOSED

        if not (0 < pool_size < self.MAX_POOL_SIZE):
            raise ValueError(
                f"pool_size must be between 1 and {self.MAX_POOL_SIZE - 1}, "
                f"got {pool_size}",
            )
        if max_requests <= 0:
            raise ValueError(
                f"max_requests must be a positive integer, got {max_requests}",
            )

        self.__max_requests = max_requests
        self.pool = ThreadPool(pool_size)
        self._in_progress = 0
        self._lock = Lock()
        self._state = ContextState.OPEN

        if not NATIVE_PREAD_PWRITE:
            self._locks_cleaner = RLock()       # type: ignore
            self._locks = defaultdict(RLock)    # type: ignore

    @property
    def max_requests(self) -> int:
        return self.__max_requests

    @staticmethod
    def _invoke_callback(operation: "Operation", value):
        """
        Calls ``operation``'s callback with ``value``, if one was set via
        ``set_callback()``. A missing callback is a normal, expected case.
        A raising callback must not escape - this runs on ThreadPool's
        single internal result-handler thread, and an uncaught exception
        there kills that thread, silently stalling every future
        result/callback for the rest of this Context's lifetime.
        """
        callback = operation.callback
        if callback is None:
            return

        try:
            callback(value)
        except BaseException:  # noqa: BLE001 (must isolate any callback exception, including SystemExit/KeyboardInterrupt)
            threading.excepthook(
                threading.ExceptHookArgs(
                    (*sys.exc_info(), threading.current_thread()),
                ),
            )

    def _release_slot(self):
        with self._lock:
            self._in_progress -= 1

    def _rollback_claim(self, operation: "Operation"):
        """
        Undoes a claim that was never actually scheduled (e.g. a
        concurrent close() tore down the pool between the capacity
        reservation and apply_async()) - unlike genuine completion, this
        must reset operation.in_progress, since the operation never
        actually ran and must stay retryable.
        """
        with self._lock:
            self._in_progress -= 1
            operation.in_progress = False

    def _execute(self, operation: "Operation") -> bool:
        """
        Returns True if actually scheduled, False if skipped because
        ``operation`` was already in progress - matches the other three
        backends, which all silently skip (rather than raise on, or
        dispatch twice) a resubmit of an Operation still in flight.
        """
        handler = self._OP_MAP[operation.opcode]

        # operation.in_progress is deliberately NOT reset on genuine
        # completion below - one-shot forever once actually scheduled,
        # matching all three native backends: a completed Operation must
        # not be resubmittable, only a fresh one constructed for a retry.
        def on_error(exc):
            self._release_slot()
            operation.exception = exc
            operation.written = 0
            self._invoke_callback(operation, None)

        def on_success(result):
            self._release_slot()
            operation.written = result
            self._invoke_callback(operation, result)

        # operation.in_progress is checked and set under the same lock as
        # the capacity check/reservation - otherwise two concurrent
        # submits of the very same Operation (or the same object appearing
        # twice in one submit(op, op) call) could both see it unset and
        # both dispatch, running the I/O twice against one result object.
        with self._lock:
            if operation.in_progress:
                return False

            if self._state != ContextState.OPEN:
                raise RuntimeError("Context is closed")

            if self._in_progress >= self.__max_requests:
                raise RuntimeError(
                    "Maximum simultaneous requests have been reached",
                )

            self._in_progress += 1
            operation.in_progress = True

        try:
            self.pool.apply_async(
                handler, args=(self, operation),
                callback=on_success,
                error_callback=on_error,
            )
        except BaseException:
            # Scheduling itself failed (e.g. a concurrent close() already
            # tore down the pool) - the slot reserved above was never
            # actually claimed by a real job, so it must be given back
            # instead of permanently inflating _in_progress.
            self._rollback_claim(operation)
            raise

        return True

    if NATIVE_PREAD_PWRITE:
        def __pread(self, fd, size, offset):
            return os.pread(fd, size, offset)

        def __pwrite(self, fd, bytes, offset):
            return os.pwrite(fd, bytes, offset)
    else:
        def __pread(self, fd, size, offset):
            with self._locks[fd]:
                os.lseek(fd, 0, os.SEEK_SET)
                os.lseek(fd, offset, os.SEEK_SET)
                return os.read(fd, size)

        def __pwrite(self, fd, bytes, offset):
            with self._locks[fd]:
                os.lseek(fd, 0, os.SEEK_SET)
                os.lseek(fd, offset, os.SEEK_SET)
                return os.write(fd, bytes)

    def _handle_read(self, operation: "Operation"):
        # Stored directly, not copied through a BytesIO - pread() already
        # returns exactly the bytes object get_value()/payload need to hand
        # back, and buffering it through BytesIO.write() just to unwrap it
        # again later cost a full extra copy per read for no benefit.
        data = self.__pread(
            operation.fileno, operation.nbytes, operation.offset,
        )
        operation.buffer = data
        return len(data)

    def _handle_write(self, operation: "Operation"):
        # operation.buffer is the caller's own payload bytes, unwrapped -
        # no BytesIO round-trip needed to hand it to pwrite() either.
        return self.__pwrite(
            operation.fileno, operation.buffer, operation.offset,
        )

    def _handle_fsync(self, operation: "Operation"):
        return os.fsync(operation.fileno)

    def _handle_fdsync(self, operation: "Operation"):
        return fdsync(operation.fileno)

    def _handle_noop(self, operation: "Operation"):
        return

    def submit(self, *aio_operations) -> int:
        for operation in aio_operations:
            if not isinstance(operation, Operation):
                raise ValueError(f"Invalid Operation {operation!r}")  # noqa: TRY004 (pre-existing public exception type, not changing it here)

        count = 0
        for operation in aio_operations:
            if self._execute(operation):
                count += 1

        return count

    def cancel(self, *aio_operations) -> int:
        """
        Cancels multiple Operations. Returns

         Operation.cancel(aio_op1, aio_op2, aio_opN, ...) -> int

        (Always returns zero, this method exists for compatibility reasons)
        """
        return 0

    def close(self):
        if self._state != ContextState.OPEN:
            return
        with self._lock:
            if self._state != ContextState.OPEN:
                return
            self._state = ContextState.CLOSING
            self.pool.close()
            self._state = ContextState.CLOSED

    def __del__(self):
        self.close()

    _OP_MAP = MappingProxyType({
        OpCode.READ: _handle_read,
        OpCode.WRITE: _handle_write,
        OpCode.FSYNC: _handle_fsync,
        OpCode.FDSYNC: _handle_fdsync,
        OpCode.NOOP: _handle_noop,
    })


# noinspection PyPropertyDefinition
class Operation(AbstractOperation):
    """
    python aio operation implementation
    """
    def __init__(
        self,
        fd: int,
        nbytes: int | None,
        offset: int | None,
        opcode: OpCode,
        payload: bytes | None = None,
        priority: int | None = None,
    ):
        self.callback: Callable[[int], Any] | None = None
        self.in_progress = False
        # Plain bytes, not a BytesIO wrapper - for a write this is the
        # caller's own payload, handed to pwrite() as-is (falling back to
        # b"" for a falsy/omitted payload, e.g. an explicit zero-byte
        # write); for a read it starts empty and _handle_read() replaces
        # it with pread()'s result directly. Either way there's nothing to
        # unwrap later, so get_value()/payload return it with no extra
        # copy - and it's never None, so callers don't need to check.
        self.buffer: bytes = payload or b"" if opcode == OpCode.WRITE else b""

        self.opcode = opcode
        self.__fileno = fd
        self.__offset = offset or 0
        self.__opcode = opcode
        self.__nbytes = nbytes or 0
        self.__priority = priority or 0
        self.exception = None
        self.written = 0

    @classmethod
    def read(
        cls, nbytes: int, fd: int, offset: int, priority=0,
    ) -> "Operation":
        """
        Creates a new instance of Operation on read mode.
        """
        return cls(fd, nbytes, offset, opcode=OpCode.READ, priority=priority)

    @classmethod
    def write(
        cls, payload_bytes: bytes, fd: int, offset: int, priority=0,
    ) -> "Operation":
        """
        Creates a new instance of AIOOperation on write mode.
        """
        return cls(
            fd,
            len(payload_bytes),
            offset,
            payload=payload_bytes,
            opcode=OpCode.WRITE,
            priority=priority,
        )

    @classmethod
    def fsync(cls, fd: int, priority=0) -> "Operation":

        """
        Creates a new instance of AIOOperation on fsync mode.
        """
        return cls(fd, None, None, opcode=OpCode.FSYNC, priority=priority)

    @classmethod
    def fdsync(cls, fd: int, priority=0) -> "Operation":

        """
        Creates a new instance of AIOOperation on fdsync mode.
        """
        return cls(fd, None, None, opcode=OpCode.FDSYNC, priority=priority)

    def get_value(self) -> bytes | int:
        """
        Method returns a bytes value of AIOOperation's result or None.
        """
        if self.exception:
            raise self.exception

        if self.opcode == OpCode.WRITE:
            return self.written

        return self.buffer

    @property
    def fileno(self) -> int:
        return self.__fileno

    @property
    def offset(self) -> int:
        return self.__offset

    @property
    def payload(self) -> memoryview | None:
        return memoryview(self.buffer)

    @property
    def nbytes(self) -> int:
        return self.__nbytes

    def set_callback(self, callback: Callable[[int], Any]) -> bool:
        if not callable(callback):
            raise ValueError(f"callback must be callable, got {callback!r}")  # noqa: TRY004 (pre-existing public exception type, not changing it here)
        self.callback = callback
        return True

1	import os	10✔
2	import sys	10✔
3	import threading	10✔
4	from collections import defaultdict	10✔
5	from collections.abc import Callable	10✔
6	from enum import IntEnum, unique	10✔
7	from multiprocessing.pool import ThreadPool	10✔
8	from threading import Lock, RLock	10✔
9	from types import MappingProxyType	10✔
10	from typing import Any	10✔
11
12	from .abstract import AbstractContext, AbstractOperation	10✔
13
14	fdsync = getattr(os, "fdatasync", os.fsync)	10✔
15	NATIVE_PREAD_PWRITE = hasattr(os, "pread") and hasattr(os, "pwrite")	10✔
16
17
18	@unique	10✔
19	class OpCode(IntEnum):	10✔
20	READ = 0	10✔
21	WRITE = 1	10✔
22	FSYNC = 2	10✔
23	FDSYNC = 3	10✔
24	NOOP = -1	10✔
25
26
27	@unique	10✔
28	class ContextState(IntEnum):	10✔
29	OPEN = 0	10✔
30	CLOSING = 1	10✔
31	CLOSED = 2	10✔
32
33
34	class Context(AbstractContext):	10✔
35	"""
36	python aio context implementation
37	"""
38
39	MAX_POOL_SIZE = 128	10✔
40
41	def __init__(self, max_requests: int = 32, pool_size: int = 8):	10✔
42	# Set before any validation that can raise: __del__ runs even on a
43	# partially-constructed object, and close() below relies on _state
44	# existing and being non-OPEN to safely no-op before ever touching
45	# _lock/pool, which aren't created until validation passes.
46	self._state = ContextState.CLOSED	10✔
47
48	if not (0 < pool_size < self.MAX_POOL_SIZE):	10✔
49	raise ValueError(	10✔
50	f"pool_size must be between 1 and {self.MAX_POOL_SIZE - 1}, "
51	f"got {pool_size}",
52	)
53	if max_requests <= 0:	10✔
54	raise ValueError(	10✔
55	f"max_requests must be a positive integer, got {max_requests}",
56	)
57
58	self.__max_requests = max_requests	10✔
59	self.pool = ThreadPool(pool_size)	10✔
60	self._in_progress = 0	10✔
61	self._lock = Lock()	10✔
62	self._state = ContextState.OPEN	10✔
63
64	if not NATIVE_PREAD_PWRITE:	10✔
65	self._locks_cleaner = RLock() # type: ignore	×
66	self._locks = defaultdict(RLock) # type: ignore	×
67
68	@property	10✔
69	def max_requests(self) -> int:	10✔
70	return self.__max_requests	×
71
72	@staticmethod	10✔
73	def _invoke_callback(operation: "Operation", value):	10✔
74	"""
75	Calls ``operation``'s callback with ``value``, if one was set via
76	``set_callback()``. A missing callback is a normal, expected case.
77	A raising callback must not escape - this runs on ThreadPool's
78	single internal result-handler thread, and an uncaught exception
79	there kills that thread, silently stalling every future
80	result/callback for the rest of this Context's lifetime.
81	"""
82	callback = operation.callback	10✔
83	if callback is None:	10✔
84	return	10✔
85
86	try:	10✔
87	callback(value)	10✔
88	except BaseException: # noqa: BLE001 (must isolate any callback exception, including SystemExit/KeyboardInterrupt)	10✔
89	threading.excepthook(	10✔
90	threading.ExceptHookArgs(
91	(*sys.exc_info(), threading.current_thread()),
92	),
93	)
94
95	def _release_slot(self):	10✔
96	with self._lock:	10✔
97	self._in_progress -= 1	10✔
98
99	def _rollback_claim(self, operation: "Operation"):	10✔
100	"""
101	Undoes a claim that was never actually scheduled (e.g. a
102	concurrent close() tore down the pool between the capacity
103	reservation and apply_async()) - unlike genuine completion, this
104	must reset operation.in_progress, since the operation never
105	actually ran and must stay retryable.
106	"""
107	with self._lock:	10✔
108	self._in_progress -= 1	10✔
109	operation.in_progress = False	10✔
110
111	def _execute(self, operation: "Operation") -> bool:	10✔
112	"""
113	Returns True if actually scheduled, False if skipped because
114	``operation`` was already in progress - matches the other three
115	backends, which all silently skip (rather than raise on, or
116	dispatch twice) a resubmit of an Operation still in flight.
117	"""
118	handler = self._OP_MAP[operation.opcode]	10✔
119
120	# operation.in_progress is deliberately NOT reset on genuine
121	# completion below - one-shot forever once actually scheduled,
122	# matching all three native backends: a completed Operation must
123	# not be resubmittable, only a fresh one constructed for a retry.
124	def on_error(exc):	10✔
125	self._release_slot()	10✔
126	operation.exception = exc	10✔
127	operation.written = 0	10✔
128	self._invoke_callback(operation, None)	10✔
129
130	def on_success(result):	10✔
131	self._release_slot()	10✔
132	operation.written = result	10✔
133	self._invoke_callback(operation, result)	10✔
134
135	# operation.in_progress is checked and set under the same lock as
136	# the capacity check/reservation - otherwise two concurrent
137	# submits of the very same Operation (or the same object appearing
138	# twice in one submit(op, op) call) could both see it unset and
139	# both dispatch, running the I/O twice against one result object.
140	with self._lock:	10✔
141	if operation.in_progress:	10✔
142	return False	10✔
143
144	if self._state != ContextState.OPEN:	10✔
145	raise RuntimeError("Context is closed")	10✔
146
147	if self._in_progress >= self.__max_requests:	10✔
148	raise RuntimeError(	10✔
149	"Maximum simultaneous requests have been reached",
150	)
151
152	self._in_progress += 1	10✔
153	operation.in_progress = True	10✔
154
155	try:	10✔
156	self.pool.apply_async(	10✔
157	handler, args=(self, operation),
158	callback=on_success,
159	error_callback=on_error,
160	)
161	except BaseException:	10✔
162	# Scheduling itself failed (e.g. a concurrent close() already
163	# tore down the pool) - the slot reserved above was never
164	# actually claimed by a real job, so it must be given back
165	# instead of permanently inflating _in_progress.
166	self._rollback_claim(operation)	10✔
167	raise	10✔
168
169	return True	10✔
170
171	if NATIVE_PREAD_PWRITE:	10✔
172	def __pread(self, fd, size, offset):	10✔
173	return os.pread(fd, size, offset)	10✔
174
175	def __pwrite(self, fd, bytes, offset):	10✔
176	return os.pwrite(fd, bytes, offset)	10✔
177	else:
178	def __pread(self, fd, size, offset):	×
179	with self._locks[fd]:	×
180	os.lseek(fd, 0, os.SEEK_SET)	×
181	os.lseek(fd, offset, os.SEEK_SET)	×
182	return os.read(fd, size)	×
183
184	def __pwrite(self, fd, bytes, offset):	×
185	with self._locks[fd]:	×
186	os.lseek(fd, 0, os.SEEK_SET)	×
187	os.lseek(fd, offset, os.SEEK_SET)	×
188	return os.write(fd, bytes)	×
189
190	def _handle_read(self, operation: "Operation"):	10✔
191	# Stored directly, not copied through a BytesIO - pread() already
192	# returns exactly the bytes object get_value()/payload need to hand
193	# back, and buffering it through BytesIO.write() just to unwrap it
194	# again later cost a full extra copy per read for no benefit.
195	data = self.__pread(	10✔
196	operation.fileno, operation.nbytes, operation.offset,
197	)
198	operation.buffer = data	10✔
199	return len(data)	10✔
200
201	def _handle_write(self, operation: "Operation"):	10✔
202	# operation.buffer is the caller's own payload bytes, unwrapped -
203	# no BytesIO round-trip needed to hand it to pwrite() either.
204	return self.__pwrite(	10✔
205	operation.fileno, operation.buffer, operation.offset,
206	)
207
208	def _handle_fsync(self, operation: "Operation"):	10✔
209	return os.fsync(operation.fileno)	10✔
210
211	def _handle_fdsync(self, operation: "Operation"):	10✔
212	return fdsync(operation.fileno)	10✔
213
214	def _handle_noop(self, operation: "Operation"):	10✔
215	return	×
216
217	def submit(self, *aio_operations) -> int:	10✔
218	for operation in aio_operations:	10✔
219	if not isinstance(operation, Operation):	10✔
NEW 220	raise ValueError(f"Invalid Operation {operation!r}") # noqa: TRY004 (pre-existing public exception type, not changing it here)	×
221
222	count = 0	10✔
223	for operation in aio_operations:	10✔
224	if self._execute(operation):	10✔
225	count += 1	10✔
226
227	return count	10✔
228
229	def cancel(self, *aio_operations) -> int:	10✔
230	"""
231	Cancels multiple Operations. Returns
232
233	Operation.cancel(aio_op1, aio_op2, aio_opN, ...) -> int
234
235	(Always returns zero, this method exists for compatibility reasons)
236	"""
237	return 0	10✔
238
239	def close(self):	10✔
240	if self._state != ContextState.OPEN:	10✔
241	return	10✔
242	with self._lock:	10✔
243	if self._state != ContextState.OPEN:	10✔
NEW 244	return	×
245	self._state = ContextState.CLOSING	10✔
246	self.pool.close()	10✔
247	self._state = ContextState.CLOSED	10✔
248
249	def __del__(self):	10✔
250	self.close()	10✔
251
252	_OP_MAP = MappingProxyType({	10✔
253	OpCode.READ: _handle_read,
254	OpCode.WRITE: _handle_write,
255	OpCode.FSYNC: _handle_fsync,
256	OpCode.FDSYNC: _handle_fdsync,
257	OpCode.NOOP: _handle_noop,
258	})
259
260
261	# noinspection PyPropertyDefinition
262	class Operation(AbstractOperation):	10✔
263	"""
264	python aio operation implementation
265	"""
266	def __init__(	10✔
267	self,
268	fd: int,
269	nbytes: int \| None,
270	offset: int \| None,
271	opcode: OpCode,
272	payload: bytes \| None = None,
273	priority: int \| None = None,
274	):
275	self.callback: Callable[[int], Any] \| None = None	10✔
276	self.in_progress = False	10✔
277	# Plain bytes, not a BytesIO wrapper - for a write this is the
278	# caller's own payload, handed to pwrite() as-is (falling back to
279	# b"" for a falsy/omitted payload, e.g. an explicit zero-byte
280	# write); for a read it starts empty and _handle_read() replaces
281	# it with pread()'s result directly. Either way there's nothing to
282	# unwrap later, so get_value()/payload return it with no extra
283	# copy - and it's never None, so callers don't need to check.
284	self.buffer: bytes = payload or b"" if opcode == OpCode.WRITE else b""	10✔
285
286	self.opcode = opcode	10✔
287	self.__fileno = fd	10✔
288	self.__offset = offset or 0	10✔
289	self.__opcode = opcode	10✔
290	self.__nbytes = nbytes or 0	10✔
291	self.__priority = priority or 0	10✔
292	self.exception = None	10✔
293	self.written = 0	10✔
294
295	@classmethod	10✔
296	def read(	10✔
297	cls, nbytes: int, fd: int, offset: int, priority=0,
298	) -> "Operation":
299	"""
300	Creates a new instance of Operation on read mode.
301	"""
302	return cls(fd, nbytes, offset, opcode=OpCode.READ, priority=priority)	10✔
303
304	@classmethod	10✔
305	def write(	10✔
306	cls, payload_bytes: bytes, fd: int, offset: int, priority=0,
307	) -> "Operation":
308	"""
309	Creates a new instance of AIOOperation on write mode.
310	"""
311	return cls(	10✔
312	fd,
313	len(payload_bytes),
314	offset,
315	payload=payload_bytes,
316	opcode=OpCode.WRITE,
317	priority=priority,
318	)
319
320	@classmethod	10✔
321	def fsync(cls, fd: int, priority=0) -> "Operation":	10✔
322
323	"""
324	Creates a new instance of AIOOperation on fsync mode.
325	"""
326	return cls(fd, None, None, opcode=OpCode.FSYNC, priority=priority)	10✔
327
328	@classmethod	10✔
329	def fdsync(cls, fd: int, priority=0) -> "Operation":	10✔
330
331	"""
332	Creates a new instance of AIOOperation on fdsync mode.
333	"""
334	return cls(fd, None, None, opcode=OpCode.FDSYNC, priority=priority)	10✔
335
336	def get_value(self) -> bytes \| int:	10✔
337	"""
338	Method returns a bytes value of AIOOperation's result or None.
339	"""
340	if self.exception:	10✔
341	raise self.exception	10✔
342
343	if self.opcode == OpCode.WRITE:	10✔
344	return self.written	10✔
345
346	return self.buffer	10✔
347
348	@property	10✔
349	def fileno(self) -> int:	10✔
350	return self.__fileno	10✔
351
352	@property	10✔
353	def offset(self) -> int:	10✔
354	return self.__offset	10✔
355
356	@property	10✔
357	def payload(self) -> memoryview \| None:	10✔
358	return memoryview(self.buffer)	10✔
359
360	@property	10✔
361	def nbytes(self) -> int:	10✔
362	return self.__nbytes	10✔
363
364	def set_callback(self, callback: Callable[[int], Any]) -> bool:	10✔
365	if not callable(callback):	10✔
366	raise ValueError(f"callback must be callable, got {callback!r}") # noqa: TRY004 (pre-existing public exception type, not changing it here)	10✔
367	self.callback = callback	10✔
368	return True	10✔

mosquito / caio / 30852224785

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous