xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(diskann): io uring backend (#599) 

Co-authored-by: ZeFeng Yin <yinzefeng.yzf@alibaba-inc.com>
This commit is contained in:
rayx 2026-08-07 17:44:58 +08:00 committed by GitHub
parent e1f11e29fe
commit 3b1a10c38d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 1001 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
.github/workflows/03-macos-linux-build.yml vendored
View File

@ -177,9 +177,31 @@ jobs:
shell: bash
# ------------------------------------------------------------------ #
# DiskAnn libaio round: install libaio and re-run tests
# DiskAnn io_uring round: GitHub-hosted runners are full VMs on
# kernel 6.8+, so raw io_uring syscalls are available and win the
# backend priority (io_uring > libaio > pread). The main test
# rounds above therefore already exercised the io_uring path;
# this step asserts it so a silent probe regression fails CI.
# ------------------------------------------------------------------ #
- name: Install libaio runtime
- name: Verify io_uring backend is active
if: runner.os == 'Linux'
run: |
uname -r
python -c "
import zvec
from zvec.typing import IOBackendType
t = zvec.io_backend_type()
print('Active I/O backend:', t, '-', zvec.io_backend_description())
assert t == IOBackendType.IO_URING, f'expected IO_URING, got {t}'
"
shell: bash
# ------------------------------------------------------------------ #
# DiskAnn libaio round: disable io_uring via sysctl (kernel 6.6+,
# makes io_uring_setup() fail with -EPERM), install libaio, and
# re-run tests so the libaio fallback path is genuinely exercised.
# ------------------------------------------------------------------ #
- name: Install libaio runtime and disable io_uring
if: matrix.platform == 'linux-x64'
run: |
sudo apt-get update -y
@ -188,6 +210,14 @@ jobs:
echo "=== libaio status (should be present) ==="
dpkg -l | grep -i libaio || true
ldconfig -p | grep libaio || true
sudo sysctl -w kernel.io_uring_disabled=2
python -c "
import zvec
from zvec.typing import IOBackendType
t = zvec.io_backend_type()
print('Active I/O backend:', t)
assert t == IOBackendType.LIBAIO, f'expected LIBAIO, got {t}'
"
shell: bash
- name: Run DiskAnn C++ Tests (w/ libaio)
@ -204,6 +234,32 @@ jobs:
python -m pytest python/tests/test_collection_diskann.py -v
shell: bash
# ------------------------------------------------------------------ #
# DiskAnn pread round: with io_uring still disabled, remove libaio
# so the final synchronous pread() fallback is exercised too.
# ------------------------------------------------------------------ #
- name: Remove libaio (force pread fallback)
if: matrix.platform == 'linux-x64'
run: |
sudo apt-get remove -y libaio1t64 libaio1 2>/dev/null || \
sudo apt-get remove -y libaio1t64 || \
sudo apt-get remove -y libaio1
python -c "
import zvec
from zvec.typing import IOBackendType
t = zvec.io_backend_type()
print('Active I/O backend:', t)
assert t == IOBackendType.PREAD, f'expected PREAD, got {t}'
"
shell: bash
- name: Run DiskAnn C++ Tests (w/ pread)
if: matrix.platform == 'linux-x64'
run: |
cd "$GITHUB_WORKSPACE/build"
ctest -R diskann --output-on-failure --parallel $NPROC
shell: bash
# Verify installing libaio does not affect existing non-DiskAnn tests.
- name: Run HNSW Tests (w/ libaio)
if: matrix.platform == 'linux-x64'

2
python/tests/test_typing.py
View File

@ -112,7 +112,7 @@ def test_index_type_has_member(member):
assert member in IndexType.__members__
@pytest.mark.parametrize("member", ["PREAD", "LIBAIO"])
@pytest.mark.parametrize("member", ["PREAD", "LIBAIO", "IO_URING"])
def test_io_backend_type_has_member(member):
assert member in IOBackendType.__members__

4
python/zvec/__init__.pyi
View File

@ -52,7 +52,9 @@ from .zvec import create_and_open, init, open
def io_backend_type() -> IOBackendType:
"""Returns the current I/O backend type for DiskAnn async disk reads
as an IOBackendType enum (zvec.typing.IOBackendType).
IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise."""
IOBackendType.IO_URING if io_uring is available,
IOBackendType.LIBAIO if libaio is available,
IOBackendType.PREAD otherwise."""
def io_backend_description() -> str:
"""Returns a human-readable description of the current I/O backend.

6
python/zvec/typing/__init__.pyi
View File

@ -130,6 +130,7 @@ class IOBackendType:
- PREAD: Synchronous pread() — no async I/O.
- LIBAIO: libaio loaded at runtime via dlopen().
- IO_URING: io_uring via raw kernel syscalls (zero dependency).
Examples:
>>> from zvec.typing import IOBackendType
@ -142,13 +143,16 @@ class IOBackendType:
PREAD
LIBAIO
IO_URING
"""
IO_URING: typing.ClassVar[IOBackendType] # value = <IOBackendType.IO_URING: 2>
LIBAIO: typing.ClassVar[IOBackendType] # value = <IOBackendType.LIBAIO: 1>
PREAD: typing.ClassVar[IOBackendType] # value = <IOBackendType.PREAD: 0>
__members__: typing.ClassVar[
dict[str, IOBackendType]
] # value = {'PREAD': <IOBackendType.PREAD: 0>, 'LIBAIO': <IOBackendType.LIBAIO: 1>}
] # value = {'PREAD': <IOBackendType.PREAD: 0>, 'LIBAIO': <IOBackendType.LIBAIO: 1>, 'IO_URING': <IOBackendType.IO_URING: 2>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...

102
src/ailego/io/io_backend_def.h
View File

@ -12,32 +12,36 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// Abstract I/O backend selector.
// Abstract I/O backend selector — internal header.
//
// Wraps the low-level loaders (LibAioLoader for libaio) and provides a uniform
// way to initialize, query, and report the active I/O backend. The actual I/O
// operations are still performed by the underlying loaders; this class is
// responsible only for backend initialization and reporting.
// Wraps the low-level backends (io_uring via raw syscalls, LibAioLoader for
// libaio) and provides a uniform way to initialize, query, and report the
// active I/O backend. The actual I/O operations are still performed by the
// underlying backends; this class is responsible only for backend
// initialization and reporting.
//
// When no async backend is available, the caller should fall back to
// synchronous pread().
//
// Usage:
// auto& backend = ailego::IOBackend::Instance();
// if (!backend.is_pread()) { ... }
// LOG_INFO("I/O backend: %s", backend.name());
#pragma once
#include <ailego/io/libaio_loader.h>
#include <zvec/ailego/io/io_backend.h>
#if defined(__linux) || defined(__linux__)
#include <unistd.h> // ::syscall(), ::close() — POSIX only
#include <cstring> // std::memset
#include <ailego/io/iouring_def.h> // io_uring_params, __NR_io_uring_setup
#endif
namespace zvec {
namespace ailego {
// Returns a human-readable name for the given backend type.
inline const char *IOBackendTypeName(IOBackendType type) {
switch (type) {
case IOBackendType::kIoUring:
return "io_uring";
case IOBackendType::kLibAio:
return "libaio";
case IOBackendType::kPread:
@ -50,6 +54,9 @@ inline const char *IOBackendTypeName(IOBackendType type) {
// When the backend is kPread, includes installation guidance for libaio.
inline const char *IOBackendDescription(IOBackendType type) {
switch (type) {
case IOBackendType::kIoUring:
return "io_uring async I/O backend (raw kernel syscalls, zero "
"dependency).";
case IOBackendType::kLibAio:
return "libaio async I/O backend loaded at runtime via dlopen().";
case IOBackendType::kPread:
@ -61,12 +68,12 @@ inline const char *IOBackendDescription(IOBackendType type) {
return "Unknown I/O backend.";
}
// Singleton that loads and queries an I/O backend on demand.
// Singleton that probes and caches the I/O backend on first use.
//
// available() (no arg) tries the best backend with priority (libaio > pread)
// and returns the loaded backend type.
// available(IOBackendType) tries a specific backend.
// Use type() / name() to query the loaded backend without triggering a load.
// available() probes backends by priority (io_uring > libaio > pread)
// exactly once and caches the result — including the pread-only outcome,
// so systems without async I/O don't re-probe on every call.
// Use type() / name() to query the cached backend without probing.
class IOBackend {
public:
static IOBackend &Instance() {
@ -74,33 +81,26 @@ class IOBackend {
return instance;
}
// Try to load the best available backend (libaio > pread).
// Returns the loaded backend type.
// Idempotent — if already loaded, returns immediately.
// Returns the active backend, probing on the first call
// (io_uring > libaio > pread). Idempotent — later calls return the
// cached result immediately, even when the outcome is kPread.
IOBackendType available() {
if (type_ != IOBackendType::kPread) {
return type_;
}
return available(IOBackendType::kLibAio);
}
// Try to load the requested backend. Returns the loaded backend type
// (may differ from requested if the load failed — falls back to kPread).
// Idempotent — if the same backend is already loaded, returns immediately.
IOBackendType available(IOBackendType requested) {
if (type_ == requested && type_ != IOBackendType::kPread) {
if (probed_) {
return type_;
}
#if defined(__linux) || defined(__linux__)
if (requested == IOBackendType::kLibAio) {
if (LibAioLoader::Instance().load() &&
LibAioLoader::Instance().is_available()) {
type_ = IOBackendType::kLibAio;
return type_;
}
if (io_uring_supported()) {
type_ = IOBackendType::kIoUring;
} else if (LibAioLoader::Instance().load() &&
LibAioLoader::Instance().is_available()) {
type_ = IOBackendType::kLibAio;
} else {
type_ = IOBackendType::kPread;
}
#endif
#else
type_ = IOBackendType::kPread;
#endif
probed_ = true;
return type_;
}
@ -112,7 +112,11 @@ class IOBackend {
return available() == IOBackendType::kLibAio;
}
// Returns the loaded backend type.
bool is_io_uring() {
return available() == IOBackendType::kIoUring;
}
// Returns the cached backend type without triggering the probe.
IOBackendType type() const {
return type_;
}
@ -130,7 +134,31 @@ class IOBackend {
private:
IOBackend() = default;
#if defined(__linux) || defined(__linux__)
// Probe io_uring availability with a minimal ring setup using only raw
// syscalls — no dependency on liburing. A successful setup alone is NOT
// sufficient: io_uring_setup() exists since Linux 5.1, but the read path
// uses IORING_OP_READ, which was only added in 5.6. We therefore also
// require IORING_FEAT_RW_CUR_POS in params.features — a feature flag
// introduced in the same 5.6 release — so kernels 5.1–5.5 fall back to
// libaio/pread instead of failing every read with -EINVAL.
static bool io_uring_supported() {
struct io_uring_params params;
std::memset(&params, 0, sizeof(params));
int fd = static_cast<int>(::syscall(__NR_io_uring_setup, 1, &params));
if (fd < 0) {
return false;
}
::close(fd);
return (params.features & IORING_FEAT_RW_CUR_POS) != 0;
}
#endif
// kPread doubles as the pre-probe default; probed_ marks whether the
// one-shot probe has run so that a pread-only outcome is cached too.
// (IOBackendType values are C ABI — no kNone sentinel is added there.)
IOBackendType type_{IOBackendType::kPread};
bool probed_{false};
};
} // namespace ailego

200
src/ailego/io/iouring_def.h Normal file
View File

@ -0,0 +1,200 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Private header defining the io_uring kernel ABI structures and constants.
//
// This header is the io_uring analogue of libaio_def.h: it declares *only*
// the types, constants, and inline helpers that zvec needs from the io_uring
// kernel interface. By defining these structures ourselves we avoid any
// build-time dependency on <linux/io_uring.h> or liburing-dev, mirroring the
// project's zero-dependency philosophy established by the libaio dlopen
// approach.
//
// The struct layouts (io_uring_sqe, io_uring_cqe, io_uring_params,
// io_sqring_offsets, io_cqring_offsets) are part of the Linux kernel ABI
// and are copied verbatim from <linux/io_uring.h>.
#pragma once
#include <cstdint>
#if defined(__linux) || defined(__linux__)
// ---------------------------------------------------------------------------
// Syscall numbers
// ---------------------------------------------------------------------------
// io_uring was introduced in Linux 5.1 (2019). The three syscalls share the
// same numbers across all supported architectures. We prefer the values
// from <sys/syscall.h> when available and fall back to hardcoded numbers.
#include <sys/syscall.h>
#ifndef __NR_io_uring_setup
#define __NR_io_uring_setup 425
#endif
#ifndef __NR_io_uring_enter
#define __NR_io_uring_enter 426
#endif
#ifndef __NR_io_uring_register
#define __NR_io_uring_register 427
#endif
// ---------------------------------------------------------------------------
// Constants (from <linux/io_uring.h>)
// ---------------------------------------------------------------------------
// mmap offsets for the three shared regions.
#define IORING_OFF_SQ_RING 0ULL
#define IORING_OFF_CQ_RING 0x8000000ULL
#define IORING_OFF_SQES 0x10000000ULL
// io_uring_enter flags.
#define IORING_ENTER_GETEVENTS (1U << 0)
// io_uring_params.features flags reported by io_uring_setup().
// IORING_FEAT_RW_CUR_POS was introduced in Linux 5.6 — the same release
// that added IORING_OP_READ — so its presence proves the kernel supports
// the opcode our read path relies on.
#define IORING_FEAT_RW_CUR_POS (1U << 3)
// SQE opcode values.
#define IORING_OP_NOP 0
#define IORING_OP_READV 1
#define IORING_OP_WRITEV 2
#define IORING_OP_FSYNC 3
#define IORING_OP_READ_FIXED 4
#define IORING_OP_WRITE_FIXED 5
#define IORING_OP_POLL_ADD 6
#define IORING_OP_POLL_REMOVE 7
#define IORING_OP_SYNC_FILE_RANGE 8
#define IORING_OP_SENDMSG 9
#define IORING_OP_RECVMSG 10
#define IORING_OP_TIMEOUT 11
#define IORING_OP_TIMEOUT_REMOVE 12
#define IORING_OP_ACCEPT 13
#define IORING_OP_ASYNC_CANCEL 14
#define IORING_OP_LINK_TIMEOUT 15
#define IORING_OP_CONNECT 16
#define IORING_OP_FALLOCATE 17
#define IORING_OP_OPENAT 18
#define IORING_OP_CLOSE 19
#define IORING_OP_FILES_UPDATE 20
#define IORING_OP_STATX 21
#define IORING_OP_READ 22
#define IORING_OP_WRITE 23
// ---------------------------------------------------------------------------
// Struct definitions (copied verbatim from <linux/io_uring.h>)
// ---------------------------------------------------------------------------
// Submission queue entry — 64 bytes.
struct io_uring_sqe {
uint8_t opcode; // type of operation for this sqe
uint8_t flags; // IOSQE_ flags
uint16_t ioprio; // ioprio for the request
int32_t fd; // file descriptor to do IO on
union {
uint64_t off; // offset into file
uint64_t addr2;
};
union {
uint64_t addr; // buffer or iovecs
uint64_t splice_off_in;
};
uint32_t len; // buffer size or number of iovecs
union {
uint32_t rw_flags; // read/write flags (union of all flag types)
};
uint64_t user_data; // data to be passed back at completion time
union {
struct {
uint16_t buf_index; // index into fixed buffers, if used
uint16_t personality;
} buf;
uint64_t __pad2[3];
};
};
// Completion queue entry — 16 bytes.
struct io_uring_cqe {
uint64_t user_data; // sqe->user_data
int32_t res; // result code for this event
uint32_t flags;
};
// SQ ring offsets — returned by io_uring_setup in io_uring_params.
struct io_sqring_offsets {
uint32_t head;
uint32_t tail;
uint32_t ring_mask;
uint32_t ring_entries;
uint32_t flags;
uint32_t dropped;
uint32_t array;
uint32_t resv1;
uint64_t resv2;
};
// CQ ring offsets — returned by io_uring_setup in io_uring_params.
struct io_cqring_offsets {
uint32_t head;
uint32_t tail;
uint32_t ring_mask;
uint32_t ring_entries;
uint32_t overflow;
uint32_t cqes;
uint32_t flags;
uint32_t resv1;
uint64_t resv2;
};
// Parameters passed to io_uring_setup().
struct io_uring_params {
uint32_t sq_entries;
uint32_t cq_entries;
uint32_t flags;
uint32_t sq_thread_cpu;
uint32_t sq_thread_idle;
uint32_t features;
uint32_t wq_fd;
uint32_t resv[3];
struct io_sqring_offsets sq_off;
struct io_cqring_offsets cq_off;
};
// ---------------------------------------------------------------------------
// Inline helper — prepare an SQE for a read operation.
// ---------------------------------------------------------------------------
static inline void io_uring_prep_read(struct io_uring_sqe *sqe, int fd,
void *buf, uint32_t nbytes,
uint64_t offset) {
sqe->opcode = IORING_OP_READ;
sqe->flags = 0;
sqe->ioprio = 0;
sqe->fd = fd;
sqe->off = offset;
sqe->addr = reinterpret_cast<uint64_t>(buf);
sqe->len = nbytes;
sqe->rw_flags = 0;
sqe->user_data = 0;
sqe->buf.buf_index = 0;
sqe->buf.personality = 0;
}
// ---------------------------------------------------------------------------
// End: struct and constant definitions from <linux/io_uring.h>
// ---------------------------------------------------------------------------
#endif // __linux__

201
src/ailego/io/iouring_loader.cc Normal file
View File

@ -0,0 +1,201 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(__linux) || defined(__linux__)
#include <sys/syscall.h> // syscall(), __NR_io_uring_setup
#include <unistd.h> // close()
#include <cerrno>
#include <cstring>
#include <ailego/io/iouring_loader.h>
#include <zvec/ailego/logger/logger.h>
namespace zvec {
namespace core {
bool IoUringRing::setup(uint32_t entries) {
struct io_uring_params params;
std::memset(&params, 0, sizeof(params));
// io_uring_setup is a raw syscall — returns fd (>=0) or -1 with errno.
ring_fd_ = static_cast<int>(
syscall(__NR_io_uring_setup, static_cast<int>(entries), &params));
if (ring_fd_ < 0) {
// ENOSYS = kernel doesn't support io_uring.
// EPERM = io_uring disabled via sysctl.
// EINVAL = invalid parameters.
if (errno != ENOSYS) {
LOG_WARN("io_uring_setup failed; errno=%d, %s", errno, ::strerror(errno));
}
return false;
}
sq_entries_ = params.sq_entries;
cq_entries_ = params.cq_entries;
// The read path uses IORING_OP_READ (Linux 5.6+). On older kernels
// (5.1–5.5) io_uring_setup() succeeds but every read would fail with
// -EINVAL, so reject the ring here and let callers fall back.
if ((params.features & IORING_FEAT_RW_CUR_POS) == 0) {
LOG_WARN(
"io_uring lacks IORING_OP_READ support (kernel < 5.6); falling "
"back");
teardown();
return false;
}
// --- mmap the three shared regions ---
// 1. SQ ring (includes head, tail, mask, entries, flags, dropped, array).
size_t sq_ring_sz =
static_cast<size_t>(params.sq_off.array) + sq_entries_ * sizeof(uint32_t);
sq_ring_ptr_ = ::mmap(nullptr, sq_ring_sz, PROT_READ | PROT_WRITE, MAP_SHARED,
ring_fd_, IORING_OFF_SQ_RING);
if (sq_ring_ptr_ == MAP_FAILED) {
LOG_ERROR("mmap SQ ring failed: %s", ::strerror(errno));
sq_ring_ptr_ = nullptr;
teardown();
return false;
}
// 2. SQE array.
size_t sqes_sz = sq_entries_ * sizeof(struct io_uring_sqe);
sqes_ptr_ = reinterpret_cast<struct io_uring_sqe *>(
::mmap(nullptr, sqes_sz, PROT_READ | PROT_WRITE, MAP_SHARED, ring_fd_,
IORING_OFF_SQES));
if (sqes_ptr_ == MAP_FAILED) {
LOG_ERROR("mmap SQEs failed: %s", ::strerror(errno));
sqes_ptr_ = nullptr;
teardown();
return false;
}
// 3. CQ ring (includes head, tail, mask, entries, overflow, cqes[]).
size_t cq_ring_sz = static_cast<size_t>(params.cq_off.cqes) +
cq_entries_ * sizeof(struct io_uring_cqe);
cq_ring_ptr_ = ::mmap(nullptr, cq_ring_sz, PROT_READ | PROT_WRITE, MAP_SHARED,
ring_fd_, IORING_OFF_CQ_RING);
if (cq_ring_ptr_ == MAP_FAILED) {
LOG_ERROR("mmap CQ ring failed: %s", ::strerror(errno));
cq_ring_ptr_ = nullptr;
teardown();
return false;
}
// --- Set up typed pointers into the mmap'd regions ---
// SQ ring fields.
sq_head_ = reinterpret_cast<unsigned *>(static_cast<char *>(sq_ring_ptr_) +
params.sq_off.head);
sq_tail_ = reinterpret_cast<unsigned *>(static_cast<char *>(sq_ring_ptr_) +
params.sq_off.tail);
sq_ring_mask_ = reinterpret_cast<unsigned *>(
static_cast<char *>(sq_ring_ptr_) + params.sq_off.ring_mask);
sq_ring_entries_ = reinterpret_cast<unsigned *>(
static_cast<char *>(sq_ring_ptr_) + params.sq_off.ring_entries);
sq_flags_ = reinterpret_cast<unsigned *>(static_cast<char *>(sq_ring_ptr_) +
params.sq_off.flags);
sq_dropped_ = reinterpret_cast<unsigned *>(static_cast<char *>(sq_ring_ptr_) +
params.sq_off.dropped);
sq_array_ = reinterpret_cast<unsigned *>(static_cast<char *>(sq_ring_ptr_) +
params.sq_off.array);
// CQ ring fields.
cq_head_ = reinterpret_cast<unsigned *>(static_cast<char *>(cq_ring_ptr_) +
params.cq_off.head);
cq_tail_ = reinterpret_cast<unsigned *>(static_cast<char *>(cq_ring_ptr_) +
params.cq_off.tail);
cq_ring_mask_ = reinterpret_cast<unsigned *>(
static_cast<char *>(cq_ring_ptr_) + params.cq_off.ring_mask);
cq_ring_entries_ = reinterpret_cast<unsigned *>(
static_cast<char *>(cq_ring_ptr_) + params.cq_off.ring_entries);
cq_overflow_ = reinterpret_cast<unsigned *>(
static_cast<char *>(cq_ring_ptr_) + params.cq_off.overflow);
cqes_ = reinterpret_cast<struct io_uring_cqe *>(
static_cast<char *>(cq_ring_ptr_) + params.cq_off.cqes);
// SQE array.
sqes_ = sqes_ptr_;
// Initialize the SQ array to identity mapping so that logical index ==
// physical SQE index. This is the simplest and most common configuration.
for (uint32_t i = 0; i < sq_entries_; i++) {
sq_array_[i] = i;
}
return true;
}
void IoUringRing::teardown() {
if (sq_ring_ptr_ && sq_ring_ptr_ != MAP_FAILED) {
// We don't track the exact mmap size; munmap with a large enough size
// is safe because the kernel only unmaps what was actually mapped.
// However, to be correct we use the page-aligned size.
size_t sz = static_cast<size_t>(sq_entries_) * sizeof(uint32_t) + 4096;
::munmap(sq_ring_ptr_, sz);
}
if (sqes_ptr_ && sqes_ptr_ != MAP_FAILED) {
size_t sz = static_cast<size_t>(sq_entries_) * sizeof(struct io_uring_sqe);
::munmap(sqes_ptr_, sz);
}
if (cq_ring_ptr_ && cq_ring_ptr_ != MAP_FAILED) {
size_t sz =
static_cast<size_t>(cq_entries_) * sizeof(struct io_uring_cqe) + 4096;
::munmap(cq_ring_ptr_, sz);
}
sq_ring_ptr_ = nullptr;
sqes_ptr_ = nullptr;
cq_ring_ptr_ = nullptr;
sqes_ = nullptr;
cqes_ = nullptr;
sq_head_ = sq_tail_ = sq_ring_mask_ = sq_ring_entries_ = nullptr;
sq_flags_ = sq_dropped_ = sq_array_ = nullptr;
cq_head_ = cq_tail_ = cq_ring_mask_ = cq_ring_entries_ = nullptr;
cq_overflow_ = nullptr;
if (ring_fd_ >= 0) {
::close(ring_fd_);
ring_fd_ = -1;
}
::free(staging_);
staging_ = nullptr;
staging_size_ = 0;
sq_entries_ = 0;
cq_entries_ = 0;
}
bool IoUringRing::ensure_staging(size_t bytes) {
if (staging_size_ >= bytes) {
return true;
}
void *ptr = nullptr;
int err = ::posix_memalign(&ptr, kIoUringStagingAlign, bytes);
if (err != 0) {
LOG_WARN("io_uring staging allocation failed; size=%zu, %s", bytes,
::strerror(err));
return false;
}
::free(staging_);
staging_ = static_cast<char *>(ptr);
staging_size_ = bytes;
return true;
}
} // namespace core
} // namespace zvec
#endif // __linux__

129
src/ailego/io/iouring_loader.h Normal file
View File

@ -0,0 +1,129 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Raw-syscall wrapper for Linux io_uring: queue lifecycle via
// io_uring_setup/io_uring_enter + mmap, with zero dependency on liburing.
//
// setup() returns false when the kernel lacks the io_uring features we
// need (pre-5.6, missing IORING_OP_READ, or disabled), letting callers
// fall back to libaio or pread.
//
// Not thread-safe: each I/O thread must own its IoUringRing instance.
#pragma once
#if defined(__linux) || defined(__linux__)
#include <sys/mman.h>
#include <unistd.h>
#include <cstdlib>
#include <vector>
#include <ailego/io/iouring_def.h>
namespace zvec {
namespace core {
// AlignedRead lives in diskann_file_reader.h; a forward declaration
// suffices since execute() takes it by reference.
struct AlignedRead;
// Max SQEs submitted per io_uring_enter() call.
static constexpr uint32_t kIoUringMaxBatch = 128;
// Staging slot alignment — keeps O_DIRECT reads legal.
static constexpr size_t kIoUringStagingAlign = 4096;
class IoUringRing {
public:
IoUringRing() = default;
~IoUringRing() {
teardown();
}
IoUringRing(const IoUringRing &) = delete;
IoUringRing &operator=(const IoUringRing &) = delete;
// Create an io_uring with `entries` queue slots. Returns false if the
// kernel lacks io_uring or setup failed. In iouring_loader.cc.
bool setup(uint32_t entries);
// munmap all regions, close the ring fd, free staging. Closing the fd
// only *starts* asynchronous cancellation of in-flight requests, so call
// this only when the ring is quiesced — or after abandon_staging().
// In iouring_loader.cc.
void teardown();
// Grow the ring-owned staging pool to at least `bytes`. Only safe while
// the ring is quiesced (the old pool is freed here). In iouring_loader.cc.
bool ensure_staging(size_t bytes);
// Deliberately leak the staging pool when in-flight requests cannot be
// drained: the kernel may keep writing into it after the fd is closed.
// Caller buffers are never handed to the kernel and remain safe.
void abandon_staging() {
staging_ = nullptr;
staging_size_ = 0;
}
bool is_valid() const {
return ring_fd_ >= 0;
}
// Execute a batch of aligned reads via io_uring. Returns 0 on success,
// -1 on failure — the caller may always fall back to pread, since the
// kernel only writes into the staging pool. In diskann_file_reader.cc
// (AlignedRead is defined there).
int execute(int fd, std::vector<AlignedRead> &read_reqs);
private:
int ring_fd_{-1};
// mmap'd region bases (needed for munmap).
void *sq_ring_ptr_{nullptr};
struct io_uring_sqe *sqes_ptr_{nullptr};
void *cq_ring_ptr_{nullptr};
// SQ ring field pointers (into sq_ring_ptr_).
unsigned *sq_head_{nullptr};
unsigned *sq_tail_{nullptr};
unsigned *sq_ring_mask_{nullptr};
unsigned *sq_ring_entries_{nullptr};
unsigned *sq_flags_{nullptr};
unsigned *sq_dropped_{nullptr};
unsigned *sq_array_{nullptr};
// CQ ring field pointers (into cq_ring_ptr_).
unsigned *cq_head_{nullptr};
unsigned *cq_tail_{nullptr};
unsigned *cq_ring_mask_{nullptr};
unsigned *cq_ring_entries_{nullptr};
unsigned *cq_overflow_{nullptr};
struct io_uring_cqe *cqes_{nullptr};
struct io_uring_sqe *sqes_{nullptr};
// Ring-owned staging pool: the kernel reads into it and execute() copies
// verified completions out, decoupling caller buffer lifetime from async
// io_uring teardown (see abandon_staging()).
char *staging_{nullptr};
size_t staging_size_{0};
unsigned sq_entries_{0};
unsigned cq_entries_{0};
};
} // namespace core
} // namespace zvec
#endif // __linux__

1
src/binding/python/model/common/python_config.cc
View File

@ -228,6 +228,7 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) {
},
"Returns the current I/O backend type for DiskAnn async disk reads "
"as an IOBackendType enum (zvec.typing.IOBackendType). "
"IOBackendType.IO_URING if io_uring is available, "
"IOBackendType.LIBAIO if libaio is available, "
"IOBackendType.PREAD otherwise.");

4
src/binding/python/typing/python_type.cc
View File

@ -146,6 +146,7 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads.
- PREAD: Synchronous pread() \u2014 no async I/O.
- LIBAIO: libaio loaded at runtime via dlopen().
- IO_URING: io_uring via raw kernel syscalls (zero dependency).
Examples:
>>> from zvec.typing import IOBackendType
@ -153,7 +154,8 @@ Examples:
IOBackendType.LIBAIO
)pbdoc")
.value("PREAD", ailego::IOBackendType::kPread)
.value("LIBAIO", ailego::IOBackendType::kLibAio);
.value("LIBAIO", ailego::IOBackendType::kLibAio)
.value("IO_URING", ailego::IOBackendType::kIoUring);
}
void ZVecPyTyping::bind_status(py::module_ &m) {

336
src/core/algorithm/diskann/diskann_file_reader.cc
View File

@ -16,9 +16,11 @@
#include <algorithm>
#include <cassert>
#include <cerrno>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <thread>
#include <ailego/io/io_backend_def.h>
#include <zvec/ailego/io/io_backend.h>
#include <zvec/ailego/logger/logger.h>
@ -32,6 +34,10 @@ namespace core {
typedef struct io_event io_event_t;
typedef struct iocb iocb_t;
// Retry budget for draining in-flight io_uring requests when the kernel
// keeps returning EAGAIN/EBUSY (100 us sleep per retry, ~1 s total).
static constexpr size_t kIoUringDrainRetries = 10000;
// Ensures the I/O backend selection is logged exactly once per process,
// regardless of which entry point (setup_io_ctx or register_thread)
// triggers it first.
@ -58,11 +64,37 @@ int setup_io_ctx(IOContext &ctx) {
#if (defined(__linux) || defined(__linux__))
std::call_once(g_io_backend_log_once, log_diskann_io_backend);
if (ailego::IOBackend::Instance().is_pread()) {
// No async backend available — leave ctx null so callers fall back to
// synchronous pread().
return 0;
}
int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx);
return ret;
ctx = new IoBackend();
ailego::IOBackendType selected = ailego::IOBackend::Instance().available();
// Priority 1: io_uring (raw kernel syscalls — zero dependency).
if (selected == ailego::IOBackendType::kIoUring &&
ctx->ring.setup(MAX_EVENTS)) {
ctx->backend = IoBackend::IO_URING;
return 0;
}
// Priority 2: libaio (dlopen — soft dependency).
if (selected != ailego::IOBackendType::kPread &&
LibAioLoader::Instance().load() &&
LibAioLoader::Instance().is_available()) {
int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx->aio_ctx);
if (ret == 0) {
ctx->backend = IoBackend::LIBAIO;
return 0;
}
LOG_WARN("io_setup failed; returned: %d, %s. falling back to pread", ret,
::strerror(-ret));
}
// Priority 3: synchronous pread (always available).
ctx->backend = IoBackend::NONE;
return 0;
#else
return 0;
#endif
@ -70,15 +102,21 @@ int setup_io_ctx(IOContext &ctx) {
int destroy_io_ctx(IOContext &ctx) {
#if (defined(__linux) || defined(__linux__))
if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) {
if (ctx == nullptr) {
return 0;
}
int ret = LibAioLoader::Instance().io_destroy(ctx);
if (ret == 0) {
ctx = nullptr;
}
return ret;
if (ctx->backend == IoBackend::IO_URING) {
ctx->ring.teardown();
} else if (ctx->backend == IoBackend::LIBAIO &&
LibAioLoader::Instance().is_available()) {
LibAioLoader::Instance().io_destroy(ctx->aio_ctx);
}
// IoUringRing destructor also calls teardown() — idempotent and safe.
delete ctx;
ctx = nullptr;
return 0;
#else
return 0;
#endif
@ -107,7 +145,7 @@ static int execute_io_pread(int fd, std::vector<AlignedRead> &read_reqs) {
// invalid arguments. If that happens after submission, io_destroy() is the
// only safe way to quiesce the context before synchronous I/O touches the same
// destination buffers. Recreate the context so later reads can still use AIO.
static bool reset_aio_context(IOContext &ctx) {
static bool reset_aio_context(io_context_t &ctx) {
auto &loader = LibAioLoader::Instance();
int ret;
do {
@ -121,7 +159,7 @@ static bool reset_aio_context(IOContext &ctx) {
}
ctx = nullptr;
IOContext replacement = nullptr;
io_context_t replacement = nullptr;
ret = loader.io_setup(MAX_EVENTS, &replacement);
if (ret != 0) {
LOG_ERROR(
@ -134,7 +172,7 @@ static bool reset_aio_context(IOContext &ctx) {
return true;
}
int execute_io_libaio(IOContext &ctx, int fd,
int execute_io_libaio(io_context_t &ctx, int fd,
std::vector<AlignedRead> &read_reqs, uint64_t n_retries) {
uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS);
@ -254,18 +292,250 @@ int execute_io_libaio(IOContext &ctx, int fd,
}
#endif
int execute_io(IOContext &ctx, int fd, std::vector<AlignedRead> &read_reqs,
int execute_io(IOContext ctx, int fd, std::vector<AlignedRead> &read_reqs,
uint64_t n_retries = 0) {
#if (defined(__linux) || defined(__linux__))
if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) {
// Guard against null or sentinel contexts.
if (ctx == nullptr || ctx == (IOContext)-1) {
return execute_io_pread(fd, read_reqs);
}
return execute_io_libaio(ctx, fd, read_reqs, n_retries);
// Dispatch based on the active backend.
if (ctx->backend == IoBackend::IO_URING) {
int ret = ctx->ring.execute(fd, read_reqs);
if (ret == 0) {
return 0;
}
// The kernel only ever writes into the ring-owned staging pool, never
// into the caller's buffers, so a pread fallback can never race with
// requests that are still in flight.
LOG_WARN("io_uring execute failed; falling back to pread");
return execute_io_pread(fd, read_reqs);
}
if (ctx->backend == IoBackend::LIBAIO) {
return execute_io_libaio(ctx->aio_ctx, fd, read_reqs, n_retries);
}
// NONE backend — synchronous pread.
return execute_io_pread(fd, read_reqs);
#else
return execute_io_pread(fd, read_reqs);
#endif
}
// ---------------------------------------------------------------------------
// IoUringRing::execute — defined here (not in iouring_loader.h) because it
// accesses AlignedRead members, and AlignedRead is defined in
// diskann_file_reader.h after iouring_loader.h is included.
// ---------------------------------------------------------------------------
#if (defined(__linux) || defined(__linux__))
int IoUringRing::execute(int fd, std::vector<AlignedRead> &read_reqs) {
if (!is_valid()) {
return -1;
}
if (read_reqs.empty()) {
return 0;
}
// Process in batches limited by the SQ ring size.
uint32_t batch_size =
std::min(sq_entries_, static_cast<uint32_t>(kIoUringMaxBatch));
uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), batch_size);
for (uint64_t iter = 0; iter < iters; iter++) {
uint64_t n_ops =
std::min(static_cast<uint64_t>(read_reqs.size()) - iter * batch_size,
static_cast<uint64_t>(batch_size));
// --- Phase 1: Fill SQEs ---
//
// Reads land in the ring-owned staging pool, never in the caller's
// buffers. io_uring teardown is asynchronous — closing the ring fd
// only initiates cancellation — so the kernel may still write into
// request buffers after execute() has returned an error. Staging
// memory can simply be leaked in that case (abandon_staging()), while
// the caller's buffers stay safe to reuse or free. The copy-out below
// costs one sector-scale memcpy per read, negligible next to the I/O.
std::vector<size_t> slot_off(n_ops);
size_t staging_bytes = 0;
for (uint64_t j = 0; j < n_ops; j++) {
slot_off[j] = staging_bytes;
size_t len = read_reqs[j + iter * batch_size].len;
// Round every slot up so each staging pointer stays O_DIRECT-legal.
staging_bytes +=
(len + kIoUringStagingAlign - 1) & ~(kIoUringStagingAlign - 1);
}
// Safe: the previous batch is fully drained before we get here, so no
// in-flight request can reference the old pool being freed on growth.
if (!ensure_staging(staging_bytes)) {
return -1; // nothing submitted; pread fallback is safe
}
unsigned tail = __atomic_load_n(sq_tail_, __ATOMIC_ACQUIRE);
unsigned mask = *sq_ring_mask_;
for (uint64_t j = 0; j < n_ops; j++) {
unsigned idx = (tail + static_cast<unsigned>(j)) & mask;
unsigned sqe_idx = sq_array_[idx];
struct io_uring_sqe *sqe = &sqes_[sqe_idx];
uint64_t req_idx = j + iter * batch_size;
io_uring_prep_read(sqe, fd, staging_ + slot_off[j],
static_cast<uint32_t>(read_reqs[req_idx].len),
read_reqs[req_idx].offset);
// Store the request index so we can verify the completion.
sqe->user_data = req_idx;
}
// Memory barrier: ensure SQE contents are visible before tail update.
__sync_synchronize();
__atomic_store_n(sq_tail_, tail + static_cast<unsigned>(n_ops),
__ATOMIC_RELEASE);
// --- Phase 2: Submit and reap completions ---
//
// io_uring_enter() returns the number of SQEs consumed, not the number
// of CQEs available. A partial submission returns before the wait
// phase, and a signal can interrupt the wait while preserving a
// positive submission count, so IORING_ENTER_GETEVENTS guarantees
// min_complete completions only when the call finishes normally.
// Completions must therefore be counted against cq_tail instead of
// assuming n_ops CQEs are ready.
uint64_t submitted = 0;
uint64_t completed = 0;
bool all_ok = true;
// Consume every CQE the kernel has published so far and verify it.
// Completion order is unspecified, so use cqe->user_data to find the
// request instead of assuming submission order.
auto reap_available = [&]() {
unsigned chead = *cq_head_; // single consumer — plain load is enough
unsigned ctail = __atomic_load_n(cq_tail_, __ATOMIC_ACQUIRE);
unsigned cq_mask = *cq_ring_mask_;
if (chead == ctail) {
return;
}
while (chead != ctail) {
struct io_uring_cqe *cqe = &cqes_[chead & cq_mask];
uint64_t req_idx = cqe->user_data;
if (req_idx < iter * batch_size ||
req_idx >= iter * batch_size + n_ops) {
LOG_WARN("io_uring completion referenced unknown request: %lu",
(unsigned long)req_idx);
all_ok = false;
} else if (cqe->res < 0) {
LOG_WARN("io_uring read failed: req=%lu, res=%d, offset=%lu",
(unsigned long)req_idx, cqe->res,
(unsigned long)read_reqs[req_idx].offset);
all_ok = false;
} else if (static_cast<uint64_t>(cqe->res) != read_reqs[req_idx].len) {
LOG_WARN("io_uring short read: req=%lu, got=%d, expected=%lu",
(unsigned long)req_idx, cqe->res,
(unsigned long)read_reqs[req_idx].len);
all_ok = false;
} else {
// Verified completion — copy from staging into the caller's
// buffer. This is the only place caller memory is written.
std::memcpy(read_reqs[req_idx].buf,
staging_ + slot_off[req_idx - iter * batch_size],
read_reqs[req_idx].len);
}
chead++;
completed++;
}
// Release: CQE reads must complete before the kernel may reuse slots.
__atomic_store_n(cq_head_, chead, __ATOMIC_RELEASE);
};
while (completed < n_ops) {
reap_available();
if (completed >= n_ops) {
break;
}
unsigned to_submit = static_cast<unsigned>(n_ops - submitted);
int ret = static_cast<int>(syscall(
__NR_io_uring_enter, ring_fd_, to_submit, 1u, IORING_ENTER_GETEVENTS,
static_cast<void *>(nullptr), static_cast<size_t>(0)));
if (ret >= 0) {
submitted += static_cast<uint64_t>(ret);
continue;
}
if (errno == EINTR) {
// Interrupted during submit or wait; the SQEs already consumed are
// tracked in `submitted`, so simply retry.
continue;
}
if ((errno == EAGAIN || errno == EBUSY) && completed < submitted) {
// Kernel resources are exhausted, but in-flight requests will free
// them as they complete; keep reaping and retrying.
continue;
}
// Unrecoverable failure (or EAGAIN with nothing in flight).
LOG_WARN(
"io_uring_enter failed; errno=%d, %s, submitted=%lu/%lu, "
"completed=%lu. draining before falling back to pread",
errno, ::strerror(errno), (unsigned long)submitted,
(unsigned long)n_ops, (unsigned long)completed);
// Un-publish the SQEs the kernel never consumed so a later batch
// cannot submit them against stale buffers.
__atomic_store_n(sq_tail_, tail + static_cast<unsigned>(submitted),
__ATOMIC_RELEASE);
// Drain every in-flight request before the staging pool may be
// freed or reused by a later batch. CQEs are posted to the shared
// ring by the kernel on its own, so completions can still be reaped
// here even when io_uring_enter() keeps failing.
size_t drain_retries = 0;
while (completed < submitted) {
reap_available();
if (completed >= submitted) {
break;
}
int wret = static_cast<int>(syscall(
__NR_io_uring_enter, ring_fd_, 0u, 1u, IORING_ENTER_GETEVENTS,
static_cast<void *>(nullptr), static_cast<size_t>(0)));
if (wret >= 0 || errno == EINTR) {
continue;
}
if ((errno == EAGAIN || errno == EBUSY) &&
drain_retries++ < kIoUringDrainRetries) {
// Give in-flight requests time to complete; entering the kernel
// via the sleep also lets pending completion task-work run.
std::this_thread::sleep_for(std::chrono::microseconds(100));
continue;
}
// The ring cannot be drained. Leak the staging pool — the kernel
// may keep writing into it through the asynchronous teardown — and
// disable io_uring for this context. The caller's buffers were
// never exposed to the kernel, so the pread fallback stays safe.
LOG_ERROR(
"io_uring drain failed; errno=%d, %s. leaking the staging pool "
"and disabling io_uring for this context",
errno, ::strerror(errno));
abandon_staging();
teardown();
return -1;
}
return -1;
}
if (!all_ok) {
// Every request completed and the staging pool is quiesced, but at
// least one read failed or was short — let the caller retry with
// pread.
return -1;
}
}
return 0;
}
#endif // __linux__
LinuxAlignedFileReader::LinuxAlignedFileReader(int file_desc) {
this->file_desc = file_desc;
}
@ -286,7 +556,7 @@ IOContext &LinuxAlignedFileReader::get_ctx() {
std::unique_lock<std::mutex> lk(ctx_mut);
auto it = ctx_map.find(std::this_thread::get_id());
if (it == ctx_map.end()) {
LOG_ERROR("bad thread access; returning -1 as io_context_t");
LOG_ERROR("bad thread access; returning invalid IOContext");
return this->bad_ctx;
} else {
return it->second;
@ -299,32 +569,20 @@ void LinuxAlignedFileReader::register_thread() {
std::unique_lock<std::mutex> lk(ctx_mut);
if (ctx_map.find(thread_id) != ctx_map.end()) {
LOG_ERROR("multiple calls to register_thread from the same thread");
return;
}
IOContext ctx = nullptr;
std::call_once(g_io_backend_log_once, log_diskann_io_backend);
if (ailego::IOBackend::Instance().is_pread()) {
int ret = setup_io_ctx(ctx);
if (ret != 0) {
LOG_ERROR("setup_io_ctx failed; returned: %d", ret);
lk.unlock();
return;
}
int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx);
if (ret != 0) {
if (ret == -EAGAIN) {
LOG_ERROR(
"io_setup failed with EAGAIN: Consider increasing "
"/proc/sys/fs/aio-max-nr");
} else {
LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret));
}
} else {
if (ctx != nullptr) {
LOG_INFO("allocating ctx: %lu", (uint64_t)ctx);
ctx_map[thread_id] = ctx;
}
ctx_map[thread_id] = ctx;
lk.unlock();
#endif
}
@ -345,11 +603,8 @@ void LinuxAlignedFileReader::deregister_thread() {
ctx_map.erase(it);
}
// io_destroy is a syscall; keep it outside the lock to avoid blocking others
if (ailego::IOBackend::Instance().available() !=
ailego::IOBackendType::kPread) {
LibAioLoader::Instance().io_destroy(ctx);
}
// Teardown is a syscall; keep it outside the lock to avoid blocking others.
destroy_io_ctx(ctx);
LOG_INFO("returned ctx from thread");
#endif
}
@ -357,13 +612,8 @@ void LinuxAlignedFileReader::deregister_thread() {
void LinuxAlignedFileReader::deregister_all_threads() {
#if (defined(__linux) || defined(__linux__))
std::unique_lock<std::mutex> lk(ctx_mut);
bool aio_available = ailego::IOBackend::Instance().available() !=
ailego::IOBackendType::kPread;
for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) {
IOContext ctx = x->second;
if (aio_available) {
LibAioLoader::Instance().io_destroy(ctx);
}
destroy_io_ctx(x->second);
}
ctx_map.clear();
#endif

27
src/core/algorithm/diskann/diskann_file_reader.h
View File

@ -18,6 +18,7 @@
#include <fcntl.h>
#if (defined(__linux) || defined(__linux__))
#include <ailego/io/iouring_loader.h> // raw-syscall io_uring wrapper (IoUringRing)
#include <ailego/io/libaio_loader.h> // dlopen-based libaio wrapper
#endif
@ -31,7 +32,31 @@ namespace zvec {
namespace core {
#if (defined(__linux) || defined(__linux__))
typedef io_context_t IOContext;
// IoBackend holds the per-thread I/O context for whichever async backend
// was successfully initialised at setup time. The priority is:
// 1. io_uring (raw kernel syscalls — zero dependency)
// 2. libaio (dlopen — soft dependency)
// 3. pread (always available — synchronous fallback)
//
// IOContext is a *pointer* to IoBackend, which preserves the existing
// sentinel conventions: nullptr means uninitialised and (IOContext)-1 is
// the invalid-handle sentinel returned by get_ctx() for unregistered
// threads.
struct IoBackend {
enum Backend : uint8_t {
NONE = 0, // synchronous pread
IO_URING = 1, // io_uring via raw syscalls
LIBAIO = 2, // libaio via dlopen
};
Backend backend{NONE};
IoUringRing ring{};
io_context_t aio_ctx{nullptr};
};
typedef IoBackend *IOContext;
#else
typedef uint32_t IOContext;
#endif

10
src/include/zvec/ailego/io/io_backend.h
View File

@ -29,13 +29,17 @@ namespace zvec {
namespace ailego {
// Supported I/O backend types.
//
// Numeric values are part of the C ABI (see zvec_io_backend_type_t in c_api.h):
// kPread = 0, kLibAio = 1, kIoUring = 2.
enum class IOBackendType {
kPread, // Synchronous pread() — no async I/O
kLibAio, // libaio loaded at runtime via dlopen()
kPread = 0, // Synchronous pread() — no async I/O
kLibAio = 1, // libaio loaded at runtime via dlopen()
kIoUring = 2, // io_uring via raw kernel syscalls (zero dependency)
};
// Returns the currently active I/O backend type.
// Triggers backend initialization on first call (libaio > pread).
// Triggers backend initialization on first call (io_uring > libaio > pread).
IOBackendType current_io_backend_type();
// Returns a human-readable description of the currently active I/O backend.

7
src/include/zvec/c_api.h
View File

@ -789,6 +789,8 @@ typedef uint32_t zvec_io_backend_type_t;
0 /**< Synchronous pread() \u2014 no async I/O */
#define ZVEC_IO_BACKEND_TYPE_LIBAIO \
1 /**< libaio loaded at runtime via dlopen() */
#define ZVEC_IO_BACKEND_TYPE_IO_URING \
2 /**< io_uring via raw kernel syscalls (zero dependency) */
/**
* @brief Get the current I/O backend type for DiskAnn async disk reads.
@ -796,7 +798,8 @@ typedef uint32_t zvec_io_backend_type_t;
* Pure introspection \u2014 no side effects, no install hints.
*
* @return zvec_io_backend_type_t The loaded backend type
* (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD).
* (ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO,
* or ZVEC_IO_BACKEND_TYPE_PREAD).
*/
ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void);
@ -805,7 +808,7 @@ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void);
*
* @param type The backend type code.
* @return Thread-local string valid until the next call on this thread;
* "libaio", "pread", or "unknown".
* "io_uring", "libaio", "pread", or "unknown".
*/
ZVEC_EXPORT const char *ZVEC_CALL
zvec_get_io_backend_type_name(zvec_io_backend_type_t type);

8
tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc
View File

@ -30,7 +30,7 @@
namespace zvec {
namespace core {
int execute_io_libaio(IOContext &ctx, int fd,
int execute_io_libaio(io_context_t &ctx, int fd,
std::vector<AlignedRead> &read_reqs,
uint64_t n_retries = 0);
} // namespace core
@ -192,7 +192,7 @@ TEST(DiskAnnLinuxAioTest, AccumulatesPartialSubmissionsAndCompletions) {
state.submit_results = {2, 2};
state.completion_results = {1, 2, 1};
FakeAioGuard guard(&state);
IOContext ctx = reinterpret_cast<IOContext>(static_cast<uintptr_t>(1));
io_context_t ctx = reinterpret_cast<io_context_t>(static_cast<uintptr_t>(1));
// An invalid fd makes any accidental pread fallback fail the test.
EXPECT_EQ(execute_io_libaio(ctx, -1, requests), 0);
@ -221,7 +221,7 @@ TEST(DiskAnnLinuxAioTest, DrainsPartialSubmissionBeforePreadFallback) {
state.submit_results = {2, -EAGAIN};
state.completion_results = {1, 1};
FakeAioGuard guard(&state);
IOContext ctx = reinterpret_cast<IOContext>(static_cast<uintptr_t>(1));
io_context_t ctx = reinterpret_cast<io_context_t>(static_cast<uintptr_t>(1));
EXPECT_EQ(execute_io_libaio(ctx, file.fd(), requests), 0);
EXPECT_EQ(state.submit_sizes, (std::vector<long>{4, 2}));
@ -248,7 +248,7 @@ TEST(DiskAnnLinuxAioTest, DrainsAllCompletionsBeforePreadFallback) {
state.completion_results = {1, 1, 2};
state.short_completion = 0;
FakeAioGuard guard(&state);
IOContext ctx = reinterpret_cast<IOContext>(static_cast<uintptr_t>(1));
io_context_t ctx = reinterpret_cast<io_context_t>(static_cast<uintptr_t>(1));
EXPECT_EQ(execute_io_libaio(ctx, file.fd(), requests), 0);
EXPECT_EQ(state.completion_sizes, (std::vector<long>{4, 3, 2}));