67 lines
2.8 KiB
Python
67 lines
2.8 KiB
Python
"""GCC atomic operations on aligned shared-memory fields.
|
|
|
|
The libatomic ABI omits the builtin compare-exchange ``weak`` argument:
|
|
https://github.com/gcc-mirror/gcc/blob/master/libatomic/cas_n.c
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
from functools import cache
|
|
import mmap
|
|
|
|
_ACQUIRE, _RELEASE, _ACQ_REL, _SEQ_CST = 2, 3, 4, 5
|
|
|
|
|
|
@cache
|
|
def _library() -> ctypes.CDLL:
|
|
try:
|
|
library = ctypes.CDLL("libatomic.so.1")
|
|
except OSError as exc:
|
|
raise RuntimeError("Shared memory requires libatomic.so.1 (install libatomic1)") from exc
|
|
library.__atomic_is_lock_free.argtypes = [ctypes.c_size_t, ctypes.c_void_p]
|
|
library.__atomic_is_lock_free.restype = ctypes.c_bool
|
|
for width, number in ((4, ctypes.c_uint32), (8, ctypes.c_uint64)):
|
|
if not library.__atomic_is_lock_free(width, None):
|
|
raise RuntimeError(f"Shared memory requires lock-free {width}-byte atomics")
|
|
load = getattr(library, f"__atomic_load_{width}")
|
|
load.argtypes, load.restype = [ctypes.c_void_p, ctypes.c_int], number
|
|
store = getattr(library, f"__atomic_store_{width}")
|
|
store.argtypes, store.restype = [ctypes.c_void_p, number, ctypes.c_int], None
|
|
library.__atomic_compare_exchange_8.argtypes = [
|
|
ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint64), ctypes.c_uint64, ctypes.c_int, ctypes.c_int,
|
|
]
|
|
library.__atomic_compare_exchange_8.restype = ctypes.c_bool
|
|
library.atomic_thread_fence.argtypes, library.atomic_thread_fence.restype = [ctypes.c_int], None
|
|
return library
|
|
|
|
|
|
def _address(buffer: mmap.mmap, offset: int, width: int) -> int:
|
|
if offset % width or offset < 0 or offset + width > len(buffer):
|
|
raise ValueError("Atomic field must be aligned and inside the mapping")
|
|
return ctypes.addressof(ctypes.c_char.from_buffer(buffer, offset))
|
|
|
|
|
|
def load(buffer: mmap.mmap, offset: int, width: int) -> int:
|
|
return getattr(_library(), f"__atomic_load_{width}")(_address(buffer, offset, width), _ACQUIRE)
|
|
|
|
|
|
def store(buffer: mmap.mmap, offset: int, value: int, width: int) -> None:
|
|
if not 0 <= value < 1 << (8 * width):
|
|
raise ValueError("Atomic field value exceeds its unsigned wire range")
|
|
getattr(_library(), f"__atomic_store_{width}")(_address(buffer, offset, width), value, _RELEASE)
|
|
|
|
|
|
def advance(buffer: mmap.mmap, offset: int, desired: int) -> None:
|
|
"""Advance a shared cursor with CAS; concurrent progress can never be undone."""
|
|
address = _address(buffer, offset, 8)
|
|
expected = ctypes.c_uint64(load(buffer, offset, 8))
|
|
while expected.value < desired:
|
|
if _library().__atomic_compare_exchange_8(address, ctypes.byref(expected), desired, _ACQ_REL, _ACQUIRE):
|
|
break
|
|
|
|
|
|
def fence() -> None:
|
|
"""Order invalidation/copy/validation on CPUs with weak memory ordering."""
|
|
_library().atomic_thread_fence(_SEQ_CST)
|