From 85f89dce2cbd4823b75cba59731e487fe7b3b744 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Wed, 20 May 2026 16:59:09 +0800 Subject: [PATCH 01/46] fix --- src/core/algorithm/flat/CMakeLists.txt | 7 +++++++ src/core/algorithm/flat_sparse/CMakeLists.txt | 9 +++++++++ src/core/algorithm/hnsw/CMakeLists.txt | 6 ++++++ src/core/algorithm/hnsw_rabitq/CMakeLists.txt | 6 ++++++ src/core/algorithm/hnsw_sparse/CMakeLists.txt | 6 ++++++ src/core/algorithm/ivf/CMakeLists.txt | 6 ++++++ src/core/algorithm/vamana/CMakeLists.txt | 6 ++++++ src/core/metric/CMakeLists.txt | 6 ++++++ src/core/mixed_reducer/CMakeLists.txt | 6 ++++++ src/core/quantizer/CMakeLists.txt | 6 ++++++ src/core/utility/CMakeLists.txt | 6 ++++++ 11 files changed, 70 insertions(+) diff --git a/src/core/algorithm/flat/CMakeLists.txt b/src/core/algorithm/flat/CMakeLists.txt index 4564d8ef0..60814960e 100644 --- a/src/core/algorithm/flat/CMakeLists.txt +++ b/src/core/algorithm/flat/CMakeLists.txt @@ -1,11 +1,18 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) #message(STATUS "PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}") + +if(NOT APPLE) + set(CORE_KNN_FLAT_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_flat STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ${PROJECT_ROOT_DIR}/src/core/framework + LDFLAGS "${CORE_KNN_FLAT_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/algorithm/flat_sparse/CMakeLists.txt b/src/core/algorithm/flat_sparse/CMakeLists.txt index e27d2d3ee..44766138d 100644 --- a/src/core/algorithm/flat_sparse/CMakeLists.txt +++ b/src/core/algorithm/flat_sparse/CMakeLists.txt @@ -1,11 +1,20 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +# --exclude-libs is GNU ld / LLVM lld only; Apple ld does not support it. +# On macOS (Mach-O), symbol interposition works differently and the +# Arrow/Parquet double-free issue does not apply. +if(NOT APPLE) + set(CORE_KNN_FLAT_SPARSE_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_flat_sparse STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_FLAT_SPARSE_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/algorithm/hnsw/CMakeLists.txt b/src/core/algorithm/hnsw/CMakeLists.txt index f4a105402..cfd1147f4 100644 --- a/src/core/algorithm/hnsw/CMakeLists.txt +++ b/src/core/algorithm/hnsw/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_KNN_HNSW_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_hnsw STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework sparsehash INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/algorithm/hnsw_rabitq/CMakeLists.txt b/src/core/algorithm/hnsw_rabitq/CMakeLists.txt index ed547dc76..09ce72f55 100644 --- a/src/core/algorithm/hnsw_rabitq/CMakeLists.txt +++ b/src/core/algorithm/hnsw_rabitq/CMakeLists.txt @@ -11,11 +11,17 @@ if(AUTO_DETECT_ARCH) endforeach() endif() +if(NOT APPLE) + set(CORE_KNN_HNSW_RABITQ_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_hnsw_rabitq STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework rabitqlib sparsehash INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_RABITQ_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) \ No newline at end of file diff --git a/src/core/algorithm/hnsw_sparse/CMakeLists.txt b/src/core/algorithm/hnsw_sparse/CMakeLists.txt index fe26d10e1..15295b485 100644 --- a/src/core/algorithm/hnsw_sparse/CMakeLists.txt +++ b/src/core/algorithm/hnsw_sparse/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_KNN_HNSW_SPARSE_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_hnsw_sparse STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework sparsehash INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_HNSW_SPARSE_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/algorithm/ivf/CMakeLists.txt b/src/core/algorithm/ivf/CMakeLists.txt index ffcf30949..8e3872f31 100644 --- a/src/core/algorithm/ivf/CMakeLists.txt +++ b/src/core/algorithm/ivf/CMakeLists.txt @@ -1,10 +1,16 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_KNN_IVF_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_ivf STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS zvec_ailego core_framework core_knn_cluster INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_IVF_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/algorithm/vamana/CMakeLists.txt b/src/core/algorithm/vamana/CMakeLists.txt index 8e5bbda1e..b2feaf9c1 100644 --- a/src/core/algorithm/vamana/CMakeLists.txt +++ b/src/core/algorithm/vamana/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_KNN_VAMANA_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_vamana STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS core_framework core_knn_hnsw sparsehash INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_VAMANA_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/metric/CMakeLists.txt b/src/core/metric/CMakeLists.txt index 55dfc901e..2918b909b 100644 --- a/src/core/metric/CMakeLists.txt +++ b/src/core/metric/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_METRIC_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_metric STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS zvec_ailego zvec_turbo core_framework INCS . ${PROJECT_ROOT_DIR}/src/core + LDFLAGS "${CORE_METRIC_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/mixed_reducer/CMakeLists.txt b/src/core/mixed_reducer/CMakeLists.txt index e9566456e..e7204f0f7 100644 --- a/src/core/mixed_reducer/CMakeLists.txt +++ b/src/core/mixed_reducer/CMakeLists.txt @@ -1,10 +1,16 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_MIX_REDUCER_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_mix_reducer STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS zvec_ailego core_framework INCS . ${PROJECT_ROOT_DIR}/src/core + LDFLAGS "${CORE_MIX_REDUCER_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/quantizer/CMakeLists.txt b/src/core/quantizer/CMakeLists.txt index 21a03e449..80b4f612a 100644 --- a/src/core/quantizer/CMakeLists.txt +++ b/src/core/quantizer/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_QUANTIZER_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_quantizer STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS zvec_ailego core_framework INCS . ${PROJECT_ROOT_DIR}/src/core + LDFLAGS "${CORE_QUANTIZER_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/utility/CMakeLists.txt b/src/core/utility/CMakeLists.txt index 99cf87ca2..7c3adf702 100644 --- a/src/core/utility/CMakeLists.txt +++ b/src/core/utility/CMakeLists.txt @@ -1,11 +1,17 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +if(NOT APPLE) + set(CORE_UTILITY_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_utility STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS zvec_ailego core_framework INCS . ${PROJECT_ROOT_DIR}/src/core + LDFLAGS "${CORE_UTILITY_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) From 74da77aac08965ae47aa50f472cd3e74e1aac5c2 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 24 Jun 2026 21:19:40 +0800 Subject: [PATCH 02/46] refactor: use dlopen --- src/core/algorithm/diskann/CMakeLists.txt | 9 +++-- .../algorithm/diskann/diskann_file_reader.cc | 36 ++++++++++++++----- .../algorithm/diskann/diskann_file_reader.h | 2 +- 3 files changed, 33 insertions(+), 14 deletions(-) diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 1b2252497..e2fa911f5 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -13,14 +13,13 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # before invoking dlopen's RTLD_GLOBAL symbol sharing, and the Python wheel # (which only ships _zvec.so + the plugin) would fail to load. # -# We therefore link the plugin only against its system-level dependency -# (libaio) and rely on the global include dirs plus -Wl,--unresolved-symbols -# =ignore-all to let the linker build a shared library with unresolved -# references to internal APIs. +# libaio is loaded at runtime via dlopen()/dlsym() (see libaio_loader.h), so +# we do NOT link against -laio. The only system-level link dependency is +# ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. set(CORE_KNN_DISKANN_LIBS "") if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") - list(APPEND CORE_KNN_DISKANN_LIBS aio) + list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() cc_library( diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 4268d5156..34584d935 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -32,7 +32,10 @@ typedef struct iocb iocb_t; int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - int ret = io_setup(MAX_EVENTS, &ctx); + if (!LibAioLoader::Instance().Load()) { + return -ENOSYS; + } + int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); return ret; #else @@ -42,7 +45,10 @@ int setup_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - int ret = io_destroy(ctx); + if (!LibAioLoader::Instance().IsAvailable()) { + return 0; + } + int ret = LibAioLoader::Instance().io_destroy(ctx); return ret; #else @@ -71,6 +77,9 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { int execute_io(IOContext ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) + if (!LibAioLoader::Instance().Load()) { + return execute_io_pread(fd, read_reqs); + } uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); for (uint64_t iter = 0; iter < iters; iter++) { @@ -93,7 +102,8 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, size_t n_tries = 0; // Phase 1: io_submit with retry. while (true) { - int ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); + int ret = + LibAioLoader::Instance().io_submit(ctx, (int64_t)n_ops, cbs.data()); if (ret == (int)n_ops) { break; } @@ -111,8 +121,8 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, // Phase 2: io_getevents with retry (never re-submits). n_tries = 0; while (true) { - int ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), - nullptr); + int ret = LibAioLoader::Instance().io_getevents( + ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), nullptr); if (ret == (int)n_ops) { break; } @@ -188,7 +198,12 @@ void LinuxAlignedFileReader::register_thread() { IOContext ctx = nullptr; - int ret = io_setup(MAX_EVENTS, &ctx); + if (!LibAioLoader::Instance().Load()) { + LOG_ERROR("libaio not available; cannot register thread for async I/O"); + lk.unlock(); + return; + } + int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); if (ret != 0) { lk.unlock(); if (ret == -EAGAIN) { @@ -226,7 +241,9 @@ void LinuxAlignedFileReader::deregister_thread() { } // io_destroy is a syscall; keep it outside the lock to avoid blocking others - io_destroy(ctx); + if (LibAioLoader::Instance().IsAvailable()) { + LibAioLoader::Instance().io_destroy(ctx); + } LOG_INFO("returned ctx from thread"); #endif } @@ -234,9 +251,12 @@ void LinuxAlignedFileReader::deregister_thread() { void LinuxAlignedFileReader::deregister_all_threads() { #if (defined(__linux) || defined(__linux__)) std::unique_lock lk(ctx_mut); + bool aio_available = LibAioLoader::Instance().IsAvailable(); for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; - io_destroy(ctx); + if (aio_available) { + LibAioLoader::Instance().io_destroy(ctx); + } } ctx_map.clear(); #endif diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 432247e79..2abe8c1aa 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -18,7 +18,7 @@ #include #if (defined(__linux) || defined(__linux__)) -#include +#include "libaio_loader.h" // dlopen-based libaio wrapper #endif #include From 17a7fd50f788450072300216095c202128fd2069 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 25 Jun 2026 16:16:16 +0800 Subject: [PATCH 03/46] refactor: use dlopen to load libaio --- CMakeLists.txt | 13 +- python/tests/detail/fixture_helper.py | 36 ++- python/tests/test_collection_diskann.py | 12 +- python/zvec/__init__.py | 2 +- src/binding/python/CMakeLists.txt | 36 +-- src/binding/python/binding.cc | 16 +- src/binding/python/exports.map | 19 +- src/core/CMakeLists.txt | 7 +- src/core/algorithm/diskann/CMakeLists.txt | 50 +-- .../algorithm/diskann/diskann_file_reader.cc | 5 +- src/core/algorithm/diskann/libaio_def.h | 190 ++++++++++++ src/core/algorithm/diskann/libaio_loader.h | 144 +++++++++ src/core/interface/indexes/diskann_index.cc | 53 +--- src/core/plugin/CMakeLists.txt | 6 +- src/core/plugin/diskann_plugin.cc | 289 +++--------------- src/db/index/common/schema.cc | 19 +- src/include/zvec/ailego/pattern/factory.h | 21 +- src/include/zvec/plugin/diskann_plugin.h | 34 +-- .../algorithm/diskann/diskann_builder_test.cc | 8 +- 19 files changed, 508 insertions(+), 452 deletions(-) create mode 100644 src/core/algorithm/diskann/libaio_def.h create mode 100644 src/core/algorithm/diskann/libaio_loader.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 392dbdac8..fe8e2e89a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -190,17 +190,8 @@ if(BUILD_PYTHON_BINDINGS) message(STATUS "Zvec install path: ${ZVEC_PY_INSTALL_DIR}") install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}) - # DiskAnn ships as a runtime-loaded shared module - # (libzvec_diskann_plugin.so) that is brought online implicitly the - # first time a DiskAnn index is created — users never call any load - # function. The Python extension resolves the module next to _zvec.so - # (see the $ORIGIN rpath in src/binding/python/CMakeLists.txt); the - # module must therefore be installed alongside _zvec.so in the wheel. - # The target exists only on platforms where DiskAnn is buildable - # (currently Linux x86_64 with libaio). - if(TARGET core_knn_diskann) - install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}) - endif() + # DiskAnn is now statically linked into _zvec.so via --whole-archive, + # so no separate runtime .so needs to be installed in the wheel. # Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works # out of the box. python/zvec/__init__.py resolves this directory via # importlib.resources and registers it with set_default_jieba_dict_dir(). diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index 63aac799c..2054546cf 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -14,18 +14,21 @@ import zvec -# Cache the DiskAnn plugin preload status so we pay the load cost once per -# test session. The plugin normally auto-loads on first DiskAnn use, but we -# preload it explicitly here so a missing libaio / misplaced plugin .so -# surfaces as a clear pytest skip instead of a confusing -# "Create vector column indexer failed" deep inside the collection code path. +# Cache the DiskAnn runtime preload status so we pay the init cost once per +# test session. The runtime normally auto-loads on first DiskAnn use, but we +# preload it explicitly here so a missing libaio surfaces as a clear pytest +# skip instead of a confusing "Create vector column indexer failed" deep +# inside the collection code path. +# Note: DiskAnn is now compiled directly into _zvec.so — there is no separate +# plugin .so to locate. If libaio is missing, DiskAnn falls back to +# synchronous pread() with degraded performance. _DISKANN_PRELOAD_REASON: str | None = None _DISKANN_PRELOAD_DONE: bool = False def _ensure_diskann_runtime_or_reason() -> str | None: - """Preload the DiskAnn plugin and return None on success or a human-readable - skip reason on failure. Idempotent across calls.""" + """Initialize the DiskAnn runtime and return None on success or a + human-readable skip reason on failure. Idempotent across calls.""" global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON if _DISKANN_PRELOAD_DONE: return _DISKANN_PRELOAD_REASON @@ -35,19 +38,20 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" return _DISKANN_PRELOAD_REASON - if not zvec.is_libaio_available(): + status = zvec.load_diskann_plugin() + if status == zvec.DISKANN_PLUGIN_UNSUPPORTED_PLATFORM: _DISKANN_PRELOAD_REASON = ( - "libaio is not available on this host; DiskAnn cannot run. " - "Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry." + f"DiskAnn is not supported on this platform (status={status})." ) return _DISKANN_PRELOAD_REASON - - status = zvec.load_diskann_plugin() - if status != zvec.DISKANN_PLUGIN_OK: + # DISKANN_PLUGIN_OK and DISKANN_PLUGIN_LIBAIO_MISSING are both acceptable: + # the latter means DiskAnn will use synchronous pread() instead of async I/O. + if status not in ( + zvec.DISKANN_PLUGIN_OK, + zvec.DISKANN_PLUGIN_LIBAIO_MISSING, + ): _DISKANN_PRELOAD_REASON = ( - f"Failed to load DiskAnn plugin (status={status}); " - "check that libzvec_diskann_plugin.so is installed alongside " - "_zvec.so in the Python site-packages directory." + f"Failed to initialize DiskAnn runtime (status={status})." ) return _DISKANN_PRELOAD_REASON diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index b0e12ce7e..92cbc74e4 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -13,17 +13,17 @@ # limitations under the License. """End-to-end collection tests for the DiskAnn index. -Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin. +Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index. Two platform-level prerequisites are enforced at module import time: 1. DiskAnn is currently built only for Linux x86_64 — other platforms are skipped wholesale. -2. The DiskAnn backend lives in a *runtime-loaded* plugin - (``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL | - RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory`` - singleton is unified with the one inside ``_zvec.so``. After ``import - zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once. +2. The DiskAnn runtime is initialized via ``zvec.load_diskann_plugin()``, + which eagerly loads libaio via dlopen(). If libaio is missing, DiskAnn + falls back to synchronous pread() — the tests still run but with degraded + performance. DiskAnn is compiled directly into ``_zvec.so``; there is no + separate plugin .so to locate. If either prerequisite fails the whole module is skipped so the rest of the test-suite is not affected. diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index 5fdf9732c..a78fcdd75 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -201,7 +201,7 @@ "StatusCode", # Tools "require_module", - # DiskAnn plugin + # DiskAnn runtime "load_diskann_plugin", "is_diskann_plugin_loaded", "is_libaio_available", diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index 0db6d75ff..4cc53b6da 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -22,13 +22,12 @@ pybind11_add_module(_zvec ${SRC_LISTS}) # pybind11_add_module() defaults to CXX_VISIBILITY_PRESET=hidden + # VISIBILITY_INLINES_HIDDEN=ON, which hides the compiler-generated helper # symbols attached to inline functions (guard variables for static locals, -# vtables, typeinfo ...). The DiskAnn runtime plugin has its own copy of -# Factory::Instance()'s guard variable; if _zvec.so's copy -# is hidden, the two guards are separate and the factory constructor runs -# twice during plugin load, wiping out the registrations that happened -# during _zvec.so import. We switch to default visibility here and rely on -# the version script (exports.map / exports.mac) to keep the dynamic -# symbol table small by exporting only zvec::* and PyInit_*. +# vtables, typeinfo ...). DiskAnn is now whole-archived directly into +# _zvec.so (no separate runtime .so), so guard-variable duplication across +# DSOs is no longer a concern. We still keep default visibility so that +# the version script (exports.map / exports.mac) — not the compiler — is +# the sole gatekeeper of the dynamic symbol table, exporting only zvec::* +# and PyInit_*. set_target_properties(_zvec PROPERTIES CXX_VISIBILITY_PRESET default C_VISIBILITY_PRESET default @@ -57,6 +56,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") $ $ $ + $ $ $ $ @@ -69,21 +69,11 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") target_link_options(_zvec PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map" ) - # DiskAnn is x86-only (it depends on libaio, unavailable on ARM64), so the - # runtime plugin (libzvec_diskann_plugin.so) is only produced on non-ARM - # builds. It is shipped as a runtime-loaded shared module and brought up - # implicitly the first time a DiskAnn index is created — users never need to - # call any load function. If libaio is missing at runtime the auto-load - # fails cleanly and the error is surfaced only when DiskAnn is actually - # used; other index types (HNSW/IVF/Flat/Vamana) remain fully functional. - # The .so must therefore be discoverable next to the extension module, - # hence the $ORIGIN rpath below. - if (NOT CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") - set_target_properties(_zvec PROPERTIES - BUILD_RPATH "$ORIGIN" - INSTALL_RPATH "$ORIGIN" - ) - endif() + # DiskAnn is now statically linked into _zvec.so via --whole-archive, so + # no separate runtime .so is shipped in the wheel. libaio is still loaded + # at runtime via dlopen() (see libaio_loader.h); if it is missing, DiskAnn + # fails cleanly with an actionable error while other index types + # (HNSW/IVF/Flat/Vamana) remain fully functional. elseif (APPLE) target_link_libraries(_zvec PRIVATE -Wl,-force_load,$ @@ -94,6 +84,7 @@ elseif (APPLE) -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ + -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ @@ -113,6 +104,7 @@ elseif (MSVC) core_knn_ivf_static core_knn_vamana_static core_knn_cluster_static + core_knn_diskann_static core_mix_reducer_static core_metric_static core_utility_static diff --git a/src/binding/python/binding.cc b/src/binding/python/binding.cc index ecb4d91bb..4e257472f 100644 --- a/src/binding/python/binding.cc +++ b/src/binding/python/binding.cc @@ -26,18 +26,20 @@ namespace zvec { namespace { -// Expose DiskAnn plugin management to Python. The DiskAnn runtime normally -// auto-loads on first use, but tests (and diagnostic tooling) need a way to -// force a load up-front and get actionable errors when libaio is missing or -// the plugin shared object cannot be located. +// Expose DiskAnn runtime management to Python. DiskAnn is compiled directly +// into _zvec.so, so "loading" just means eagerly dlopen()-ing libaio and +// caching the result. Tests (and diagnostic tooling) use these entry points +// to force the load up-front and get actionable warnings when libaio is +// missing. void InitializeDiskAnnPluginBindings(pybind11::module_ &m) { m.def( "load_diskann_plugin", [](const std::string &path) { return ::zvec::LoadDiskAnnPlugin(path); }, pybind11::arg("path") = std::string(), - "Load the DiskAnn runtime plugin. Returns 0 on success or a negative " - "DiskAnnPluginStatus code on failure (unsupported platform, libaio " - "missing, or dlopen failure)."); + "Load libaio for the DiskAnn runtime. Returns DISKANN_PLUGIN_OK (0) " + "on success, or DISKANN_PLUGIN_LIBAIO_MISSING if libaio is not " + "available (DiskAnn falls back to synchronous pread in that case). " + "Returns a negative code for unsupported platforms."); m.def("is_diskann_plugin_loaded", &::zvec::IsDiskAnnPluginLoaded, "Return True if the DiskAnn runtime plugin is currently loaded."); m.def("is_libaio_available", &::zvec::IsLibAioAvailable, diff --git a/src/binding/python/exports.map b/src/binding/python/exports.map index 553da0055..76cc27953 100644 --- a/src/binding/python/exports.map +++ b/src/binding/python/exports.map @@ -3,23 +3,20 @@ # Python module entry point(s). PyInit_*; - # Expose the full zvec C++ namespace so the DiskAnn runtime plugin - # (libzvec_diskann_plugin.so), which is dlopen()ed with - # RTLD_NOW | RTLD_GLOBAL after the interpreter has loaded _zvec.so, - # can resolve its undefined references against this module. Without - # this, the plugin fails to load with errors like - # undefined symbol: _ZN4zvec6ailego6Logger10LEVEL_INFOE - # because the default version script hides every internal symbol. + # Expose the full zvec C++ namespace so that the C API binding + # (and any C++ consumer that links against _zvec.so) can resolve + # zvec:: symbols. DiskAnn is now statically linked into _zvec.so, + # so these exports are no longer needed for runtime plugin loading, + # but they are still required for the C API and for C++ consumers + # that link against the shared module. extern "C++" { "zvec::*"; zvec::*; # Also export the compiler-generated helper symbols that live # alongside symbols in the zvec namespace (guard variables for # static locals, vtables, typeinfo, VTT, construction vtables, - # thunks). Without these, the DiskAnn plugin and _zvec.so each - # get their own copy of e.g. the guard for Factory::Instance's - # static local, which causes the factory constructor to run - # twice - wiping out registrations done during _zvec.so load. + # thunks). These are needed so that factory singletons and + # vtables are unified across translation units. "guard variable for zvec::*"; "vtable for zvec::*"; "VTT for zvec::*"; diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index 6405c3220..15e684188 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -63,9 +63,10 @@ endif() # Always exclude algorithm/diskann implementation files from zvec_core. # The DiskAnn algorithm is provided by the separate core_knn_diskann library -# (real on Linux x86_64, stub on other platforms). Including them here causes -# duplicate symbols and missing -laio when test binaries link both zvec_core -# (via zvec) and core_knn_diskann. +# (STATIC+SHARED, real on Linux x86_64, stub on other platforms). The static +# variant is whole-archived into _zvec.so for the Python wheel; the shared +# variant is used by C++ tools and tests. Including the sources here would +# cause duplicate symbols. list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*") if(NOT DISKANN_SUPPORTED) list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc") diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index e2fa911f5..5f1b7bf7f 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -1,56 +1,34 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) file(GLOB_RECURSE ALL_SRCS *.cc *.c) -# The DiskAnn plugin is loaded at runtime via zvec::LoadDiskAnnPlugin() with -# RTLD_GLOBAL, so its undefined references to internal zvec symbols -# (core_framework, core_knn_cluster, zvec_ailego, ...) are resolved at load -# time against the hosting binary (_zvec.so for the Python extension, the -# test executable for gtest, or libzvec_core for tools). -# -# As a consequence the plugin .so must NOT carry NEEDED entries for those -# libs: otherwise the dynamic loader would try to resolve them from disk -# before invoking dlopen's RTLD_GLOBAL symbol sharing, and the Python wheel -# (which only ships _zvec.so + the plugin) would fail to load. +# DiskAnn is compiled as a STATIC+SHARED library (matching the pattern used by +# core_knn_hnsw, core_knn_cluster, etc.). The static variant +# (core_knn_diskann_static) is whole-archived into _zvec.so for the Python +# wheel, so no separate runtime .so is needed. The shared variant is used by +# C++ tools and tests that link against it directly. # # libaio is loaded at runtime via dlopen()/dlsym() (see libaio_loader.h), so # we do NOT link against -laio. The only system-level link dependency is # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. -set(CORE_KNN_DISKANN_LIBS "") +set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster) if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() +if(NOT APPLE) + set(CORE_KNN_DISKANN_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + cc_library( NAME core_knn_diskann - SHARED STRICT + STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc LIBS ${CORE_KNN_DISKANN_LIBS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" -) - -# Internal zvec libs are referenced only for header availability; the actual -# symbol resolution happens at dlopen time. Expose their public include dirs -# without adding them as link dependencies. -foreach(_dep zvec_ailego core_framework core_knn_cluster) - if(TARGET ${_dep}) - target_include_directories(core_knn_diskann PRIVATE - $) - add_dependencies(core_knn_diskann ${_dep}) - endif() -endforeach() - -# Allow the plugin to have unresolved symbols that will be satisfied at -# dlopen(RTLD_NOW | RTLD_GLOBAL) time by the hosting binary. -if(CMAKE_SYSTEM_NAME STREQUAL "Linux") - target_link_options(core_knn_diskann PRIVATE - "LINKER:--unresolved-symbols=ignore-all") -endif() - -# Rename the artifact to libzvec_diskann_plugin.so so it is clearly identified -# as an optional plugin that users load at runtime via zvec::LoadDiskAnnPlugin(). -set_target_properties(core_knn_diskann PROPERTIES - OUTPUT_NAME zvec_diskann_plugin ) \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 34584d935..c8d8bdd0b 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -33,7 +33,8 @@ typedef struct iocb iocb_t; int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) if (!LibAioLoader::Instance().Load()) { - return -ENOSYS; + LOG_WARN("libaio not available; falling back to synchronous pread"); + return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); @@ -199,7 +200,7 @@ void LinuxAlignedFileReader::register_thread() { IOContext ctx = nullptr; if (!LibAioLoader::Instance().Load()) { - LOG_ERROR("libaio not available; cannot register thread for async I/O"); + LOG_WARN("libaio not available; async I/O disabled, will use pread"); lk.unlock(); return; } diff --git a/src/core/algorithm/diskann/libaio_def.h b/src/core/algorithm/diskann/libaio_def.h new file mode 100644 index 000000000..ce95bc982 --- /dev/null +++ b/src/core/algorithm/diskann/libaio_def.h @@ -0,0 +1,190 @@ +// 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 replacement for . +// +// This header declares *only* the types, constants, and inline helpers that +// zvec needs from libaio. By doing so the project is completely decoupled +// from the system libaio-dev header: there is no `#include ` anywhere +// in the source tree, which means libaio-dev does not need to be installed at +// build time and the code is portable to cross-compilation environments that +// lack the header. +// +// The struct layouts (struct iocb, struct io_event, ...) are part of the Linux +// kernel ABI. They are copied verbatim from the upstream , including +// the PADDED macros that handle architecture-specific padding. The inline +// helper io_prep_pread() is likewise copied — it only manipulates struct fields +// and does not call into the library. + +#pragma once + +#include // struct timespec (used by io_getevents signature) +#include // memset() — used by io_prep_pread() inline helper + +#if defined(__linux) || defined(__linux__) + +struct sockaddr; +struct iovec; + +// --------------------------------------------------------------------------- +// Type and struct definitions copied from +// --------------------------------------------------------------------------- + +typedef struct io_context *io_context_t; + +typedef enum io_iocb_cmd { + IO_CMD_PREAD = 0, + IO_CMD_PWRITE = 1, + IO_CMD_FSYNC = 2, + IO_CMD_FDSYNC = 3, + IO_CMD_POLL = 5, + IO_CMD_NOOP = 6, + IO_CMD_PREADV = 7, + IO_CMD_PWRITEV = 8, +} io_iocb_cmd_t; + +// PADDED macros — copied verbatim from to guarantee ABI-compatible +// struct layout on every supported architecture. + +/* little endian, 32 bits */ +#if defined(__i386__) || (defined(__x86_64__) && defined(__ILP32__)) || \ + (defined(__arm__) && !defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__LITTLE_ENDIAN__)) || defined(__bfin__) || \ + (defined(__MIPSEL__) && !defined(__mips64)) || defined(__cris__) || \ + defined(__loongarch32) || (defined(__riscv) && __riscv_xlen == 32) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDptr(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDul(x, y) \ + unsigned long x; \ + unsigned y + +/* little endian, 64 bits */ +#elif defined(__ia64__) || defined(__x86_64__) || defined(__alpha__) || \ + (defined(__mips64) && defined(__MIPSEL__)) || \ + (defined(__aarch64__) && defined(__AARCH64EL__)) || \ + defined(__loongarch64) || (defined(__riscv) && __riscv_xlen == 64) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) x, y +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +/* big endian, 64 bits */ +#elif defined(__powerpc64__) || defined(__s390x__) || \ + (defined(__hppa__) && defined(__arch64__)) || \ + (defined(__sparc__) && defined(__arch64__)) || \ + (defined(__mips64) && defined(__MIPSEB__)) || \ + (defined(__aarch64__) && defined(__AARCH64EB__)) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +/* big endian, 32 bits */ +#elif defined(__PPC__) || defined(__s390__) || \ + (defined(__arm__) && defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__BIG_ENDIAN__)) || defined(__sparc__) || \ + defined(__MIPSEB__) || defined(__m68k__) || defined(__hppa__) || \ + defined(__frv__) || defined(__avr32__) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDul(x, y) \ + unsigned y; \ + unsigned long x + +#else +#error endianness? +#endif + +struct io_iocb_poll { + AIO_PADDED(int events, __pad1); +}; + +struct io_iocb_sockaddr { + AIO_PADDEDptr(struct sockaddr *addr, __pad1); + AIO_PADDEDul(len, __pad2); +}; + +struct io_iocb_common { + AIO_PADDEDptr(void *buf, __pad1); + AIO_PADDEDul(nbytes, __pad2); + long long offset; + long long __pad3; + unsigned flags; + unsigned resfd; +}; + +struct io_iocb_vector { + AIO_PADDEDptr(const struct iovec *vec, __pad1); + AIO_PADDEDul(nr, __pad2); + long long offset; +}; + +struct iocb { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDED(unsigned key, aio_rw_flags); + short aio_lio_opcode; + short aio_reqprio; + int aio_fildes; + union { + struct io_iocb_common c; + struct io_iocb_vector v; + struct io_iocb_poll poll; + struct io_iocb_sockaddr saddr; + } u; +}; + +struct io_event { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDEDptr(struct iocb *obj, __pad2); + AIO_PADDEDul(res, __pad3); + AIO_PADDEDul(res2, __pad4); +}; + +#undef AIO_PADDED +#undef AIO_PADDEDptr +#undef AIO_PADDEDul + +// Inline helper — copied from . Only manipulates struct fields. +static inline void io_prep_pread(struct iocb *iocb, int fd, void *buf, + size_t count, long long offset) { + memset(iocb, 0, sizeof(*iocb)); + iocb->aio_fildes = fd; + iocb->aio_lio_opcode = IO_CMD_PREAD; + iocb->aio_reqprio = 0; + iocb->u.c.buf = buf; + iocb->u.c.nbytes = count; + iocb->u.c.offset = offset; +} + +// --------------------------------------------------------------------------- +// End: type and struct definitions from +// --------------------------------------------------------------------------- + +#endif // __linux__ diff --git a/src/core/algorithm/diskann/libaio_loader.h b/src/core/algorithm/diskann/libaio_loader.h new file mode 100644 index 000000000..dca400ff9 --- /dev/null +++ b/src/core/algorithm/diskann/libaio_loader.h @@ -0,0 +1,144 @@ +// 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. + +// dlopen-based wrapper for libaio. Instead of linking against -laio at build +// time, the DiskAnn plugin loads libaio at runtime via dlopen()/dlsym(). This +// removes libaio as a hard build- and load-time dependency: the plugin .so no +// longer carries a NEEDED entry for libaio.so.1, so it can be dlopen'ed on +// hosts that don't have libaio installed. Callers that actually exercise the +// async-I/O path will fall back to synchronous pread() when libaio is absent. +// +// All ABI-stable type and struct definitions (struct iocb, struct io_event, +// io_context_t, PADDED macros, io_prep_pread(), ...) now live in libaio_def.h — +// a private header that replaces so the project has zero build-time +// dependency on libaio-dev. + +#pragma once + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include +#include "libaio_def.h" // ABI-stable struct definitions (replaces ) + +// Function-pointer typedefs for the four libaio syscalls used by DiskAnn. +typedef int (*aio_setup_fn)(int maxevents, io_context_t *ctxp); +typedef int (*aio_destroy_fn)(io_context_t ctx); +typedef int (*aio_submit_fn)(io_context_t ctx, long nr, struct iocb *ios[]); +typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr, + struct io_event *events, + struct timespec *timeout); + +// Runtime loader for libaio. Thread-safe singleton that dlopen()'s libaio +// once and caches the function pointers. If libaio is not present on the +// host, all pointers remain nullptr and callers should fall back to +// synchronous I/O (pread). +// +// Usage: +// if (LibAioLoader::Instance().Load()) { +// LibAioLoader::Instance().io_setup(...); +// } +class LibAioLoader { + public: + static LibAioLoader &Instance() { + static LibAioLoader instance; + return instance; + } + + // Load (or confirm already loaded) libaio. Returns true on success. + // Thread-safe and idempotent. + bool Load() { + if (available_.load(std::memory_order_acquire)) { + return true; + } + std::call_once(once_, [this] { this->TryLoad(); }); + return available_.load(std::memory_order_relaxed); + } + + bool IsAvailable() const { + return available_.load(std::memory_order_acquire); + } + + // Function pointers — nullptr until Load() succeeds. + aio_setup_fn io_setup; + aio_destroy_fn io_destroy; + aio_submit_fn io_submit; + aio_getevents_fn io_getevents; + + private: + LibAioLoader() + : io_setup(nullptr), + io_destroy(nullptr), + io_submit(nullptr), + io_getevents(nullptr) {} + + ~LibAioLoader() { + if (handle_ != nullptr) { + dlclose(handle_); + } + } + + LibAioLoader(const LibAioLoader &) = delete; + LibAioLoader &operator=(const LibAioLoader &) = delete; + + void TryLoad() { + // On Ubuntu 24.04 the libaio package was renamed with the t64 suffix + // (64-bit time_t transition), so probe both spellings. + static constexpr const char *kSonames[] = { + "libaio.so.1", + "libaio.so.1t64", + }; + + for (const char *soname : kSonames) { + void *h = dlopen(soname, RTLD_LAZY); + if (h == nullptr) { + continue; + } + + io_setup = reinterpret_cast(dlsym(h, "io_setup")); + io_destroy = reinterpret_cast(dlsym(h, "io_destroy")); + io_submit = reinterpret_cast(dlsym(h, "io_submit")); + + // io_getevents may be redirected to io_getevents_time64 on 32-bit + // platforms compiled with _TIME_BITS=64. + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents")); + if (io_getevents == nullptr) { + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents_time64")); + } + + if (io_setup && io_destroy && io_submit && io_getevents) { + handle_ = h; + available_.store(true, std::memory_order_release); + return; + } + + // Some symbols missing — try the next soname. + dlclose(h); + io_setup = nullptr; + io_destroy = nullptr; + io_submit = nullptr; + io_getevents = nullptr; + } + } + + std::once_flag once_; + std::atomic available_{false}; + void *handle_{nullptr}; +}; + +#endif // __linux__ diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 7615ed354..1ec54f265 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -16,61 +16,16 @@ #include #include #include -#include #include "algorithm/diskann/diskann_params.h" #include "holder_builder.h" namespace zvec::core_interface { -namespace { - -// Implicitly bring the DiskAnn runtime online on first use. This keeps the -// DiskAnn index an ordinary public API (users just instantiate a -// DiskAnnIndexParam) while still letting the rest of the library — HNSW, -// IVF, Flat, Vamana — run on hosts that happen to lack libaio. On such -// hosts only DiskAnn fails, with a clear, actionable error message, and -// every other index type stays fully functional. -int EnsureDiskAnnRuntimeReady() { - static std::once_flag once; - static int cached_result = 0; - std::call_once(once, []() { - const int status = ::zvec::LoadDiskAnnPlugin(); - if (status == kDiskAnnPluginOk) { - cached_result = 0; - return; - } - switch (status) { - case kDiskAnnPluginLibAioMissing: - LOG_ERROR( - "DiskAnn requires libaio at runtime, but it was not found on this " - "host. Install it (e.g. 'apt-get install libaio1' on " - "Debian/Ubuntu, " - "or 'libaio1t64' on Ubuntu 24.04+) and retry."); - break; - case kDiskAnnPluginUnsupportedPlatform: - LOG_ERROR("DiskAnn is only supported on Linux x86_64."); - break; - case kDiskAnnPluginDlopenFailed: - default: - LOG_ERROR("Failed to initialize the DiskAnn runtime (status=%d).", - status); - break; - } - cached_result = core::IndexError_Runtime; - }); - return cached_result; -} - -} // namespace - int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { - // Fail fast and cleanly if the DiskAnn runtime cannot be brought up on - // this host (most commonly: libaio is missing). The rest of zvec keeps - // running; only DiskAnn is unusable. - if (int rc = EnsureDiskAnnRuntimeReady(); rc != 0) { - return rc; - } - + // Platform and libaio checks are handled earlier at schema validation + // time (schema.cc calls LoadDiskAnnPlugin). If libaio is missing, DiskAnn + // falls back to synchronous pread() in diskann_file_reader.cc — no need + // to re-check here. if (is_sparse_) { LOG_ERROR("Failed to create streamer. Sparse is not Supported."); return core::IndexError_Unsupported; diff --git a/src/core/plugin/CMakeLists.txt b/src/core/plugin/CMakeLists.txt index 510677aed..808e35f57 100644 --- a/src/core/plugin/CMakeLists.txt +++ b/src/core/plugin/CMakeLists.txt @@ -2,13 +2,13 @@ ## Copyright (C) The Software Authors. All rights reserved. ## ## \file CMakeLists.txt -## \brief Build script for the zvec internal plugin-loading glue library. +## \brief Build script for the zvec internal DiskAnn runtime-init glue. ## Lives inside the main zvec_core artifact and provides the ## implicit DiskAnn runtime bring-up used by DiskAnnIndex on first ## use. These APIs are NOT part of the public user surface: users ## simply instantiate ``DiskAnnIndexParam`` / ``DiskAnnIndex`` and -## the runtime (``libzvec_diskann_plugin.so``) is loaded behind -## the scenes. +## libaio is loaded behind the scenes via dlopen(). DiskAnn itself +## is compiled directly into the hosting binary. ## include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) diff --git a/src/core/plugin/diskann_plugin.cc b/src/core/plugin/diskann_plugin.cc index 9d4e9d0f5..d9af30276 100644 --- a/src/core/plugin/diskann_plugin.cc +++ b/src/core/plugin/diskann_plugin.cc @@ -15,7 +15,6 @@ #include #include #include -#include #include #include @@ -28,295 +27,97 @@ #include #endif +// Include the libaio dlopen wrapper so we can load libaio eagerly at DiskAnn +// bring-up time (instead of waiting for the first I/O operation). +#if defined(__linux__) || defined(__linux) +#include "algorithm/diskann/libaio_loader.h" +#endif + namespace zvec { namespace { #if defined(__linux__) || defined(__linux) -constexpr const char *kPluginFileName = "libzvec_diskann_plugin.so"; -// Candidate soname list. On Ubuntu 24.04 the libaio package was renamed with -// the t64 suffix (64-bit time_t transition), so we probe both spellings. -constexpr const char *kLibAioSoNames[] = { - "libaio.so.1", - "libaio.so.1t64", -}; -constexpr bool kPlatformSupportsDiskAnnPlugin = true; +constexpr bool kPlatformSupportsDiskAnn = true; #elif defined(__APPLE__) -[[maybe_unused]] constexpr const char *kPluginFileName = - "libzvec_diskann_plugin.dylib"; -constexpr bool kPlatformSupportsDiskAnnPlugin = false; +constexpr bool kPlatformSupportsDiskAnn = false; #else -[[maybe_unused]] constexpr const char *kPluginFileName = - "zvec_diskann_plugin.dll"; -constexpr bool kPlatformSupportsDiskAnnPlugin = false; +constexpr bool kPlatformSupportsDiskAnn = false; #endif -// Global plugin handle. Nullptr means "not loaded". -std::atomic g_plugin_handle{nullptr}; -std::mutex g_plugin_mutex; - -#if defined(__linux__) || defined(__linux) - -// Resolve the directory containing the currently running executable, so we -// can look for the plugin next to it regardless of the working directory. -std::string GetExecutableDir() { - char buf[PATH_MAX]; - ssize_t n = ::readlink("/proc/self/exe", buf, sizeof(buf) - 1); - if (n <= 0) { - return {}; - } - buf[n] = '\0'; - std::string path(buf); - auto slash = path.find_last_of('/'); - if (slash == std::string::npos) { - return {}; - } - return path.substr(0, slash); -} - -// Resolve the directory containing the shared object that hosts this -// function. For Python wheels this is the directory of -// ``_zvec.cpython-*.so``; for regular C++ binaries it is the directory of -// ``libzvec_core.so``. In either case the DiskAnn plugin is shipped -// alongside, so this is the most reliable lookup location. -// -// NOTE: we pass the address of an *exported* function (LoadDiskAnnPlugin) -// rather than one in this anonymous namespace, because dladdr() on a symbol -// with internal linkage can report the main executable instead of the -// hosting shared object when the translation unit is whole-archived into -// another .so. -std::string ResolveHostingSoDir() { - ::Dl_info info{}; - if (::dladdr(reinterpret_cast(&::zvec::LoadDiskAnnPlugin), &info) == - 0 || - info.dli_fname == nullptr) { - return {}; - } - std::string path(info.dli_fname); - auto slash = path.find_last_of('/'); - if (slash == std::string::npos) { - return {}; - } - return path.substr(0, slash); -} - -// Full path of the shared object that hosts LoadDiskAnnPlugin, or empty -// string on failure. -std::string ResolveHostingSoPath() { - ::Dl_info info{}; - if (::dladdr(reinterpret_cast(&::zvec::LoadDiskAnnPlugin), &info) == - 0 || - info.dli_fname == nullptr) { - return {}; - } - return std::string(info.dli_fname); -} - -// Promote the hosting shared object (e.g. the Python extension module -// ``_zvec.cpython-*.so``) to the global symbol scope. Python loads C -// extensions with RTLD_LOCAL by default, which means their C++ symbols are -// invisible to subsequently dlopen(RTLD_GLOBAL)ed libraries. Without this -// promotion, the DiskAnn plugin's undefined references to zvec:: symbols -// cannot be resolved against the already-loaded host module and dlopen -// fails with messages like: -// -// undefined symbol: _ZN4zvec6ailego6Logger10LEVEL_INFOE -// -// Using RTLD_NOLOAD re-opens the existing image without loading it again, -// while RTLD_GLOBAL merges its symbols into the global scope. This is a -// no-op for hosts that were already loaded with RTLD_GLOBAL. -void PromoteHostingSoToGlobal() { - const std::string host = ResolveHostingSoPath(); - if (host.empty()) { - return; - } - // When LoadDiskAnnPlugin is statically linked into the main executable, - // dladdr resolves to the executable itself. The main exe's symbols are - // already in the global scope by definition, so skip promotion. - const std::string exe_path = GetExecutableDir(); - if (!exe_path.empty()) { - char exe_buf[PATH_MAX]; - ssize_t n = ::readlink("/proc/self/exe", exe_buf, sizeof(exe_buf) - 1); - if (n > 0) { - exe_buf[n] = '\0'; - // Compare resolved real paths to handle relative vs absolute. - char host_real[PATH_MAX]; - char exe_real[PATH_MAX]; - if (::realpath(host.c_str(), host_real) != nullptr && - ::realpath(exe_buf, exe_real) != nullptr && - std::string(host_real) == std::string(exe_real)) { - // Host IS the main executable; symbols are already global. - return; - } - } - } - void *h = ::dlopen(host.c_str(), RTLD_NOW | RTLD_GLOBAL | RTLD_NOLOAD); - if (h == nullptr) { - const char *err = ::dlerror(); - LOG_WARN("Could not promote host '%s' to RTLD_GLOBAL: %s", host.c_str(), - err ? err : "unknown"); - return; - } - // We purposely keep the handle refcount incremented for the life of the - // process: there is no safe point at which to dlclose() it, and doing so - // would only decrement the count anyway (the image remains mapped because - // Python still holds its own reference). - (void)h; -} - -// Build the list of candidate paths for the plugin. -std::vector BuildCandidatePaths(const std::string &explicit_path) { - std::vector candidates; - if (!explicit_path.empty()) { - candidates.push_back(explicit_path); - return candidates; - } - - // Helper that pushes both ``/`` and ``/../lib/`` - // to handle the conventional CMake build layout where executables live in - // ``bin/`` while shared objects (including this plugin) live in ``lib/``. - auto push_dir_candidates = [&candidates](const std::string &dir) { - if (dir.empty()) { - return; - } - candidates.push_back(dir + "/" + kPluginFileName); - candidates.push_back(dir + "/../lib/" + kPluginFileName); - }; - - // 1. Directory of the library that hosts LoadDiskAnnPlugin (e.g. the - // Python extension module or libzvec_core). This works for Python, - // C++ embedding, and most packaging layouts. - const std::string own_dir = ResolveHostingSoDir(); - push_dir_candidates(own_dir); - // 2. Directory of the running executable. Useful for self-contained C++ - // tools that drop the plugin next to their binary, as well as for the - // standard CMake bin/lib split (handled by ``../lib/`` above). - const std::string exe_dir = GetExecutableDir(); - if (!exe_dir.empty() && exe_dir != own_dir) { - push_dir_candidates(exe_dir); - } - // 3. Fallback: rely on the dynamic linker's default search path - // (RPATH / LD_LIBRARY_PATH / /etc/ld.so.conf). - candidates.emplace_back(kPluginFileName); - return candidates; -} - -#endif // linux +// Tracks whether the DiskAnn runtime has been initialised. Since DiskAnn is +// now compiled directly into the hosting binary (_zvec.so for Python, the test +// executable for gtest, etc.) there is no separate .so to dlopen; the "loaded" +// flag simply means libaio has been probed and the result cached. +std::atomic g_runtime_ready{false}; +std::mutex g_runtime_mutex; } // namespace bool IsLibAioAvailable() { #if defined(__linux__) || defined(__linux) - const char *kRequiredSymbols[] = {"io_setup", "io_submit", "io_getevents", - "io_destroy"}; - for (const char *soname : kLibAioSoNames) { - // RTLD_LAZY keeps the cost low; we only need to know whether the library - // is resolvable and exposes the symbols DiskAnn actually calls. - void *handle = ::dlopen(soname, RTLD_LAZY); - if (handle == nullptr) { - continue; - } - bool ok = true; - for (const char *sym : kRequiredSymbols) { - if (::dlsym(handle, sym) == nullptr) { - ok = false; - break; - } - } - ::dlclose(handle); - if (ok) { - return true; - } - } - return false; + // Use the LibAioLoader singleton so we share the cached dlopen handle with + // the DiskAnn file reader. Load() is idempotent and thread-safe. + return LibAioLoader::Instance().Load(); #else return false; #endif } bool IsDiskAnnPluginLoaded() { - return g_plugin_handle.load(std::memory_order_acquire) != nullptr; + return g_runtime_ready.load(std::memory_order_acquire); } int LoadDiskAnnPlugin(const std::string &path) { - if (!kPlatformSupportsDiskAnnPlugin) { + (void)path; // No external path needed — DiskAnn is linked in statically. + + if (!kPlatformSupportsDiskAnn) { LOG_ERROR( - "DiskAnn plugin is not supported on this platform; it is only " + "DiskAnn is not supported on this platform; it is only " "available on Linux x86_64 with libaio."); return kDiskAnnPluginUnsupportedPlatform; } #if defined(__linux__) || defined(__linux) - // Fast path: already loaded. - if (g_plugin_handle.load(std::memory_order_acquire) != nullptr) { + // Fast path: already initialised. + if (g_runtime_ready.load(std::memory_order_acquire)) { return kDiskAnnPluginOk; } - std::lock_guard lock(g_plugin_mutex); - if (g_plugin_handle.load(std::memory_order_relaxed) != nullptr) { + std::lock_guard lock(g_runtime_mutex); + if (g_runtime_ready.load(std::memory_order_relaxed)) { return kDiskAnnPluginOk; } - if (!IsLibAioAvailable()) { - LOG_ERROR( - "libaio is not available on this host; the DiskAnn runtime cannot be " - "activated. Install libaio1 (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry. This does not affect " - "other index types (HNSW, IVF, Flat, Vamana)."); + // Eagerly load libaio at DiskAnn bring-up time so the user gets immediate + // feedback (success or failure) rather than a delayed error on the first + // async-I/O operation. + LOG_INFO("DiskAnn: initializing runtime — loading libaio ..."); + if (!LibAioLoader::Instance().Load()) { + LOG_WARN( + "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " + "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " + "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " + "to synchronous pread() — performance will be degraded. Other " + "index types (HNSW/IVF/Flat/Vamana) are unaffected."); + // We still mark the runtime as "ready" — DiskAnn code is linked in and + // the file reader will gracefully fall back to pread(). The user gets + // a clear warning now rather than a hard failure later. + g_runtime_ready.store(true, std::memory_order_release); return kDiskAnnPluginLibAioMissing; } - const std::vector candidates = BuildCandidatePaths(path); - // Ensure the hosting module's C++ symbols (zvec::*) are visible to the - // plugin at dlopen time. See PromoteHostingSoToGlobal() for the rationale. - PromoteHostingSoToGlobal(); - void *handle = nullptr; - std::string last_error; - for (const std::string &candidate : candidates) { - // RTLD_GLOBAL so the plugin's factory registrations (which live in the - // plugin's own static-init code) can reference symbols from the main - // library, and any callers that later dlsym against the process can see - // the plugin's symbols. - handle = ::dlopen(candidate.c_str(), RTLD_NOW | RTLD_GLOBAL); - if (handle != nullptr) { - LOG_INFO("Loaded DiskAnn plugin from: %s", candidate.c_str()); - break; - } - const char *err = ::dlerror(); - last_error = err ? err : "unknown dlopen error"; - LOG_DEBUG("dlopen(%s) failed: %s", candidate.c_str(), last_error.c_str()); - } - - if (handle == nullptr) { - LOG_ERROR("Failed to load DiskAnn plugin; last error: %s", - last_error.c_str()); - return kDiskAnnPluginDlopenFailed; - } - - g_plugin_handle.store(handle, std::memory_order_release); + LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + g_runtime_ready.store(true, std::memory_order_release); return kDiskAnnPluginOk; #else - (void)path; return kDiskAnnPluginUnsupportedPlatform; #endif } bool UnloadDiskAnnPlugin() { -#if defined(__linux__) || defined(__linux) - std::lock_guard lock(g_plugin_mutex); - void *handle = g_plugin_handle.exchange(nullptr, std::memory_order_acq_rel); - if (handle == nullptr) { - return false; - } - if (::dlclose(handle) != 0) { - const char *err = ::dlerror(); - LOG_WARN("dlclose for DiskAnn plugin returned non-zero: %s", - err ? err : "unknown"); - } - return true; -#else + // DiskAnn is statically linked — there is nothing to unload. return false; -#endif } } // namespace zvec diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 9471eee90..bf5b02d38 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -175,25 +175,22 @@ Status FieldSchema::validate() const { if (index_params_->type() == IndexType::DISKANN) { // Probe the DiskAnn runtime eagerly at creation time so unsupported - // platforms (non Linux x86_64), missing libaio, or a missing plugin - // .so fail fast with a clear message instead of surfacing later during - // optimize(). This reuses the same gate DiskAnnIndex applies on first - // use (zvec::LoadDiskAnnPlugin, wrapped by EnsureDiskAnnRuntimeReady). - // All validate() call sites are creation-time only, so triggering the - // plugin load here is safe (and idempotent/cached). + // platforms (non Linux x86_64) fail fast with a clear message instead + // of surfacing later during optimize(). All validate() call sites + // are creation-time only, so triggering the runtime init here is + // safe (and idempotent/cached). + // + // kDiskAnnPluginLibAioMissing is non-fatal: DiskAnn falls back to + // synchronous pread() with degraded performance. const int rc = ::zvec::LoadDiskAnnPlugin(); switch (rc) { case kDiskAnnPluginOk: + case kDiskAnnPluginLibAioMissing: break; case kDiskAnnPluginUnsupportedPlatform: return Status::NotSupported( "DiskAnn is not supported on this platform (Linux x86_64 " "only)"); - case kDiskAnnPluginLibAioMissing: - return Status::NotSupported( - "DiskAnn requires libaio at runtime, but it was not found on " - "this host. Install it (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry."); default: return Status::NotSupported( "DiskAnn runtime could not be initialized on this host"); diff --git a/src/include/zvec/ailego/pattern/factory.h b/src/include/zvec/ailego/pattern/factory.h index 9b292b90a..3ffbfcc1f 100644 --- a/src/include/zvec/ailego/pattern/factory.h +++ b/src/include/zvec/ailego/pattern/factory.h @@ -120,15 +120,18 @@ class Factory { //! a guard variable (`_ZGVZN...E7factory`) alongside the object, and both //! must be unified across DSOs for the singleton to be shared. //! - //! In the Python extension build the _zvec.so version script exports the - //! storage (`zvec::*` matches its demangled name) while the guard variable, - //! whose demangled form is `guard variable for zvec::...`, ends up hidden - //! (compilers emit the guard in a COMDAT group whose visibility is not - //! upgraded by our version script). When libzvec_diskann_plugin.so is - //! loaded, _zvec.so and the plugin then share the `factory` storage but - //! each have their own guard; the plugin's still-zero guard triggers a - //! second run of the Factory constructor on the shared storage, wiping - //! all registrations performed during _zvec.so import (e.g. FlatStreamer). + //! Historically (when DiskAnn was a runtime-loaded shared plugin) the + //! _zvec.so version script exported the storage (`zvec::*` matches its + //! demangled name) while the guard variable, whose demangled form is + //! `guard variable for zvec::...`, ended up hidden (compilers emit the + //! guard in a COMDAT group whose visibility is not upgraded by our version + //! script). When the plugin was loaded, _zvec.so and the plugin shared the + //! `factory` storage but each had its own guard; the plugin's still-zero + //! guard triggered a second run of the Factory constructor on the shared + //! storage, wiping out registrations performed during _zvec.so import. + //! DiskAnn is now statically linked into _zvec.so, so this DSO-splitting + //! issue no longer applies, but the atomic-pointer pattern is retained as + //! a defensive measure. //! //! A constant-initialized static std::atomic has NO guard variable //! (its zero init is compile-time), so we use a leaked heap singleton with diff --git a/src/include/zvec/plugin/diskann_plugin.h b/src/include/zvec/plugin/diskann_plugin.h index 988187068..ef1437692 100644 --- a/src/include/zvec/plugin/diskann_plugin.h +++ b/src/include/zvec/plugin/diskann_plugin.h @@ -20,9 +20,14 @@ // invoked implicitly by ``DiskAnnIndex`` on first use so that DiskAnn works // out of the box, without users ever calling a ``load_diskann_plugin()`` / // ``is_libaio_available()`` entry point. External callers should not depend -// on these symbols; they may change or be removed in future releases. On -// hosts missing libaio the bring-up fails cleanly and only DiskAnn becomes -// unavailable — other index types (HNSW / IVF / Flat / Vamana) keep working. +// on these symbols; they may change or be removed in future releases. +// +// DiskAnn is compiled directly into the hosting binary (_zvec.so for the +// Python wheel, the test executable for gtest, or the tool binary for C++ +// tools). ``LoadDiskAnnPlugin`` no longer dlopens a separate .so — it simply +// loads libaio eagerly via dlopen()/dlsym() and logs the result. On hosts +// missing libaio, DiskAnn falls back to synchronous pread() (with a warning) +// while other index types (HNSW / IVF / Flat / Vamana) keep working. #if defined(_WIN32) || defined(__CYGWIN__) #ifdef ZVEC_BUILD_SHARED @@ -58,27 +63,22 @@ enum DiskAnnPluginStatus { // part of the user-facing API. ZVEC_PLUGIN_EXPORT bool IsLibAioAvailable(); -// Load the DiskAnn runtime shared library (libzvec_diskann_plugin.so) via -// dlopen(). Invoked implicitly by ``DiskAnnIndex::CreateAndInitStreamer`` on -// first use; callers should not invoke it directly. The call is idempotent -// and returns ``kDiskAnnPluginOk`` when the runtime is already active. +// Load libaio at runtime via dlopen()/dlsym() and mark the DiskAnn runtime +// as ready. DiskAnn code is compiled directly into the hosting binary, so +// no separate .so is loaded. Invoked implicitly by +// ``DiskAnnIndex::CreateAndInitStreamer`` on first use; callers should not +// invoke it directly. The call is idempotent and returns +// ``kDiskAnnPluginOk`` when the runtime is already active. // -// Search order when ``path`` is empty: -// 1. next to the currently running executable (``/proc/self/exe`` on Linux); -// 2. next to the hosting shared object (e.g. ``_zvec.cpython-*.so``); -// 3. the platform default dynamic-linker search path (RPATH, -// ``LD_LIBRARY_PATH``, ``/etc/ld.so.conf``, …). +// The ``path`` parameter is retained for API compatibility but is ignored. ZVEC_PLUGIN_EXPORT int LoadDiskAnnPlugin(const std::string &path = ""); // Returns true if the DiskAnn runtime is currently loaded in this process. // Internal diagnostic; not a user-facing API. ZVEC_PLUGIN_EXPORT bool IsDiskAnnPluginLoaded(); -// Unload the DiskAnn runtime. Internal only — unloading a library that has -// registered itself into global factory singletons is inherently racy; -// callers must guarantee that no DiskAnn objects are still alive and no -// background threads are executing DiskAnn code before calling this. Returns -// true when a live handle was released. +// Unload the DiskAnn runtime. Since DiskAnn is statically linked, this is +// a no-op that always returns false. Retained for API compatibility. ZVEC_PLUGIN_EXPORT bool UnloadDiskAnnPlugin(); } // namespace zvec diff --git a/tests/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 098765c3d..42b04bf7e 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -107,10 +107,10 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) { // test via the ``core_knn_diskann`` target), its factory entries are // registered automatically and the global ``IndexFactory`` can hand out a // ``DiskAnnBuilder`` without any explicit setup step. On hosts missing -// libaio, DiskAnn would fail at the index-creation layer with a clear error -// while other index types (HNSW/IVF/Flat/Vamana) remain unaffected; that -// runtime branch lives in ``DiskAnnIndex::CreateAndInitStreamer`` and is -// covered by the higher-level interface tests. +// libaio, DiskAnn falls back to synchronous pread() with a warning while +// other index types (HNSW/IVF/Flat/Vamana) remain unaffected; that runtime +// branch lives in ``DiskAnnIndex::CreateAndInitStreamer`` and is covered by +// the higher-level interface tests. TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) { IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); ASSERT_NE(builder, nullptr) From 8ef5d2d8d4681ac174fb6cfad36abff7613b6778 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 26 Jun 2026 14:05:40 +0800 Subject: [PATCH 04/46] fix: log out the warning only once --- src/core/algorithm/diskann/diskann_file_reader.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index c8d8bdd0b..0d4d69cae 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -28,12 +28,18 @@ namespace core { #if (defined(__linux) || defined(__linux__)) typedef struct io_event io_event_t; typedef struct iocb iocb_t; + +// Ensures the libaio-unavailable warning is logged only once per process, +// regardless of how many threads or readers trigger the fallback path. +static std::once_flag g_libaio_warn_once; #endif int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) if (!LibAioLoader::Instance().Load()) { - LOG_WARN("libaio not available; falling back to synchronous pread"); + std::call_once(g_libaio_warn_once, [] { + LOG_WARN("libaio not available; falling back to synchronous pread"); + }); return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); @@ -200,7 +206,9 @@ void LinuxAlignedFileReader::register_thread() { IOContext ctx = nullptr; if (!LibAioLoader::Instance().Load()) { - LOG_WARN("libaio not available; async I/O disabled, will use pread"); + std::call_once(g_libaio_warn_once, [] { + LOG_WARN("libaio not available; async I/O disabled, will use pread"); + }); lk.unlock(); return; } From 50c6c8d9989b6332c19d80ea93f081f0b3d13e85 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 26 Jun 2026 15:16:44 +0800 Subject: [PATCH 05/46] fix: add log --- .../algorithm/diskann/diskann_file_reader.cc | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 0d4d69cae..bc19d8eb5 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -29,17 +29,23 @@ namespace core { typedef struct io_event io_event_t; typedef struct iocb iocb_t; -// Ensures the libaio-unavailable warning is logged only once per process, -// regardless of how many threads or readers trigger the fallback path. -static std::once_flag g_libaio_warn_once; +// 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. +static std::once_flag g_io_backend_log_once; #endif int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - if (!LibAioLoader::Instance().Load()) { - std::call_once(g_libaio_warn_once, [] { - LOG_WARN("libaio not available; falling back to synchronous pread"); - }); + LibAioLoader::Instance().Load(); + std::call_once(g_io_backend_log_once, [] { + if (LibAioLoader::Instance().IsAvailable()) { + LOG_INFO("DiskAnn I/O backend: libaio (async I/O enabled)"); + } else { + LOG_WARN("DiskAnn I/O backend: synchronous pread (libaio not available)"); + } + }); + if (!LibAioLoader::Instance().IsAvailable()) { return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); @@ -205,10 +211,15 @@ void LinuxAlignedFileReader::register_thread() { IOContext ctx = nullptr; - if (!LibAioLoader::Instance().Load()) { - std::call_once(g_libaio_warn_once, [] { - LOG_WARN("libaio not available; async I/O disabled, will use pread"); - }); + LibAioLoader::Instance().Load(); + std::call_once(g_io_backend_log_once, [] { + if (LibAioLoader::Instance().IsAvailable()) { + LOG_INFO("DiskAnn I/O backend: libaio (async I/O enabled)"); + } else { + LOG_WARN("DiskAnn I/O backend: synchronous pread (libaio not available)"); + } + }); + if (!LibAioLoader::Instance().IsAvailable()) { lk.unlock(); return; } From ceee226125f6192077c006f250478d7b55eeaf96 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 26 Jun 2026 15:39:07 +0800 Subject: [PATCH 06/46] fix: add log --- tools/core/bench_original.cc | 1 + tools/core/recall_original.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/tools/core/bench_original.cc b/tools/core/bench_original.cc index 65a36d7ba..796ea399f 100644 --- a/tools/core/bench_original.cc +++ b/tools/core/bench_original.cc @@ -886,6 +886,7 @@ int main(int argc, char *argv[]) { transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); + zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]); } // Calculate Bench diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index c8b2068ef..bf9b2c8a4 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -1883,6 +1883,7 @@ int main(int argc, char *argv[]) { transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); + zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]); } // Calculate Recall From a2524ebd6f1c7ade392e57306ca110575c72d72d Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 10:53:22 +0800 Subject: [PATCH 07/46] fix: fix windows build --- src/core/algorithm/CMakeLists.txt | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index bed772984..f874eba62 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -23,7 +23,12 @@ else() if(MSVC) # MSVC: STATIC-only stub to avoid creating an empty DLL with no exports - # (MSVC linker fails to produce an import library when there are zero exports) + # (MSVC linker fails to produce an import library when there are zero exports). + # cc_library with STATIC-only creates target "core_knn_diskann" but NOT the + # "core_knn_diskann_static" variant (that only happens with STATIC+SHARED). + # The Python binding references core_knn_diskann_static on all platforms, so + # create an ALIAS so the same target name works under MSVC as well. + # ($ supports ALIAS targets since CMake 3.18.) cc_library( NAME core_knn_diskann STATIC STRICT ALWAYS_LINK @@ -32,6 +37,7 @@ else() INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm VERSION "${PROXIMA_ZVEC_VERSION}" ) + add_library(core_knn_diskann_static ALIAS core_knn_diskann) else() cc_library( NAME core_knn_diskann From 6688e1ab5baecefac3c5d6e30b7ddae0a7f299f7 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 20:31:59 +0800 Subject: [PATCH 08/46] fix: remove libaio dependency --- .github/workflows/03-macos-linux-build.yml | 8 -------- .github/workflows/08-cmake-subproject-integration.yml | 7 ------- .github/workflows/clang_tidy.yml | 2 +- .github/workflows/nightly_coverage.yml | 2 +- CMakeLists.txt | 2 +- pyproject.toml | 6 +----- 6 files changed, 4 insertions(+), 23 deletions(-) diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index 4e4f1626f..03fe2cafa 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -60,14 +60,6 @@ jobs: sudo apt-get install -y clang libomp-dev shell: bash - - name: Install AIO - if: runner.os == 'Linux' && runner.arch == 'X64' - run: | - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - libaio-dev - shell: bash - - name: Print CPU info if: runner.os == 'Linux' run: lscpu diff --git a/.github/workflows/08-cmake-subproject-integration.yml b/.github/workflows/08-cmake-subproject-integration.yml index 564a30484..3c5d50e17 100644 --- a/.github/workflows/08-cmake-subproject-integration.yml +++ b/.github/workflows/08-cmake-subproject-integration.yml @@ -34,13 +34,6 @@ jobs: cache: 'pip' cache-dependency-path: 'pyproject.toml' - - name: Install system dependencies - run: | - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - libaio-dev - shell: bash - - name: Set up environment variables run: | NPROC=$(nproc 2>/dev/null || echo 2) diff --git a/.github/workflows/clang_tidy.yml b/.github/workflows/clang_tidy.yml index 20a302d73..5ab872d58 100644 --- a/.github/workflows/clang_tidy.yml +++ b/.github/workflows/clang_tidy.yml @@ -47,7 +47,7 @@ jobs: if: steps.changed_files.outputs.any_changed == 'true' run: | sudo apt-get update - sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev libaio-dev + sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev - name: Setup ccache if: steps.changed_files.outputs.any_changed == 'true' diff --git a/.github/workflows/nightly_coverage.yml b/.github/workflows/nightly_coverage.yml index 81429b1eb..4c9aca50b 100644 --- a/.github/workflows/nightly_coverage.yml +++ b/.github/workflows/nightly_coverage.yml @@ -49,7 +49,7 @@ jobs: run: | sudo apt-get update sudo apt-get install -y --no-install-recommends \ - lcov libaio-dev + lcov shell: bash - name: Install dependencies diff --git a/CMakeLists.txt b/CMakeLists.txt index 2ba9b597d..956ee5599 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -122,7 +122,7 @@ else() endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") -# DiskAnn support (Linux x86_64 only, requires libaio) +# DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen) if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) set(DISKANN_SUPPORTED ON) add_definitions(-DDISKANN_SUPPORTED=1) diff --git a/pyproject.toml b/pyproject.toml index 4f7add507..bce1670d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,11 +181,7 @@ test-command = "cd {project} && pytest python/tests -v --tb=short" build-verbosity = 1 [tool.cibuildwheel.linux] -# libaio is required by the C++ backend; install it inside the manylinux -# container (manylinux_2_28 is AlmaLinux 8 based, so use dnf/libaio-devel). -# libaio-devel lives in BaseOS; disable EPEL so a flaky EPEL mirror cannot -# break metadata refresh. -before-all = "dnf install -y --disablerepo=epel libaio-devel" +# libaio is loaded at runtime via dlopen (no build-time dependency). archs = ["auto"] environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" } manylinux-x86_64-image = "manylinux_2_28" From f2559b56734b2b30dc00b923317dd3b35d8bac93 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 19:20:48 +0800 Subject: [PATCH 09/46] fix: remove no use code --- tools/core/bench_original.cc | 1 - tools/core/recall_original.cc | 1 - 2 files changed, 2 deletions(-) diff --git a/tools/core/bench_original.cc b/tools/core/bench_original.cc index 796ea399f..c787b883b 100644 --- a/tools/core/bench_original.cc +++ b/tools/core/bench_original.cc @@ -885,7 +885,6 @@ int main(int argc, char *argv[]) { : "debug"; transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { - IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]); } diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index bf9b2c8a4..997ef5ad7 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -1882,7 +1882,6 @@ int main(int argc, char *argv[]) { : "debug"; transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { - IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]); } From d7530802df74abaab63d4c17f9a574087ee414bb Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 19:39:29 +0800 Subject: [PATCH 10/46] fix: fix makefile --- thirdparty/arrow/CMakeLists.txt | 45 +++++---------------------------- 1 file changed, 7 insertions(+), 38 deletions(-) diff --git a/thirdparty/arrow/CMakeLists.txt b/thirdparty/arrow/CMakeLists.txt index 30d77545b..a20fa8816 100644 --- a/thirdparty/arrow/CMakeLists.txt +++ b/thirdparty/arrow/CMakeLists.txt @@ -42,45 +42,14 @@ endif() set(CONFIGURE_ENV_LIST "") if(USE_OSS_MIRROR) - set(_OSS_BASE "https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty") - - # Prefer local tarballs (manually placed in repo root) over network - # downloads to avoid flaky mirrors, HTTP/2 PROTOCOL_ERROR, or hash - # mismatches. Falls back to the OSS mirror when the local file is - # absent. Verified SHA-256 hashes match Arrow's expectations: - # boost-1.88.0-cmake.tar.gz dcea50f4... - # thrift-0.22.0.tar.gz 794a0e45... - # 13.0.0.tar.gz (xsimd) 8bdbbad0... - set(_LOCAL_BOOST "${PROJECT_ROOT_DIR}/boost-1.88.0-cmake.tar.gz") - if(EXISTS "${_LOCAL_BOOST}") - list(APPEND CONFIGURE_ENV_LIST "ARROW_BOOST_URL=file://${_LOCAL_BOOST}") - message(STATUS "Arrow: using local Boost tarball") - else() - list(APPEND CONFIGURE_ENV_LIST "ARROW_BOOST_URL=${_OSS_BASE}/boost-1.88.0-cmake.tar.gz") - endif() - - set(_LOCAL_THRIFT "${PROJECT_ROOT_DIR}/thrift-0.22.0.tar.gz") - if(EXISTS "${_LOCAL_THRIFT}") - list(APPEND CONFIGURE_ENV_LIST "ARROW_THRIFT_URL=file://${_LOCAL_THRIFT}") - message(STATUS "Arrow: using local Thrift tarball") - else() - list(APPEND CONFIGURE_ENV_LIST "ARROW_THRIFT_URL=${_OSS_BASE}/thrift-0.22.0.tar.gz") - endif() - - set(_LOCAL_XSIMD "${PROJECT_ROOT_DIR}/13.0.0.tar.gz") - if(EXISTS "${_LOCAL_XSIMD}") - list(APPEND CONFIGURE_ENV_LIST "ARROW_XSIMD_URL=file://${_LOCAL_XSIMD}") - message(STATUS "Arrow: using local xsimd tarball") - else() - list(APPEND CONFIGURE_ENV_LIST "ARROW_XSIMD_URL=${_OSS_BASE}/xsimd-13.0.0.tar.gz") - endif() - - # Dependencies without local copies — always use OSS mirror list(APPEND CONFIGURE_ENV_LIST - "ARROW_RAPIDJSON_URL=${_OSS_BASE}/rapidjson-232389d4f1012dddec4ef84861face2d2ba85709.tar.gz" - "ARROW_RE2_URL=${_OSS_BASE}/re2-2022-06-01.tar.gz" - "ARROW_UTF8PROC_URL=${_OSS_BASE}/utf8proc-2.10.0.tar.gz" - "ARROW_ZLIB_URL=${_OSS_BASE}/zlib-1.3.1.tar.gz" + "ARROW_BOOST_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/boost-1.88.0-cmake.tar.gz" + "ARROW_RAPIDJSON_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/rapidjson-232389d4f1012dddec4ef84861face2d2ba85709.tar.gz" + "ARROW_RE2_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/re2-2022-06-01.tar.gz" + "ARROW_THRIFT_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/thrift-0.22.0.tar.gz" + "ARROW_UTF8PROC_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/utf8proc-2.10.0.tar.gz" + "ARROW_XSIMD_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/xsimd-13.0.0.tar.gz" + "ARROW_ZLIB_URL=https://zvec-bj.oss-cn-beijing.aliyuncs.com/thirdparty/zlib-1.3.1.tar.gz" ) message(STATUS "Using OSS mirror for third-party downloads") endif() From 292322c115ce478cd6c8ca9aa1c3aa4ed4bb22e0 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 3 Jul 2026 11:28:47 +0800 Subject: [PATCH 11/46] fix: fix ut --- python/tests/test_collection_diskann.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 92cbc74e4..93382aba4 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -32,7 +32,6 @@ from __future__ import annotations import math -import os import platform import sys @@ -46,13 +45,6 @@ reason="DiskAnn plugin is only supported on Linux x86_64", ) -# Promote all symbols in subsequently-loaded DSOs to the global namespace and -# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see -# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa. -# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement. -if sys.platform == "linux": - sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW) - import zvec # noqa: E402 from zvec import ( # noqa: E402 From a0a9ce98b441843340d0f5bb02da03fbf0f4fbaa Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 3 Jul 2026 11:40:49 +0800 Subject: [PATCH 12/46] fix: remove export map --- src/binding/python/CMakeLists.txt | 14 -------------- src/binding/python/exports.map | 27 +-------------------------- 2 files changed, 1 insertion(+), 40 deletions(-) diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index 4cc53b6da..e6439a55a 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -19,20 +19,6 @@ set(SRC_LISTS pybind11_add_module(_zvec ${SRC_LISTS}) -# pybind11_add_module() defaults to CXX_VISIBILITY_PRESET=hidden + -# VISIBILITY_INLINES_HIDDEN=ON, which hides the compiler-generated helper -# symbols attached to inline functions (guard variables for static locals, -# vtables, typeinfo ...). DiskAnn is now whole-archived directly into -# _zvec.so (no separate runtime .so), so guard-variable duplication across -# DSOs is no longer a concern. We still keep default visibility so that -# the version script (exports.map / exports.mac) — not the compiler — is -# the sole gatekeeper of the dynamic symbol table, exporting only zvec::* -# and PyInit_*. -set_target_properties(_zvec PROPERTIES - CXX_VISIBILITY_PRESET default - C_VISIBILITY_PRESET default - VISIBILITY_INLINES_HIDDEN OFF) - # Ensure any change to the linker version script (exports.map) triggers a # re-link of _zvec.so. target_link_options() alone is a command-line flag # and does not register the script file as a build dependency, so stale diff --git a/src/binding/python/exports.map b/src/binding/python/exports.map index 76cc27953..66d92e04f 100644 --- a/src/binding/python/exports.map +++ b/src/binding/python/exports.map @@ -1,31 +1,6 @@ { global: - # Python module entry point(s). - PyInit_*; - - # Expose the full zvec C++ namespace so that the C API binding - # (and any C++ consumer that links against _zvec.so) can resolve - # zvec:: symbols. DiskAnn is now statically linked into _zvec.so, - # so these exports are no longer needed for runtime plugin loading, - # but they are still required for the C API and for C++ consumers - # that link against the shared module. - extern "C++" { - "zvec::*"; - zvec::*; - # Also export the compiler-generated helper symbols that live - # alongside symbols in the zvec namespace (guard variables for - # static locals, vtables, typeinfo, VTT, construction vtables, - # thunks). These are needed so that factory singletons and - # vtables are unified across translation units. - "guard variable for zvec::*"; - "vtable for zvec::*"; - "VTT for zvec::*"; - "typeinfo for zvec::*"; - "typeinfo name for zvec::*"; - "construction vtable for zvec::*"; - "non-virtual thunk to zvec::*"; - "virtual thunk to zvec::*"; - }; + PyInit_*; # export only python related functions local: *; }; From 69b1e23fad75a1a8043c8433848861c9d4860438 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 3 Jul 2026 15:37:45 +0800 Subject: [PATCH 13/46] fix: clean code --- python/tests/detail/fixture_helper.py | 10 ++-- python/tests/test_collection_diskann.py | 2 +- python/zvec/__init__.py | 34 ++++++------- src/binding/python/CMakeLists.txt | 3 -- src/binding/python/binding.cc | 41 ++++++++-------- src/core/CMakeLists.txt | 3 +- src/core/interface/CMakeLists.txt | 7 ++- .../diskann_runtime.cc} | 23 ++++----- src/core/interface/indexes/diskann_index.cc | 2 +- src/core/plugin/CMakeLists.txt | 29 ----------- src/db/index/common/schema.cc | 12 ++--- .../interface/diskann_runtime.h} | 48 ++++++------------- tests/core/algorithm/diskann/CMakeLists.txt | 2 +- .../algorithm/diskann/diskann_builder_test.cc | 4 +- 14 files changed, 84 insertions(+), 136 deletions(-) rename src/core/{plugin/diskann_plugin.cc => interface/diskann_runtime.cc} (89%) delete mode 100644 src/core/plugin/CMakeLists.txt rename src/include/zvec/{plugin/diskann_plugin.h => core/interface/diskann_runtime.h} (61%) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index 2054546cf..e37fc307a 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -38,17 +38,17 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" return _DISKANN_PRELOAD_REASON - status = zvec.load_diskann_plugin() - if status == zvec.DISKANN_PLUGIN_UNSUPPORTED_PLATFORM: + status = zvec.init_diskann_runtime() + if status == zvec.DISKANN_RUNTIME_UNSUPPORTED_PLATFORM: _DISKANN_PRELOAD_REASON = ( f"DiskAnn is not supported on this platform (status={status})." ) return _DISKANN_PRELOAD_REASON - # DISKANN_PLUGIN_OK and DISKANN_PLUGIN_LIBAIO_MISSING are both acceptable: + # DISKANN_RUNTIME_OK and DISKANN_RUNTIME_LIBAIO_MISSING are both acceptable: # the latter means DiskAnn will use synchronous pread() instead of async I/O. if status not in ( - zvec.DISKANN_PLUGIN_OK, - zvec.DISKANN_PLUGIN_LIBAIO_MISSING, + zvec.DISKANN_RUNTIME_OK, + zvec.DISKANN_RUNTIME_LIBAIO_MISSING, ): _DISKANN_PRELOAD_REASON = ( f"Failed to initialize DiskAnn runtime (status={status})." diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 93382aba4..e90751a04 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -19,7 +19,7 @@ 1. DiskAnn is currently built only for Linux x86_64 — other platforms are skipped wholesale. -2. The DiskAnn runtime is initialized via ``zvec.load_diskann_plugin()``, +2. The DiskAnn runtime is initialized via ``zvec.init_diskann_runtime()``, which eagerly loads libaio via dlopen(). If libaio is missing, DiskAnn falls back to synchronous pread() — the tests still run but with degraded performance. DiskAnn is compiled directly into ``_zvec.so``; there is no diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index 630602be2..cef509077 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -43,19 +43,19 @@ # Public API — grouped by category # ============================== -# —— DiskAnn runtime plugin —— -# Re-export the plugin management entry points defined by the C++ extension. -# DiskAnn normally auto-loads on first use; these APIs let tests and -# diagnostic tools preload the plugin and get a clear error if libaio is -# missing or the plugin shared object cannot be located. +# —— DiskAnn runtime —— +# Re-export the runtime management entry points defined by the C++ extension. +# DiskAnn normally auto-initializes on first use; these APIs let tests and +# diagnostic tools force initialization up-front and get a clear error if +# libaio is missing. from zvec._zvec import ( - DISKANN_PLUGIN_DLOPEN_FAILED, - DISKANN_PLUGIN_LIBAIO_MISSING, - DISKANN_PLUGIN_OK, - DISKANN_PLUGIN_UNSUPPORTED_PLATFORM, - is_diskann_plugin_loaded, + DISKANN_RUNTIME_DLOPEN_FAILED, + DISKANN_RUNTIME_LIBAIO_MISSING, + DISKANN_RUNTIME_OK, + DISKANN_RUNTIME_UNSUPPORTED_PLATFORM, + init_diskann_runtime, + is_diskann_runtime_ready, is_libaio_available, - load_diskann_plugin, ) from . import model as model @@ -204,13 +204,13 @@ # Tools "require_module", # DiskAnn runtime - "load_diskann_plugin", - "is_diskann_plugin_loaded", + "init_diskann_runtime", + "is_diskann_runtime_ready", "is_libaio_available", - "DISKANN_PLUGIN_OK", - "DISKANN_PLUGIN_UNSUPPORTED_PLATFORM", - "DISKANN_PLUGIN_LIBAIO_MISSING", - "DISKANN_PLUGIN_DLOPEN_FAILED", + "DISKANN_RUNTIME_OK", + "DISKANN_RUNTIME_UNSUPPORTED_PLATFORM", + "DISKANN_RUNTIME_LIBAIO_MISSING", + "DISKANN_RUNTIME_DLOPEN_FAILED", ] # ============================== diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index e6439a55a..7524502be 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -47,7 +47,6 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") $ $ $ - $ -Wl,--no-whole-archive zvec ${CMAKE_DL_LIBS} @@ -75,7 +74,6 @@ elseif (APPLE) -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ - -Wl,-force_load,$ zvec ) target_link_libraries(_zvec PRIVATE @@ -95,7 +93,6 @@ elseif (MSVC) core_metric_static core_utility_static core_quantizer_static - core_plugin ) target_link_libraries(_zvec PRIVATE ${_zvec_whole_archive_libs} diff --git a/src/binding/python/binding.cc b/src/binding/python/binding.cc index 4e257472f..2d6332644 100644 --- a/src/binding/python/binding.cc +++ b/src/binding/python/binding.cc @@ -13,7 +13,7 @@ // limitations under the License. #include -#include +#include #include "python_collection.h" #include "python_config.h" #include "python_doc.h" @@ -27,34 +27,35 @@ namespace zvec { namespace { // Expose DiskAnn runtime management to Python. DiskAnn is compiled directly -// into _zvec.so, so "loading" just means eagerly dlopen()-ing libaio and +// into _zvec.so, so "initializing" just means eagerly dlopen()-ing libaio and // caching the result. Tests (and diagnostic tooling) use these entry points -// to force the load up-front and get actionable warnings when libaio is +// to force the init up-front and get actionable warnings when libaio is // missing. -void InitializeDiskAnnPluginBindings(pybind11::module_ &m) { +void InitializeDiskAnnRuntimeBindings(pybind11::module_ &m) { m.def( - "load_diskann_plugin", - [](const std::string &path) { return ::zvec::LoadDiskAnnPlugin(path); }, + "init_diskann_runtime", + [](const std::string &path) { return ::zvec::InitDiskAnnRuntime(path); }, pybind11::arg("path") = std::string(), - "Load libaio for the DiskAnn runtime. Returns DISKANN_PLUGIN_OK (0) " - "on success, or DISKANN_PLUGIN_LIBAIO_MISSING if libaio is not " - "available (DiskAnn falls back to synchronous pread in that case). " - "Returns a negative code for unsupported platforms."); - m.def("is_diskann_plugin_loaded", &::zvec::IsDiskAnnPluginLoaded, - "Return True if the DiskAnn runtime plugin is currently loaded."); + "Initialize the DiskAnn runtime by loading libaio via dlopen(). " + "Returns DISKANN_RUNTIME_OK (0) on success, or " + "DISKANN_RUNTIME_LIBAIO_MISSING if libaio is not available (DiskAnn " + "falls back to synchronous pread in that case). Returns a negative " + "code for unsupported platforms."); + m.def("is_diskann_runtime_ready", &::zvec::IsDiskAnnRuntimeReady, + "Return True if the DiskAnn runtime has been initialized."); m.def("is_libaio_available", &::zvec::IsLibAioAvailable, "Return True if libaio is resolvable on this host (required by the " "DiskAnn runtime)."); // Status constants so callers can compare against well-known codes without // hard-coding integers. - m.attr("DISKANN_PLUGIN_OK") = static_cast(::zvec::kDiskAnnPluginOk); - m.attr("DISKANN_PLUGIN_UNSUPPORTED_PLATFORM") = - static_cast(::zvec::kDiskAnnPluginUnsupportedPlatform); - m.attr("DISKANN_PLUGIN_LIBAIO_MISSING") = - static_cast(::zvec::kDiskAnnPluginLibAioMissing); - m.attr("DISKANN_PLUGIN_DLOPEN_FAILED") = - static_cast(::zvec::kDiskAnnPluginDlopenFailed); + m.attr("DISKANN_RUNTIME_OK") = static_cast(::zvec::kDiskAnnRuntimeOk); + m.attr("DISKANN_RUNTIME_UNSUPPORTED_PLATFORM") = + static_cast(::zvec::kDiskAnnRuntimeUnsupportedPlatform); + m.attr("DISKANN_RUNTIME_LIBAIO_MISSING") = + static_cast(::zvec::kDiskAnnRuntimeLibAioMissing); + m.attr("DISKANN_RUNTIME_DLOPEN_FAILED") = + static_cast(::zvec::kDiskAnnRuntimeDlopenFailed); } } // namespace @@ -69,6 +70,6 @@ PYBIND11_MODULE(_zvec, m) { ZVecPyConfig::Initialize(m); ZVecPyDoc::Initialize(m); ZVecPyCollection::Initialize(m); - InitializeDiskAnnPluginBindings(m); + InitializeDiskAnnRuntimeBindings(m); } } // namespace zvec diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index 15e684188..13e835ef9 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -49,7 +49,6 @@ cc_directory(quantizer) cc_directory(utility) cc_directory(interface) cc_directory(mixed_reducer) -cc_directory(plugin) git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR}) file(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h) @@ -73,7 +72,7 @@ if(NOT DISKANN_SUPPORTED) endif() set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) -# The plugin loader uses dlopen/dlsym, so link libdl on Linux. +# The DiskAnn runtime loader uses dlopen/dlsym, so link libdl on Linux. if(CMAKE_SYSTEM_NAME STREQUAL "Linux") list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) endif() diff --git a/src/core/interface/CMakeLists.txt b/src/core/interface/CMakeLists.txt index d5475a4b5..ffb3cb043 100644 --- a/src/core/interface/CMakeLists.txt +++ b/src/core/interface/CMakeLists.txt @@ -1,10 +1,15 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) +set(CORE_INTERFACE_LIBS zvec_ailego core_framework sparsehash magic_enum rabitqlib) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + list(APPEND CORE_INTERFACE_LIBS ${CMAKE_DL_LIBS}) +endif() + cc_library( NAME core_interface STATIC STRICT ALWAYS_LINK SRCS *.cc indexes/*.cc INCS . ${PROJECT_ROOT_DIR}/src/ ${PROJECT_ROOT_DIR}/src/core - LIBS zvec_ailego core_framework core_plugin sparsehash magic_enum rabitqlib + LIBS ${CORE_INTERFACE_LIBS} VERSION "${PROXIMA_ZVEC_VERSION}" ) diff --git a/src/core/plugin/diskann_plugin.cc b/src/core/interface/diskann_runtime.cc similarity index 89% rename from src/core/plugin/diskann_plugin.cc rename to src/core/interface/diskann_runtime.cc index d9af30276..9fef28690 100644 --- a/src/core/plugin/diskann_plugin.cc +++ b/src/core/interface/diskann_runtime.cc @@ -16,7 +16,7 @@ #include #include #include -#include +#include #if defined(__linux__) || defined(__linux) || defined(__APPLE__) #include @@ -64,29 +64,29 @@ bool IsLibAioAvailable() { #endif } -bool IsDiskAnnPluginLoaded() { +bool IsDiskAnnRuntimeReady() { return g_runtime_ready.load(std::memory_order_acquire); } -int LoadDiskAnnPlugin(const std::string &path) { +int InitDiskAnnRuntime(const std::string &path) { (void)path; // No external path needed — DiskAnn is linked in statically. if (!kPlatformSupportsDiskAnn) { LOG_ERROR( "DiskAnn is not supported on this platform; it is only " "available on Linux x86_64 with libaio."); - return kDiskAnnPluginUnsupportedPlatform; + return kDiskAnnRuntimeUnsupportedPlatform; } #if defined(__linux__) || defined(__linux) // Fast path: already initialised. if (g_runtime_ready.load(std::memory_order_acquire)) { - return kDiskAnnPluginOk; + return kDiskAnnRuntimeOk; } std::lock_guard lock(g_runtime_mutex); if (g_runtime_ready.load(std::memory_order_relaxed)) { - return kDiskAnnPluginOk; + return kDiskAnnRuntimeOk; } // Eagerly load libaio at DiskAnn bring-up time so the user gets immediate @@ -104,20 +104,15 @@ int LoadDiskAnnPlugin(const std::string &path) { // the file reader will gracefully fall back to pread(). The user gets // a clear warning now rather than a hard failure later. g_runtime_ready.store(true, std::memory_order_release); - return kDiskAnnPluginLibAioMissing; + return kDiskAnnRuntimeLibAioMissing; } LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); g_runtime_ready.store(true, std::memory_order_release); - return kDiskAnnPluginOk; + return kDiskAnnRuntimeOk; #else - return kDiskAnnPluginUnsupportedPlatform; + return kDiskAnnRuntimeUnsupportedPlatform; #endif } -bool UnloadDiskAnnPlugin() { - // DiskAnn is statically linked — there is nothing to unload. - return false; -} - } // namespace zvec diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 1ec54f265..36f0c20fb 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -23,7 +23,7 @@ namespace zvec::core_interface { int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { // Platform and libaio checks are handled earlier at schema validation - // time (schema.cc calls LoadDiskAnnPlugin). If libaio is missing, DiskAnn + // time (schema.cc calls InitDiskAnnRuntime). If libaio is missing, DiskAnn // falls back to synchronous pread() in diskann_file_reader.cc — no need // to re-check here. if (is_sparse_) { diff --git a/src/core/plugin/CMakeLists.txt b/src/core/plugin/CMakeLists.txt deleted file mode 100644 index 808e35f57..000000000 --- a/src/core/plugin/CMakeLists.txt +++ /dev/null @@ -1,29 +0,0 @@ -## -## Copyright (C) The Software Authors. All rights reserved. -## -## \file CMakeLists.txt -## \brief Build script for the zvec internal DiskAnn runtime-init glue. -## Lives inside the main zvec_core artifact and provides the -## implicit DiskAnn runtime bring-up used by DiskAnnIndex on first -## use. These APIs are NOT part of the public user surface: users -## simply instantiate ``DiskAnnIndexParam`` / ``DiskAnnIndex`` and -## libaio is loaded behind the scenes via dlopen(). DiskAnn itself -## is compiled directly into the hosting binary. -## - -include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) -include(${PROJECT_ROOT_DIR}/cmake/option.cmake) - -set(CORE_PLUGIN_LIBS zvec_ailego core_framework) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux") - list(APPEND CORE_PLUGIN_LIBS ${CMAKE_DL_LIBS}) -endif() - -cc_library( - NAME core_plugin - STATIC STRICT ALWAYS_LINK - SRCS *.cc - LIBS ${CORE_PLUGIN_LIBS} - INCS . ${PROJECT_ROOT_DIR}/src/core - VERSION "${PROXIMA_ZVEC_VERSION}" -) diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index bf5b02d38..4eb673a68 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -16,11 +16,11 @@ #include #include #include +#include #include #include #include #include -#include #include "ailego/internal/cpu_features.h" #include "db/common/constants.h" #include "db/common/typedef.h" @@ -180,14 +180,14 @@ Status FieldSchema::validate() const { // are creation-time only, so triggering the runtime init here is // safe (and idempotent/cached). // - // kDiskAnnPluginLibAioMissing is non-fatal: DiskAnn falls back to + // kDiskAnnRuntimeLibAioMissing is non-fatal: DiskAnn falls back to // synchronous pread() with degraded performance. - const int rc = ::zvec::LoadDiskAnnPlugin(); + const int rc = ::zvec::InitDiskAnnRuntime(); switch (rc) { - case kDiskAnnPluginOk: - case kDiskAnnPluginLibAioMissing: + case kDiskAnnRuntimeOk: + case kDiskAnnRuntimeLibAioMissing: break; - case kDiskAnnPluginUnsupportedPlatform: + case kDiskAnnRuntimeUnsupportedPlatform: return Status::NotSupported( "DiskAnn is not supported on this platform (Linux x86_64 " "only)"); diff --git a/src/include/zvec/plugin/diskann_plugin.h b/src/include/zvec/core/interface/diskann_runtime.h similarity index 61% rename from src/include/zvec/plugin/diskann_plugin.h rename to src/include/zvec/core/interface/diskann_runtime.h index ef1437692..657e12061 100644 --- a/src/include/zvec/plugin/diskann_plugin.h +++ b/src/include/zvec/core/interface/diskann_runtime.h @@ -18,67 +18,47 @@ // NOTE: The APIs declared in this header are INTERNAL to zvec. They are // invoked implicitly by ``DiskAnnIndex`` on first use so that DiskAnn works -// out of the box, without users ever calling a ``load_diskann_plugin()`` / +// out of the box, without users ever calling a ``init_diskann_runtime()`` / // ``is_libaio_available()`` entry point. External callers should not depend // on these symbols; they may change or be removed in future releases. // // DiskAnn is compiled directly into the hosting binary (_zvec.so for the // Python wheel, the test executable for gtest, or the tool binary for C++ -// tools). ``LoadDiskAnnPlugin`` no longer dlopens a separate .so — it simply +// tools). ``InitDiskAnnRuntime`` does not dlopen a separate .so — it simply // loads libaio eagerly via dlopen()/dlsym() and logs the result. On hosts // missing libaio, DiskAnn falls back to synchronous pread() (with a warning) // while other index types (HNSW / IVF / Flat / Vamana) keep working. -#if defined(_WIN32) || defined(__CYGWIN__) -#ifdef ZVEC_BUILD_SHARED -#define ZVEC_PLUGIN_EXPORT __declspec(dllexport) -#elif defined(ZVEC_USE_SHARED) -#define ZVEC_PLUGIN_EXPORT __declspec(dllimport) -#else -#define ZVEC_PLUGIN_EXPORT -#endif -#else -#if __GNUC__ >= 4 -#define ZVEC_PLUGIN_EXPORT __attribute__((visibility("default"))) -#else -#define ZVEC_PLUGIN_EXPORT -#endif -#endif - namespace zvec { -// Return codes for LoadDiskAnnPlugin(). -enum DiskAnnPluginStatus { - kDiskAnnPluginOk = 0, - kDiskAnnPluginUnsupportedPlatform = -1, - kDiskAnnPluginLibAioMissing = -2, - kDiskAnnPluginDlopenFailed = -3, +// Return codes for InitDiskAnnRuntime(). +enum DiskAnnRuntimeStatus { + kDiskAnnRuntimeOk = 0, + kDiskAnnRuntimeUnsupportedPlatform = -1, + kDiskAnnRuntimeLibAioMissing = -2, + kDiskAnnRuntimeDlopenFailed = -3, }; // Returns true if libaio is present on the host and the minimum set of symbols // required by the DiskAnn runtime (io_setup / io_submit / io_getevents / // io_destroy) can be resolved at runtime. // -// Internal probe used by ``LoadDiskAnnPlugin`` before attempting dlopen. Not +// Internal probe used by ``InitDiskAnnRuntime`` before attempting dlopen. Not // part of the user-facing API. -ZVEC_PLUGIN_EXPORT bool IsLibAioAvailable(); +bool IsLibAioAvailable(); // Load libaio at runtime via dlopen()/dlsym() and mark the DiskAnn runtime // as ready. DiskAnn code is compiled directly into the hosting binary, so // no separate .so is loaded. Invoked implicitly by // ``DiskAnnIndex::CreateAndInitStreamer`` on first use; callers should not // invoke it directly. The call is idempotent and returns -// ``kDiskAnnPluginOk`` when the runtime is already active. +// ``kDiskAnnRuntimeOk`` when the runtime is already active. // // The ``path`` parameter is retained for API compatibility but is ignored. -ZVEC_PLUGIN_EXPORT int LoadDiskAnnPlugin(const std::string &path = ""); +int InitDiskAnnRuntime(const std::string &path = ""); -// Returns true if the DiskAnn runtime is currently loaded in this process. +// Returns true if the DiskAnn runtime has been initialized in this process. // Internal diagnostic; not a user-facing API. -ZVEC_PLUGIN_EXPORT bool IsDiskAnnPluginLoaded(); - -// Unload the DiskAnn runtime. Since DiskAnn is statically linked, this is -// a no-op that always returns false. Retained for API compatibility. -ZVEC_PLUGIN_EXPORT bool UnloadDiskAnnPlugin(); +bool IsDiskAnnRuntimeReady(); } // namespace zvec diff --git a/tests/core/algorithm/diskann/CMakeLists.txt b/tests/core/algorithm/diskann/CMakeLists.txt index 2d919f270..e6ad1af12 100644 --- a/tests/core/algorithm/diskann/CMakeLists.txt +++ b/tests/core/algorithm/diskann/CMakeLists.txt @@ -7,7 +7,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) cc_gtest( NAME ${CC_TARGET} STRICT - LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_plugin core_knn_diskann + LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_knn_diskann SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/diskann ) diff --git a/tests/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 108123b79..62d40b474 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -144,8 +144,8 @@ TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) { } // DiskAnn is now exposed implicitly: no caller ever invokes a -// ``LoadDiskAnnPlugin`` / ``IsLibAioAvailable`` API (those were removed from -// the public surface together with ``zvec.load_diskann_plugin()`` in Python). +// ``InitDiskAnnRuntime`` / ``IsLibAioAvailable`` API (those were removed from +// the public surface together with ``zvec.init_diskann_runtime()`` in Python). // The only contract this test validates is the UX guarantee: once the DiskAnn // module has been linked into the hosting binary (here, directly into the // test via the ``core_knn_diskann`` target), its factory entries are From 081aee33810be5849d360b1e846292d3e7b5facd Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 3 Jul 2026 16:21:07 +0800 Subject: [PATCH 14/46] fix: add ci job to test libaio --- .github/workflows/01-ci-pipeline.yml | 8 ++ .github/workflows/09-diskann-libaio-test.yml | 109 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 .github/workflows/09-diskann-libaio-test.yml diff --git a/.github/workflows/01-ci-pipeline.yml b/.github/workflows/01-ci-pipeline.yml index ce9f8e2ef..b0c3c7e27 100644 --- a/.github/workflows/01-ci-pipeline.yml +++ b/.github/workflows/01-ci-pipeline.yml @@ -97,6 +97,14 @@ jobs: os: ubuntu-24.04 compiler: clang + # DiskAnn w/ libaio: explicitly install libaio and run tests to exercise + # the async I/O path. The existing CI (without libaio-dev) already covers + # the w/o libaio pread() fallback path. + diskann-libaio-test: + name: DiskAnn w/ libaio (linux-x64) + needs: [lint, clang-tidy] + uses: ./.github/workflows/09-diskann-libaio-test.yml + build-android: if: github.event_name != 'push' || github.ref != 'refs/heads/main' name: Build & Test (android) diff --git a/.github/workflows/09-diskann-libaio-test.yml b/.github/workflows/09-diskann-libaio-test.yml new file mode 100644 index 000000000..1024c9ffb --- /dev/null +++ b/.github/workflows/09-diskann-libaio-test.yml @@ -0,0 +1,109 @@ +name: DiskAnn libaio Test + +# Tests DiskAnn on Linux x86_64 with libaio explicitly installed, exercising +# the async I/O path (io_setup/io_submit/io_getevents). +# +# The existing CI (03-macos-linux-build.yml without libaio-dev installed) and +# non-Linux platform jobs already cover the "w/o libaio" fallback path — +# DiskAnn gracefully degrades to synchronous pread() when libaio is absent. +# This workflow complements that by ensuring the async I/O path is also +# exercised when libaio is present. + +on: + workflow_call: + workflow_dispatch: + +permissions: + contents: read + +jobs: + libaio-test: + name: DiskAnn w/ libaio (linux-x64) + runs-on: ubuntu-24.04 + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + submodules: recursive + + - name: Install libaio runtime + run: | + sudo apt-get update -y + # libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older + sudo apt-get install -y libaio1t64 || sudo apt-get install -y libaio1 + echo "=== libaio status (should be present) ===" + dpkg -l | grep -i libaio || true + ldconfig -p | grep libaio || true + shell: bash + + - name: Setup ccache + uses: hendrikmuhs/ccache-action@v1.2 + with: + key: linux-x64-libaio-test + max-size: 150M + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: 'pyproject.toml' + + - name: Set up environment variables + run: | + NPROC=$(nproc 2>/dev/null || echo 2) + echo "NPROC=$NPROC" >> $GITHUB_ENV + echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH + shell: bash + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install \ + pybind11==3.0 \ + cmake==3.30.0 \ + ninja==1.11.1 \ + pytest \ + pytest-xdist \ + scikit-build-core \ + setuptools_scm + shell: bash + + - name: Build from source + run: | + cd "$GITHUB_WORKSPACE" + CMAKE_GENERATOR="Ninja" \ + CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ + python -m pip install -v . \ + --no-build-isolation \ + --config-settings='cmake.define.BUILD_TOOLS=ON' \ + --config-settings='cmake.define.ENABLE_WERROR=ON' \ + --config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \ + --config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache' + shell: bash + + - name: Run C++ Tests (w/ libaio) + run: | + cd "$GITHUB_WORKSPACE/build" + cmake --build . --target unittest --parallel $NPROC + shell: bash + + - name: Verify no libaio in test binary + run: | + # Confirm the test binary does NOT have a NEEDED entry for libaio. + # This verifies the dlopen decoupling — libaio is loaded at runtime, + # not linked at build time. + TEST_BIN="$GITHUB_WORKSPACE/build/bin/diskann_searcher_test" + if [ -f "$TEST_BIN" ]; then + ldd "$TEST_BIN" | grep 'libaio' \ + && { echo "ERROR: test binary links libaio directly!"; exit 1; } \ + || echo "OK: test binary does not link libaio (dlopen-based)" + fi + shell: bash + + - name: Run Python Tests (w/ libaio) + run: | + cd "$GITHUB_WORKSPACE" + python -m pytest python/tests/ + shell: bash From 058c67c556b4912201975af5071d417115322fe6 Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 6 Jul 2026 16:15:16 +0800 Subject: [PATCH 15/46] feat: add diskann on arm --- CMakeLists.txt | 19 +- python/tests/detail/fixture_helper.py | 20 +- python/tests/test_collection_diskann.py | 16 +- src/core/CMakeLists.txt | 10 +- src/core/algorithm/CMakeLists.txt | 2 +- src/core/algorithm/diskann/CMakeLists.txt | 19 +- src/core/algorithm/diskann/diskann_context.h | 8 +- .../algorithm/diskann/diskann_file_reader.cc | 179 +++++++++++++++++- .../algorithm/diskann/diskann_file_reader.h | 21 ++ src/core/algorithm/diskann/diskann_indexer.h | 8 +- src/core/algorithm/diskann/diskann_util.h | 8 +- src/core/interface/indexes/diskann_index.cc | 11 +- src/core/plugin/diskann_plugin.cc | 70 +++++-- src/db/index/common/schema.cc | 8 +- src/include/zvec/plugin/diskann_plugin.h | 5 +- 15 files changed, 336 insertions(+), 68 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index d334b857f..b417af403 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -122,14 +122,17 @@ else() endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") -# DiskAnn support (Linux x86_64 only, requires libaio) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) +# DiskAnn support: +# - Linux (x86_64, i686, i386, aarch64, arm64) with libaio +# - macOS (x86_64, ARM64/Apple Silicon) with kqueue +if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS) + OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND NOT IOS)) set(DISKANN_SUPPORTED ON) add_definitions(-DDISKANN_SUPPORTED=1) else() set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled - only supported on Linux x86_64") + message(STATUS "DiskAnn support disabled - supported on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") @@ -189,11 +192,11 @@ if(BUILD_PYTHON_BINDINGS) # zvec package directory as well. # # Gate on DISKANN_SUPPORTED, not on the target's existence: on unsupported - # platforms (e.g. macOS / ARM64) the core_knn_diskann target is still - # defined, but built from an empty stub (src/core/algorithm/CMakeLists.txt) - # with zero exported symbols and a runtime load path compiled out - # (#if DISKANN_SUPPORTED). Shipping that stub is pure dead weight, so it is - # only packaged where DiskAnn is real — currently Linux x86_64 with libaio. + # platforms the core_knn_diskann target is still defined, but built from + # an empty stub (src/core/algorithm/CMakeLists.txt) with zero exported + # symbols and a runtime load path compiled out (#if DISKANN_SUPPORTED). + # Shipping that stub is pure dead weight, so it is only packaged where + # DiskAnn is real — currently Linux (x86_64/ARM64) with libaio and macOS with kqueue. if(DISKANN_SUPPORTED) install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec COMPONENT python) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index 63aac799c..9f33317be 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,12 +2,10 @@ import logging import platform -DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( - "x86_64", - "AMD64", - "i686", - "i386", -) +DISKANN_SUPPORTED = ( + platform.system() == "Linux" + and platform.machine() in ("x86_64", "AMD64", "i686", "i386", "aarch64", "arm64") +) or platform.system() == "Darwin" from typing import Any, Generator from zvec.typing import DataType, StatusCode, MetricType, QuantizeType @@ -32,10 +30,12 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_DONE = True if not DISKANN_SUPPORTED: - _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" + _DISKANN_PRELOAD_REASON = "DiskAnn is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)" return _DISKANN_PRELOAD_REASON - if not zvec.is_libaio_available(): + # On Linux, verify libaio is available. On macOS, kqueue is always + # available as part of the system. + if platform.system() == "Linux" and not zvec.is_libaio_available(): _DISKANN_PRELOAD_REASON = ( "libaio is not available on this host; DiskAnn cannot run. " "Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry." @@ -46,8 +46,8 @@ def _ensure_diskann_runtime_or_reason() -> str | None: if status != zvec.DISKANN_PLUGIN_OK: _DISKANN_PRELOAD_REASON = ( f"Failed to load DiskAnn plugin (status={status}); " - "check that libzvec_diskann_plugin.so is installed alongside " - "_zvec.so in the Python site-packages directory." + "check that the DiskAnn plugin shared library is installed " + "alongside _zvec.so in the Python site-packages directory." ) return _DISKANN_PRELOAD_REASON diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index b0e12ce7e..f372a14a2 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -17,8 +17,8 @@ Two platform-level prerequisites are enforced at module import time: -1. DiskAnn is currently built only for Linux x86_64 — other platforms are - skipped wholesale. +1. DiskAnn is built for Linux (x86_64/ARM64 with libaio) and macOS (with kqueue) — + other platforms are skipped wholesale. 2. The DiskAnn backend lives in a *runtime-loaded* plugin (``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL | RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory`` @@ -42,15 +42,21 @@ # Platform gating (must happen BEFORE we touch zvec). # --------------------------------------------------------------------------- # pytestmark = pytest.mark.skipif( - not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), - reason="DiskAnn plugin is only supported on Linux x86_64", + not ( + ( + sys.platform == "linux" + and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64") + ) + or sys.platform == "darwin" + ), + reason="DiskAnn plugin is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)", ) # Promote all symbols in subsequently-loaded DSOs to the global namespace and # resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see # the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa. # See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement. -if sys.platform == "linux": +if sys.platform in ("linux", "darwin"): sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW) import zvec # noqa: E402 diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index c52c3b8cc..6a5b3c7ff 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -63,9 +63,9 @@ endif() # Always exclude algorithm/diskann implementation files from zvec_core. # The DiskAnn algorithm is provided by the separate core_knn_diskann library -# (real on Linux x86_64, stub on other platforms). Including them here causes -# duplicate symbols and missing -laio when test binaries link both zvec_core -# (via zvec) and core_knn_diskann. +# (real on Linux x86_64/ARM64 with libaio and macOS with kqueue, stub on other +# platforms). Including them here causes duplicate symbols and missing +# -laio when test binaries link both zvec_core (via zvec) and core_knn_diskann. list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*") if(NOT DISKANN_SUPPORTED) @@ -73,8 +73,8 @@ if(NOT DISKANN_SUPPORTED) endif() set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) -# The plugin loader uses dlopen/dlsym, so link libdl on Linux. -if(CMAKE_SYSTEM_NAME STREQUAL "Linux") +# The plugin loader uses dlopen/dlsym, so link libdl on Unix platforms. +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" OR CMAKE_SYSTEM_NAME STREQUAL "Darwin") list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) endif() diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index bed772984..db898d248 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -17,7 +17,7 @@ else() # Empty stub library for unsupported platforms file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc "// Stub implementation for unsupported platforms\n" - "// DiskAnn only supports Linux x86_64\n" + "// DiskAnn supports Linux (x86_64/ARM64 with libaio) and macOS (kqueue)\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 1b2252497..47a68fbe9 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -14,12 +14,15 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # (which only ships _zvec.so + the plugin) would fail to load. # # We therefore link the plugin only against its system-level dependency -# (libaio) and rely on the global include dirs plus -Wl,--unresolved-symbols -# =ignore-all to let the linker build a shared library with unresolved -# references to internal APIs. +# (libaio on Linux x86_64/ARM64) and rely on the global include dirs plus +# linker flags to let the linker build a shared library with unresolved +# references to internal APIs. On macOS, kqueue is part of the system +# libraries so no extra dependency is needed. set(CORE_KNN_DISKANN_LIBS "") -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") +# libaio is available on all mainstream Linux architectures (x86_64, aarch64, +# arm64, ppc64le, ...). Link it on Linux regardless of CPU architecture. +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT ANDROID) list(APPEND CORE_KNN_DISKANN_LIBS aio) endif() @@ -48,10 +51,14 @@ endforeach() if(CMAKE_SYSTEM_NAME STREQUAL "Linux") target_link_options(core_knn_diskann PRIVATE "LINKER:--unresolved-symbols=ignore-all") +elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + target_link_options(core_knn_diskann PRIVATE + "LINKER:-undefined,dynamic_lookup") endif() -# Rename the artifact to libzvec_diskann_plugin.so so it is clearly identified -# as an optional plugin that users load at runtime via zvec::LoadDiskAnnPlugin(). +# Rename the artifact so it is clearly identified as an optional plugin that +# users load at runtime via zvec::LoadDiskAnnPlugin(). The extension is +# platform-specific (.so on Linux, .dylib on macOS). set_target_properties(core_knn_diskann PROPERTIES OUTPUT_NAME zvec_diskann_plugin ) \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index f8a736cdf..6e0401f6a 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -341,7 +341,13 @@ class DiskAnnContext : public IndexContext, uint32_t group_num_{0}; std::map group_topk_heaps_{}; - IOContext io_ctx_{0}; + IOContext io_ctx_{ +#if defined(__APPLE__) || defined(__MACH__) + -1 +#else + 0 +#endif + }; SearchStats query_stats_; float *pq_table_dist_buffer_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 4268d5156..20592ac98 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -33,8 +33,18 @@ typedef struct iocb iocb_t; int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) int ret = io_setup(MAX_EVENTS, &ctx); - return ret; +#elif defined(__APPLE__) || defined(__MACH__) + // Create a kqueue for this context. On macOS the kqueue is used to + // monitor file descriptor readiness for async-style I/O. + int kq = ::kqueue(); + if (kq == -1) { + LOG_ERROR("kqueue() failed in setup_io_ctx; errno=%d, %s", errno, + ::strerror(errno)); + return IndexError_Runtime; + } + ctx = kq; + return 0; #else return 0; #endif @@ -43,8 +53,13 @@ int setup_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) int ret = io_destroy(ctx); - return ret; +#elif defined(__APPLE__) || defined(__MACH__) + if (ctx >= 0) { + ::close(ctx); + ctx = -1; + } + return 0; #else return 0; #endif @@ -68,6 +83,110 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } +#if defined(__APPLE__) || defined(__MACH__) +// Execute batch I/O on macOS using kqueue to monitor file descriptor +// readiness and pread for actual data transfer. +// +// On macOS, regular file descriptors are almost always "readable", so +// kqueue's primary value here is providing the same async I/O interface +// as Linux's libaio. For each read request we: +// 1. Attempt a non-blocking pread (via O_NONBLOCK on the fd). +// 2. If EAGAIN, wait on kqueue for EVFILT_READ readiness, then retry. +// 3. Fall back to blocking pread if kqueue encounters an error. +// +// The kqueue fd is passed in as the IOContext. If no valid kqueue is +// available, we fall back to plain blocking pread. +static int execute_io_kqueue(int kq, int fd, + std::vector &read_reqs) { + // If no kqueue available, fall back to blocking pread. + if (kq < 0) { + return execute_io_pread(fd, read_reqs); + } + + // Register the file descriptor with the kqueue for read events. + // EV_CLEAR gives edge-triggered semantics so we only get notified + // when new data becomes available. + struct kevent ke; + EV_SET(&ke, fd, EVFILT_READ, EV_ADD | EV_CLEAR, 0, 0, nullptr); + + for (auto &req : read_reqs) { + while (true) { + ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); + + if (bytes_read > 0) { + // Successfully read data; verify full read. + if ((size_t)bytes_read != req.len) { + // Partial read — retry for the remaining bytes. + // Update offset and buffer to read the rest. + char *buf_ptr = static_cast(req.buf) + bytes_read; + uint64_t new_offset = req.offset + bytes_read; + size_t remaining = req.len - bytes_read; + while (remaining > 0) { + ssize_t n = ::pread(fd, buf_ptr, remaining, new_offset); + if (n < 0) { + if (errno == EINTR) continue; + LOG_ERROR("pread retry failed; errno=%d, %s", errno, + ::strerror(errno)); + return IndexError_Runtime; + } + if (n == 0) break; + buf_ptr += n; + new_offset += n; + remaining -= n; + } + if (remaining > 0) { + LOG_ERROR("pread short read after retry; remaining=%zu", remaining); + return IndexError_Runtime; + } + } + break; // Success, move to next request. + } + + if (bytes_read == 0) { + // EOF — should not happen for a valid index file. + LOG_ERROR("pread returned 0 (EOF); offset=%lu, len=%lu", + (unsigned long)req.offset, (unsigned long)req.len); + return IndexError_Runtime; + } + + // bytes_read == -1, error + if (errno == EINTR) { + continue; // Retry on signal. + } + if (errno == EAGAIN || errno == EWOULDBLOCK) { + // Data not ready — wait on kqueue for readability. + struct kevent events[1]; + struct timespec ts; + ts.tv_sec = 5; // 5 second timeout as a safety net. + ts.tv_nsec = 0; + int n_ev = ::kevent(kq, &ke, 1, events, 1, &ts); + if (n_ev < 0) { + if (errno == EINTR) continue; + // kqueue error — fall back to blocking pread for this request. + LOG_WARN("kevent failed; errno=%d, %s, falling back to pread", errno, + ::strerror(errno)); + return execute_io_pread(fd, read_reqs); + } + if (n_ev == 0) { + // Timeout — fall back to blocking pread. + LOG_WARN("kqueue timeout, falling back to pread"); + return execute_io_pread(fd, read_reqs); + } + // Event triggered — retry pread. + continue; + } + + // Other error — fall back to blocking pread. + LOG_ERROR("pread failed; errno=%d, %s, falling back to pread", errno, + ::strerror(errno)); + return execute_io_pread(fd, read_reqs); + } + } + + return 0; +} +#endif // __APPLE__ + int execute_io(IOContext ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) @@ -144,7 +263,13 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, } return 0; +#elif defined(__APPLE__) || defined(__MACH__) + // On macOS, use kqueue-based I/O. The IOContext (ctx) is a kqueue fd. + (void)n_retries; + return execute_io_kqueue(ctx, fd, read_reqs); #else + (void)ctx; + (void)n_retries; return execute_io_pread(fd, read_reqs); #endif } @@ -182,7 +307,6 @@ void LinuxAlignedFileReader::register_thread() { std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { LOG_ERROR("multiple calls to register_thread from the same thread"); - return; } @@ -197,14 +321,30 @@ void LinuxAlignedFileReader::register_thread() { "/proc/sys/fs/aio-max-nr"); } else { LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret)); - ; } } else { LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); - ctx_map[thread_id] = ctx; } + lk.unlock(); +#elif defined(__APPLE__) || defined(__MACH__) + auto thread_id = std::this_thread::get_id(); + std::unique_lock 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 = -1; + int kq = ::kqueue(); + if (kq == -1) { + LOG_ERROR("kqueue() failed in register_thread; errno=%d, %s", errno, + ::strerror(errno)); + } else { + LOG_INFO("allocating kqueue ctx: %d", kq); + ctx = kq; + ctx_map[thread_id] = ctx; + } lk.unlock(); #endif } @@ -228,6 +368,25 @@ void LinuxAlignedFileReader::deregister_thread() { // io_destroy is a syscall; keep it outside the lock to avoid blocking others io_destroy(ctx); LOG_INFO("returned ctx from thread"); +#elif defined(__APPLE__) || defined(__MACH__) + auto thread_id = std::this_thread::get_id(); + IOContext ctx; + + { + std::lock_guard lk(ctx_mut); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("deregister_thread: thread not registered"); + return; + } + ctx = it->second; + ctx_map.erase(it); + } + + if (ctx >= 0) { + ::close(ctx); + } + LOG_INFO("returned kqueue ctx from thread"); #endif } @@ -239,6 +398,15 @@ void LinuxAlignedFileReader::deregister_all_threads() { io_destroy(ctx); } ctx_map.clear(); +#elif defined(__APPLE__) || defined(__MACH__) + std::unique_lock lk(ctx_mut); + for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { + IOContext ctx = x->second; + if (ctx >= 0) { + ::close(ctx); + } + } + ctx_map.clear(); #endif } @@ -294,6 +462,5 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } - } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 432247e79..a085a43f5 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -21,8 +21,17 @@ #include #endif +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#include +#endif + #include #include +#include +#include +#include #include #include #include "diskann_util.h" @@ -30,8 +39,13 @@ namespace zvec { namespace core { +// On Linux, IOContext is the libaio io_context_t. +// On macOS, IOContext is an int holding a kqueue file descriptor. +// On other platforms, IOContext is a uint32_t placeholder. #if (defined(__linux) || defined(__linux__)) typedef io_context_t IOContext; +#elif defined(__APPLE__) || defined(__MACH__) +typedef int IOContext; #else typedef uint32_t IOContext; #endif @@ -48,9 +62,12 @@ struct AlignedRead { AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) { +#if defined(__linux__) || defined(__linux) + // O_DIRECT requires 512-byte alignment on Linux. ailego_assert(static_cast(offset) % 512 == 0); ailego_assert(static_cast(len) % 512 == 0); ailego_assert(reinterpret_cast(buf) % 512 == 0); +#endif } }; @@ -75,6 +92,10 @@ class AlignedFileReader { bool async = false) = 0; }; +// Reader implementation used on all supported platforms. +// On Linux (x86_64 and ARM64) it uses libaio for asynchronous batch I/O. +// On macOS (including ARM/Apple Silicon) it uses kqueue to monitor file +// descriptor readiness and pread for actual data transfer. class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index c372d288f..5eea9c7c3 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -91,7 +91,13 @@ class DiskAnnIndexer { PQTable::Pointer pq_table_; - IOContext init_ctx_{0}; + IOContext init_ctx_{ +#if defined(__APPLE__) || defined(__MACH__) + -1 +#else + 0 +#endif + }; std::vector neighbor_cache_buffer_; void *coord_cache_buf_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index a02130bf0..3e2ecd652 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -35,7 +35,13 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { - *ptr = ::aligned_alloc(align, size); + // C11 aligned_alloc() requires size to be an integral multiple of + // alignment. This is true on Linux (glibc relaxes the requirement) + // but NOT on macOS, where aligned_alloc(32, 16) returns NULL. + // Use posix_memalign() which is portable and has no such restriction. + if (::posix_memalign(ptr, align, size) != 0) { + *ptr = nullptr; + } } static inline void free_aligned(void *ptr) { diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 7615ed354..59a47de2e 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -27,9 +27,10 @@ namespace { // Implicitly bring the DiskAnn runtime online on first use. This keeps the // DiskAnn index an ordinary public API (users just instantiate a // DiskAnnIndexParam) while still letting the rest of the library — HNSW, -// IVF, Flat, Vamana — run on hosts that happen to lack libaio. On such -// hosts only DiskAnn fails, with a clear, actionable error message, and -// every other index type stays fully functional. +// IVF, Flat, Vamana — run on hosts that happen to lack the DiskAnn runtime +// dependency (libaio on Linux). On such hosts only DiskAnn fails, with a +// clear, actionable error message, and every other index type stays fully +// functional. int EnsureDiskAnnRuntimeReady() { static std::once_flag once; static int cached_result = 0; @@ -48,7 +49,9 @@ int EnsureDiskAnnRuntimeReady() { "or 'libaio1t64' on Ubuntu 24.04+) and retry."); break; case kDiskAnnPluginUnsupportedPlatform: - LOG_ERROR("DiskAnn is only supported on Linux x86_64."); + LOG_ERROR( + "DiskAnn is not supported on this platform. It is available on " + "Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)."); break; case kDiskAnnPluginDlopenFailed: default: diff --git a/src/core/plugin/diskann_plugin.cc b/src/core/plugin/diskann_plugin.cc index 9d4e9d0f5..789eaab84 100644 --- a/src/core/plugin/diskann_plugin.cc +++ b/src/core/plugin/diskann_plugin.cc @@ -28,6 +28,11 @@ #include #endif +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#endif + namespace zvec { namespace { @@ -42,9 +47,10 @@ constexpr const char *kLibAioSoNames[] = { }; constexpr bool kPlatformSupportsDiskAnnPlugin = true; #elif defined(__APPLE__) -[[maybe_unused]] constexpr const char *kPluginFileName = - "libzvec_diskann_plugin.dylib"; -constexpr bool kPlatformSupportsDiskAnnPlugin = false; +constexpr const char *kPluginFileName = "libzvec_diskann_plugin.dylib"; +// On macOS, DiskAnn uses kqueue (part of the system) instead of libaio, +// so there is no external runtime dependency to probe. +constexpr bool kPlatformSupportsDiskAnnPlugin = true; #else [[maybe_unused]] constexpr const char *kPluginFileName = "zvec_diskann_plugin.dll"; @@ -55,17 +61,31 @@ constexpr bool kPlatformSupportsDiskAnnPlugin = false; std::atomic g_plugin_handle{nullptr}; std::mutex g_plugin_mutex; -#if defined(__linux__) || defined(__linux) +#if (defined(__linux__) || defined(__linux)) || \ + (defined(__APPLE__) || defined(__MACH__)) // Resolve the directory containing the currently running executable, so we // can look for the plugin next to it regardless of the working directory. std::string GetExecutableDir() { char buf[PATH_MAX]; +#if defined(__linux__) || defined(__linux) ssize_t n = ::readlink("/proc/self/exe", buf, sizeof(buf) - 1); if (n <= 0) { return {}; } buf[n] = '\0'; +#elif defined(__APPLE__) + uint32_t bufsize = sizeof(buf); + if (_NSGetExecutablePath(buf, &bufsize) != 0) { + return {}; + } + // _NSGetExecutablePath may return a path that's not absolute; resolve it. + char resolved[PATH_MAX]; + if (::realpath(buf, resolved) != nullptr) { + ::strncpy(buf, resolved, sizeof(buf) - 1); + buf[sizeof(buf) - 1] = '\0'; + } +#endif std::string path(buf); auto slash = path.find_last_of('/'); if (slash == std::string::npos) { @@ -76,9 +96,10 @@ std::string GetExecutableDir() { // Resolve the directory containing the shared object that hosts this // function. For Python wheels this is the directory of -// ``_zvec.cpython-*.so``; for regular C++ binaries it is the directory of -// ``libzvec_core.so``. In either case the DiskAnn plugin is shipped -// alongside, so this is the most reliable lookup location. +// ``_zvec.cpython-*.so`` (Linux) or ``_zvec.cpython-*.so`` (macOS); for +// regular C++ binaries it is the directory of ``libzvec_core.so`` (Linux) +// or ``libzvec_core.dylib`` (macOS). In either case the DiskAnn plugin is +// shipped alongside, so this is the most reliable lookup location. // // NOTE: we pass the address of an *exported* function (LoadDiskAnnPlugin) // rather than one in this anonymous namespace, because dladdr() on a symbol @@ -136,9 +157,19 @@ void PromoteHostingSoToGlobal() { const std::string exe_path = GetExecutableDir(); if (!exe_path.empty()) { char exe_buf[PATH_MAX]; +#if defined(__linux__) || defined(__linux) ssize_t n = ::readlink("/proc/self/exe", exe_buf, sizeof(exe_buf) - 1); if (n > 0) { exe_buf[n] = '\0'; +#elif defined(__APPLE__) + uint32_t bufsize = sizeof(exe_buf); + if (_NSGetExecutablePath(exe_buf, &bufsize) == 0) { + char resolved[PATH_MAX]; + if (::realpath(exe_buf, resolved) != nullptr) { + ::strncpy(exe_buf, resolved, sizeof(exe_buf) - 1); + exe_buf[sizeof(exe_buf) - 1] = '\0'; + } +#endif // Compare resolved real paths to handle relative vs absolute. char host_real[PATH_MAX]; char exe_real[PATH_MAX]; @@ -196,12 +227,12 @@ std::vector BuildCandidatePaths(const std::string &explicit_path) { push_dir_candidates(exe_dir); } // 3. Fallback: rely on the dynamic linker's default search path - // (RPATH / LD_LIBRARY_PATH / /etc/ld.so.conf). + // (RPATH / LD_LIBRARY_PATH / DYLD_LIBRARY_PATH / ld.so.conf). candidates.emplace_back(kPluginFileName); return candidates; } -#endif // linux +#endif // linux || apple } // namespace @@ -230,7 +261,9 @@ bool IsLibAioAvailable() { } return false; #else - return false; + // On macOS, DiskAnn uses kqueue which is part of the system libraries, + // so there is no external dependency to probe — always "available". + return true; #endif } @@ -241,12 +274,13 @@ bool IsDiskAnnPluginLoaded() { int LoadDiskAnnPlugin(const std::string &path) { if (!kPlatformSupportsDiskAnnPlugin) { LOG_ERROR( - "DiskAnn plugin is not supported on this platform; it is only " - "available on Linux x86_64 with libaio."); + "DiskAnn plugin is not supported on this platform; it is available " + "on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)."); return kDiskAnnPluginUnsupportedPlatform; } -#if defined(__linux__) || defined(__linux) +#if (defined(__linux__) || defined(__linux)) || \ + (defined(__APPLE__) || defined(__MACH__)) // Fast path: already loaded. if (g_plugin_handle.load(std::memory_order_acquire) != nullptr) { return kDiskAnnPluginOk; @@ -257,6 +291,10 @@ int LoadDiskAnnPlugin(const std::string &path) { return kDiskAnnPluginOk; } +#if defined(__linux__) || defined(__linux) + // On Linux, verify that libaio is available before attempting to load + // the plugin. The plugin links against libaio and will fail to load + // without it. if (!IsLibAioAvailable()) { LOG_ERROR( "libaio is not available on this host; the DiskAnn runtime cannot be " @@ -265,6 +303,9 @@ int LoadDiskAnnPlugin(const std::string &path) { "other index types (HNSW, IVF, Flat, Vamana)."); return kDiskAnnPluginLibAioMissing; } +#endif + // On macOS, no external library check is needed — kqueue is part of the + // system and the plugin has no external link-time dependencies. const std::vector candidates = BuildCandidatePaths(path); // Ensure the hosting module's C++ symbols (zvec::*) are visible to the @@ -302,7 +343,8 @@ int LoadDiskAnnPlugin(const std::string &path) { } bool UnloadDiskAnnPlugin() { -#if defined(__linux__) || defined(__linux) +#if (defined(__linux__) || defined(__linux)) || \ + (defined(__APPLE__) || defined(__MACH__)) std::lock_guard lock(g_plugin_mutex); void *handle = g_plugin_handle.exchange(nullptr, std::memory_order_acq_rel); if (handle == nullptr) { diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 9471eee90..247b0dc98 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -175,8 +175,8 @@ Status FieldSchema::validate() const { if (index_params_->type() == IndexType::DISKANN) { // Probe the DiskAnn runtime eagerly at creation time so unsupported - // platforms (non Linux x86_64), missing libaio, or a missing plugin - // .so fail fast with a clear message instead of surfacing later during + // platforms, missing libaio (on Linux), or a missing plugin .so fail + // fast with a clear message instead of surfacing later during // optimize(). This reuses the same gate DiskAnnIndex applies on first // use (zvec::LoadDiskAnnPlugin, wrapped by EnsureDiskAnnRuntimeReady). // All validate() call sites are creation-time only, so triggering the @@ -187,8 +187,8 @@ Status FieldSchema::validate() const { break; case kDiskAnnPluginUnsupportedPlatform: return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 " - "only)"); + "DiskAnn is not supported on this platform. It is available " + "on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)."); case kDiskAnnPluginLibAioMissing: return Status::NotSupported( "DiskAnn requires libaio at runtime, but it was not found on " diff --git a/src/include/zvec/plugin/diskann_plugin.h b/src/include/zvec/plugin/diskann_plugin.h index 988187068..016220b82 100644 --- a/src/include/zvec/plugin/diskann_plugin.h +++ b/src/include/zvec/plugin/diskann_plugin.h @@ -21,8 +21,9 @@ // out of the box, without users ever calling a ``load_diskann_plugin()`` / // ``is_libaio_available()`` entry point. External callers should not depend // on these symbols; they may change or be removed in future releases. On -// hosts missing libaio the bring-up fails cleanly and only DiskAnn becomes -// unavailable — other index types (HNSW / IVF / Flat / Vamana) keep working. +// hosts missing the runtime dependency (libaio on Linux) the bring-up fails +// cleanly and only DiskAnn becomes unavailable — other index types (HNSW / +// IVF / Flat / Vamana) keep working. #if defined(_WIN32) || defined(__CYGWIN__) #ifdef ZVEC_BUILD_SHARED From ab8901a3d5a24ee3b305da1474727b292aece4ba Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 6 Jul 2026 16:56:14 +0800 Subject: [PATCH 16/46] fix: fix on comment --- .github/workflows/nightly_coverage.yml | 2 +- pyproject.toml | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/nightly_coverage.yml b/.github/workflows/nightly_coverage.yml index 4c9aca50b..16c3e46de 100644 --- a/.github/workflows/nightly_coverage.yml +++ b/.github/workflows/nightly_coverage.yml @@ -49,7 +49,7 @@ jobs: run: | sudo apt-get update sudo apt-get install -y --no-install-recommends \ - lcov + lcov libaio-dev shell: bash - name: Install dependencies diff --git a/pyproject.toml b/pyproject.toml index bce1670d5..6296b3402 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,7 +181,6 @@ test-command = "cd {project} && pytest python/tests -v --tb=short" build-verbosity = 1 [tool.cibuildwheel.linux] -# libaio is loaded at runtime via dlopen (no build-time dependency). archs = ["auto"] environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" } manylinux-x86_64-image = "manylinux_2_28" From 63f5ff7e19d43e74a10df3e35e7babf5dfd2faec Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 6 Jul 2026 17:32:35 +0800 Subject: [PATCH 17/46] fix: remove diskann runtime --- python/tests/detail/fixture_helper.py | 34 +---- python/tests/test_collection_diskann.py | 8 +- python/zvec/__init__.py | 23 ---- src/binding/python/binding.cc | 38 ------ src/core/algorithm/diskann/diskann_builder.cc | 15 +++ .../algorithm/diskann/diskann_searcher.cc | 15 +++ .../algorithm/diskann/diskann_streamer.cc | 16 +++ src/core/interface/diskann_runtime.cc | 118 ------------------ src/core/interface/indexes/diskann_index.cc | 4 - src/db/index/common/schema.cc | 30 ++--- .../zvec/core/interface/diskann_runtime.h | 64 ---------- .../algorithm/diskann/diskann_builder_test.cc | 12 -- 12 files changed, 61 insertions(+), 316 deletions(-) delete mode 100644 src/core/interface/diskann_runtime.cc delete mode 100644 src/include/zvec/core/interface/diskann_runtime.h diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index e37fc307a..f5e6c08e4 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -13,22 +13,14 @@ from zvec.typing import DataType, StatusCode, MetricType, QuantizeType import zvec - -# Cache the DiskAnn runtime preload status so we pay the init cost once per -# test session. The runtime normally auto-loads on first DiskAnn use, but we -# preload it explicitly here so a missing libaio surfaces as a clear pytest -# skip instead of a confusing "Create vector column indexer failed" deep -# inside the collection code path. -# Note: DiskAnn is now compiled directly into _zvec.so — there is no separate -# plugin .so to locate. If libaio is missing, DiskAnn falls back to -# synchronous pread() with degraded performance. _DISKANN_PRELOAD_REASON: str | None = None _DISKANN_PRELOAD_DONE: bool = False def _ensure_diskann_runtime_or_reason() -> str | None: - """Initialize the DiskAnn runtime and return None on success or a - human-readable skip reason on failure. Idempotent across calls.""" + """Check whether DiskAnn is available on this platform and return None + on success or a human-readable skip reason on failure. Idempotent across + calls.""" global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON if _DISKANN_PRELOAD_DONE: return _DISKANN_PRELOAD_REASON @@ -38,23 +30,6 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" return _DISKANN_PRELOAD_REASON - status = zvec.init_diskann_runtime() - if status == zvec.DISKANN_RUNTIME_UNSUPPORTED_PLATFORM: - _DISKANN_PRELOAD_REASON = ( - f"DiskAnn is not supported on this platform (status={status})." - ) - return _DISKANN_PRELOAD_REASON - # DISKANN_RUNTIME_OK and DISKANN_RUNTIME_LIBAIO_MISSING are both acceptable: - # the latter means DiskAnn will use synchronous pread() instead of async I/O. - if status not in ( - zvec.DISKANN_RUNTIME_OK, - zvec.DISKANN_RUNTIME_LIBAIO_MISSING, - ): - _DISKANN_PRELOAD_REASON = ( - f"Failed to initialize DiskAnn runtime (status={status})." - ) - return _DISKANN_PRELOAD_REASON - _DISKANN_PRELOAD_REASON = None return None @@ -152,8 +127,7 @@ def full_schema_new(request) -> CollectionSchema: else: nullable, has_index, vector_index = True, False, HnswIndexParam() - # Skip DiskAnn tests on unsupported platforms or when the runtime cannot - # be brought up (missing libaio, plugin .so not installed, etc.). + # Skip DiskAnn tests on unsupported platforms. from zvec.model.param import DiskAnnIndexParam if isinstance(vector_index, DiskAnnIndexParam): diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index e90751a04..8daf8c770 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -19,11 +19,9 @@ 1. DiskAnn is currently built only for Linux x86_64 — other platforms are skipped wholesale. -2. The DiskAnn runtime is initialized via ``zvec.init_diskann_runtime()``, - which eagerly loads libaio via dlopen(). If libaio is missing, DiskAnn - falls back to synchronous pread() — the tests still run but with degraded - performance. DiskAnn is compiled directly into ``_zvec.so``; there is no - separate plugin .so to locate. +2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() / + DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to + synchronous pread() — the tests still run but with degraded performance. If either prerequisite fails the whole module is skipped so the rest of the test-suite is not affected. diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index cef509077..9e02c7575 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -43,21 +43,6 @@ # Public API — grouped by category # ============================== -# —— DiskAnn runtime —— -# Re-export the runtime management entry points defined by the C++ extension. -# DiskAnn normally auto-initializes on first use; these APIs let tests and -# diagnostic tools force initialization up-front and get a clear error if -# libaio is missing. -from zvec._zvec import ( - DISKANN_RUNTIME_DLOPEN_FAILED, - DISKANN_RUNTIME_LIBAIO_MISSING, - DISKANN_RUNTIME_OK, - DISKANN_RUNTIME_UNSUPPORTED_PLATFORM, - init_diskann_runtime, - is_diskann_runtime_ready, - is_libaio_available, -) - from . import model as model # —— Extensions —— @@ -203,14 +188,6 @@ "StatusCode", # Tools "require_module", - # DiskAnn runtime - "init_diskann_runtime", - "is_diskann_runtime_ready", - "is_libaio_available", - "DISKANN_RUNTIME_OK", - "DISKANN_RUNTIME_UNSUPPORTED_PLATFORM", - "DISKANN_RUNTIME_LIBAIO_MISSING", - "DISKANN_RUNTIME_DLOPEN_FAILED", ] # ============================== diff --git a/src/binding/python/binding.cc b/src/binding/python/binding.cc index 2d6332644..729a54d15 100644 --- a/src/binding/python/binding.cc +++ b/src/binding/python/binding.cc @@ -13,7 +13,6 @@ // limitations under the License. #include -#include #include "python_collection.h" #include "python_config.h" #include "python_doc.h" @@ -24,42 +23,6 @@ namespace zvec { -namespace { - -// Expose DiskAnn runtime management to Python. DiskAnn is compiled directly -// into _zvec.so, so "initializing" just means eagerly dlopen()-ing libaio and -// caching the result. Tests (and diagnostic tooling) use these entry points -// to force the init up-front and get actionable warnings when libaio is -// missing. -void InitializeDiskAnnRuntimeBindings(pybind11::module_ &m) { - m.def( - "init_diskann_runtime", - [](const std::string &path) { return ::zvec::InitDiskAnnRuntime(path); }, - pybind11::arg("path") = std::string(), - "Initialize the DiskAnn runtime by loading libaio via dlopen(). " - "Returns DISKANN_RUNTIME_OK (0) on success, or " - "DISKANN_RUNTIME_LIBAIO_MISSING if libaio is not available (DiskAnn " - "falls back to synchronous pread in that case). Returns a negative " - "code for unsupported platforms."); - m.def("is_diskann_runtime_ready", &::zvec::IsDiskAnnRuntimeReady, - "Return True if the DiskAnn runtime has been initialized."); - m.def("is_libaio_available", &::zvec::IsLibAioAvailable, - "Return True if libaio is resolvable on this host (required by the " - "DiskAnn runtime)."); - - // Status constants so callers can compare against well-known codes without - // hard-coding integers. - m.attr("DISKANN_RUNTIME_OK") = static_cast(::zvec::kDiskAnnRuntimeOk); - m.attr("DISKANN_RUNTIME_UNSUPPORTED_PLATFORM") = - static_cast(::zvec::kDiskAnnRuntimeUnsupportedPlatform); - m.attr("DISKANN_RUNTIME_LIBAIO_MISSING") = - static_cast(::zvec::kDiskAnnRuntimeLibAioMissing); - m.attr("DISKANN_RUNTIME_DLOPEN_FAILED") = - static_cast(::zvec::kDiskAnnRuntimeDlopenFailed); -} - -} // namespace - PYBIND11_MODULE(_zvec, m) { m.doc() = "Zvec core module"; @@ -70,6 +33,5 @@ PYBIND11_MODULE(_zvec, m) { ZVecPyConfig::Initialize(m); ZVecPyDoc::Initialize(m); ZVecPyCollection::Initialize(m); - InitializeDiskAnnRuntimeBindings(m); } } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index e344844f0..dd43fb29d 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -32,6 +32,21 @@ namespace core { int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { LOG_INFO("Begin DiskAnnBuilder::init"); +#if defined(__linux__) || defined(__linux) + // Eagerly load libaio at init time so the user gets immediate feedback + // about whether async I/O is available. LibAioLoader::Load() is idempotent + // and thread-safe. + if (!LibAioLoader::Instance().Load()) { + LOG_WARN( + "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " + "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " + "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " + "to synchronous pread() — performance will be degraded."); + } else { + LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + } +#endif + params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_); params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_); params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_); diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index 905fdf2bd..c56d23b8f 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -25,6 +25,21 @@ DiskAnnSearcher::DiskAnnSearcher() {} DiskAnnSearcher::~DiskAnnSearcher() {} int DiskAnnSearcher::init(const ailego::Params &search_params) { +#if defined(__linux__) || defined(__linux) + // Eagerly load libaio at init time so the user gets immediate feedback + // about whether async I/O is available. LibAioLoader::Load() is idempotent + // and thread-safe. + if (!LibAioLoader::Instance().Load()) { + LOG_WARN( + "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " + "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " + "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " + "to synchronous pread() — performance will be degraded."); + } else { + LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + } +#endif + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); return 0; diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 7a978b69e..37de6ec81 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -28,6 +28,22 @@ DiskAnnStreamer::~DiskAnnStreamer() {} int DiskAnnStreamer::init(const IndexMeta &meta, const ailego::Params &search_params) { meta_ = meta; + +#if defined(__linux__) || defined(__linux) + // Eagerly load libaio at init time so the user gets immediate feedback + // about whether async I/O is available. LibAioLoader::Load() is idempotent + // and thread-safe. + if (!LibAioLoader::Instance().Load()) { + LOG_WARN( + "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " + "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " + "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " + "to synchronous pread() — performance will be degraded."); + } else { + LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + } +#endif + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); return 0; diff --git a/src/core/interface/diskann_runtime.cc b/src/core/interface/diskann_runtime.cc deleted file mode 100644 index 9fef28690..000000000 --- a/src/core/interface/diskann_runtime.cc +++ /dev/null @@ -1,118 +0,0 @@ -// 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. - -#include -#include -#include -#include -#include - -#if defined(__linux__) || defined(__linux) || defined(__APPLE__) -#include -#include -#endif - -#if defined(__linux__) || defined(__linux) -#include -#endif - -// Include the libaio dlopen wrapper so we can load libaio eagerly at DiskAnn -// bring-up time (instead of waiting for the first I/O operation). -#if defined(__linux__) || defined(__linux) -#include "algorithm/diskann/libaio_loader.h" -#endif - -namespace zvec { - -namespace { - -#if defined(__linux__) || defined(__linux) -constexpr bool kPlatformSupportsDiskAnn = true; -#elif defined(__APPLE__) -constexpr bool kPlatformSupportsDiskAnn = false; -#else -constexpr bool kPlatformSupportsDiskAnn = false; -#endif - -// Tracks whether the DiskAnn runtime has been initialised. Since DiskAnn is -// now compiled directly into the hosting binary (_zvec.so for Python, the test -// executable for gtest, etc.) there is no separate .so to dlopen; the "loaded" -// flag simply means libaio has been probed and the result cached. -std::atomic g_runtime_ready{false}; -std::mutex g_runtime_mutex; - -} // namespace - -bool IsLibAioAvailable() { -#if defined(__linux__) || defined(__linux) - // Use the LibAioLoader singleton so we share the cached dlopen handle with - // the DiskAnn file reader. Load() is idempotent and thread-safe. - return LibAioLoader::Instance().Load(); -#else - return false; -#endif -} - -bool IsDiskAnnRuntimeReady() { - return g_runtime_ready.load(std::memory_order_acquire); -} - -int InitDiskAnnRuntime(const std::string &path) { - (void)path; // No external path needed — DiskAnn is linked in statically. - - if (!kPlatformSupportsDiskAnn) { - LOG_ERROR( - "DiskAnn is not supported on this platform; it is only " - "available on Linux x86_64 with libaio."); - return kDiskAnnRuntimeUnsupportedPlatform; - } - -#if defined(__linux__) || defined(__linux) - // Fast path: already initialised. - if (g_runtime_ready.load(std::memory_order_acquire)) { - return kDiskAnnRuntimeOk; - } - - std::lock_guard lock(g_runtime_mutex); - if (g_runtime_ready.load(std::memory_order_relaxed)) { - return kDiskAnnRuntimeOk; - } - - // Eagerly load libaio at DiskAnn bring-up time so the user gets immediate - // feedback (success or failure) rather than a delayed error on the first - // async-I/O operation. - LOG_INFO("DiskAnn: initializing runtime — loading libaio ..."); - if (!LibAioLoader::Instance().Load()) { - LOG_WARN( - "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " - "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " - "to synchronous pread() — performance will be degraded. Other " - "index types (HNSW/IVF/Flat/Vamana) are unaffected."); - // We still mark the runtime as "ready" — DiskAnn code is linked in and - // the file reader will gracefully fall back to pread(). The user gets - // a clear warning now rather than a hard failure later. - g_runtime_ready.store(true, std::memory_order_release); - return kDiskAnnRuntimeLibAioMissing; - } - - LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); - g_runtime_ready.store(true, std::memory_order_release); - return kDiskAnnRuntimeOk; -#else - return kDiskAnnRuntimeUnsupportedPlatform; -#endif -} - -} // namespace zvec diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 36f0c20fb..25274f727 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -22,10 +22,6 @@ namespace zvec::core_interface { int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { - // Platform and libaio checks are handled earlier at schema validation - // time (schema.cc calls InitDiskAnnRuntime). If libaio is missing, DiskAnn - // falls back to synchronous pread() in diskann_file_reader.cc — no need - // to re-check here. if (is_sparse_) { LOG_ERROR("Failed to create streamer. Sparse is not Supported."); return core::IndexError_Unsupported; diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 4eb673a68..b494f946e 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -174,27 +173,14 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // Probe the DiskAnn runtime eagerly at creation time so unsupported - // platforms (non Linux x86_64) fail fast with a clear message instead - // of surfacing later during optimize(). All validate() call sites - // are creation-time only, so triggering the runtime init here is - // safe (and idempotent/cached). - // - // kDiskAnnRuntimeLibAioMissing is non-fatal: DiskAnn falls back to - // synchronous pread() with degraded performance. - const int rc = ::zvec::InitDiskAnnRuntime(); - switch (rc) { - case kDiskAnnRuntimeOk: - case kDiskAnnRuntimeLibAioMissing: - break; - case kDiskAnnRuntimeUnsupportedPlatform: - return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 " - "only)"); - default: - return Status::NotSupported( - "DiskAnn runtime could not be initialized on this host"); - } + // DiskAnn requires Linux x86_64. The libaio runtime is loaded + // eagerly (via dlopen) inside DiskAnnBuilder::init() and + // DiskAnnStreamer::init(); if libaio is missing, DiskAnn falls back + // to synchronous pread() with degraded performance. +#if !defined(__linux__) && !defined(__linux) + return Status::NotSupported( + "DiskAnn is not supported on this platform (Linux x86_64 only)"); +#endif } diff --git a/src/include/zvec/core/interface/diskann_runtime.h b/src/include/zvec/core/interface/diskann_runtime.h deleted file mode 100644 index 657e12061..000000000 --- a/src/include/zvec/core/interface/diskann_runtime.h +++ /dev/null @@ -1,64 +0,0 @@ -// 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. - -#pragma once - -#include - -// NOTE: The APIs declared in this header are INTERNAL to zvec. They are -// invoked implicitly by ``DiskAnnIndex`` on first use so that DiskAnn works -// out of the box, without users ever calling a ``init_diskann_runtime()`` / -// ``is_libaio_available()`` entry point. External callers should not depend -// on these symbols; they may change or be removed in future releases. -// -// DiskAnn is compiled directly into the hosting binary (_zvec.so for the -// Python wheel, the test executable for gtest, or the tool binary for C++ -// tools). ``InitDiskAnnRuntime`` does not dlopen a separate .so — it simply -// loads libaio eagerly via dlopen()/dlsym() and logs the result. On hosts -// missing libaio, DiskAnn falls back to synchronous pread() (with a warning) -// while other index types (HNSW / IVF / Flat / Vamana) keep working. - -namespace zvec { - -// Return codes for InitDiskAnnRuntime(). -enum DiskAnnRuntimeStatus { - kDiskAnnRuntimeOk = 0, - kDiskAnnRuntimeUnsupportedPlatform = -1, - kDiskAnnRuntimeLibAioMissing = -2, - kDiskAnnRuntimeDlopenFailed = -3, -}; - -// Returns true if libaio is present on the host and the minimum set of symbols -// required by the DiskAnn runtime (io_setup / io_submit / io_getevents / -// io_destroy) can be resolved at runtime. -// -// Internal probe used by ``InitDiskAnnRuntime`` before attempting dlopen. Not -// part of the user-facing API. -bool IsLibAioAvailable(); - -// Load libaio at runtime via dlopen()/dlsym() and mark the DiskAnn runtime -// as ready. DiskAnn code is compiled directly into the hosting binary, so -// no separate .so is loaded. Invoked implicitly by -// ``DiskAnnIndex::CreateAndInitStreamer`` on first use; callers should not -// invoke it directly. The call is idempotent and returns -// ``kDiskAnnRuntimeOk`` when the runtime is already active. -// -// The ``path`` parameter is retained for API compatibility but is ignored. -int InitDiskAnnRuntime(const std::string &path = ""); - -// Returns true if the DiskAnn runtime has been initialized in this process. -// Internal diagnostic; not a user-facing API. -bool IsDiskAnnRuntimeReady(); - -} // namespace zvec diff --git a/tests/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 62d40b474..752759e9d 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -143,18 +143,6 @@ TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) { << " ms — likely a lost-wakeup regression in progress loops."; } -// DiskAnn is now exposed implicitly: no caller ever invokes a -// ``InitDiskAnnRuntime`` / ``IsLibAioAvailable`` API (those were removed from -// the public surface together with ``zvec.init_diskann_runtime()`` in Python). -// The only contract this test validates is the UX guarantee: once the DiskAnn -// module has been linked into the hosting binary (here, directly into the -// test via the ``core_knn_diskann`` target), its factory entries are -// registered automatically and the global ``IndexFactory`` can hand out a -// ``DiskAnnBuilder`` without any explicit setup step. On hosts missing -// libaio, DiskAnn falls back to synchronous pread() with a warning while -// other index types (HNSW/IVF/Flat/Vamana) remain unaffected; that runtime -// branch lives in ``DiskAnnIndex::CreateAndInitStreamer`` and is covered by -// the higher-level interface tests. TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) { IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); ASSERT_NE(builder, nullptr) From 30a08f84c2152464550b0ff8b2cd64d97a2b784c Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 10:24:15 +0800 Subject: [PATCH 18/46] fix: fix ut --- tests/db/crash_recovery/write_recovery_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/db/crash_recovery/write_recovery_test.cc b/tests/db/crash_recovery/write_recovery_test.cc index 4dce9985a..2e595eb87 100644 --- a/tests/db/crash_recovery/write_recovery_test.cc +++ b/tests/db/crash_recovery/write_recovery_test.cc @@ -188,7 +188,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { collection.reset(); } - RunGeneratorAndCrash("0", "10000", "insert", "0", 3); + RunGeneratorAndCrash("0", "10000", "insert", "0", 5); auto result = Collection::Open(dir_path_, options_); ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " @@ -196,7 +196,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { auto collection = result.value(); uint64_t doc_count{collection->Stats().value().doc_count}; ASSERT_GT(doc_count, 800) - << "Document count is too low after 3s of insertion and recovery"; + << "Document count is too low after 5s of insertion and recovery"; for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) { const auto expected_doc = CreateTestDoc(doc_id, 0); From 5ffc23a1d7261ac9223273fb2bc24767db336646 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 15:02:42 +0800 Subject: [PATCH 19/46] refactor: move libaio loader to ailego --- .../algorithm/diskann/diskann_file_reader.h | 2 +- src/core/algorithm/diskann/libaio_def.h | 190 ------------------ src/core/algorithm/diskann/libaio_loader.h | 144 ------------- 3 files changed, 1 insertion(+), 335 deletions(-) delete mode 100644 src/core/algorithm/diskann/libaio_def.h delete mode 100644 src/core/algorithm/diskann/libaio_loader.h diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 2abe8c1aa..a81ee401a 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -18,7 +18,7 @@ #include #if (defined(__linux) || defined(__linux__)) -#include "libaio_loader.h" // dlopen-based libaio wrapper +#include // dlopen-based libaio wrapper #endif #include diff --git a/src/core/algorithm/diskann/libaio_def.h b/src/core/algorithm/diskann/libaio_def.h deleted file mode 100644 index ce95bc982..000000000 --- a/src/core/algorithm/diskann/libaio_def.h +++ /dev/null @@ -1,190 +0,0 @@ -// 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 replacement for . -// -// This header declares *only* the types, constants, and inline helpers that -// zvec needs from libaio. By doing so the project is completely decoupled -// from the system libaio-dev header: there is no `#include ` anywhere -// in the source tree, which means libaio-dev does not need to be installed at -// build time and the code is portable to cross-compilation environments that -// lack the header. -// -// The struct layouts (struct iocb, struct io_event, ...) are part of the Linux -// kernel ABI. They are copied verbatim from the upstream , including -// the PADDED macros that handle architecture-specific padding. The inline -// helper io_prep_pread() is likewise copied — it only manipulates struct fields -// and does not call into the library. - -#pragma once - -#include // struct timespec (used by io_getevents signature) -#include // memset() — used by io_prep_pread() inline helper - -#if defined(__linux) || defined(__linux__) - -struct sockaddr; -struct iovec; - -// --------------------------------------------------------------------------- -// Type and struct definitions copied from -// --------------------------------------------------------------------------- - -typedef struct io_context *io_context_t; - -typedef enum io_iocb_cmd { - IO_CMD_PREAD = 0, - IO_CMD_PWRITE = 1, - IO_CMD_FSYNC = 2, - IO_CMD_FDSYNC = 3, - IO_CMD_POLL = 5, - IO_CMD_NOOP = 6, - IO_CMD_PREADV = 7, - IO_CMD_PWRITEV = 8, -} io_iocb_cmd_t; - -// PADDED macros — copied verbatim from to guarantee ABI-compatible -// struct layout on every supported architecture. - -/* little endian, 32 bits */ -#if defined(__i386__) || (defined(__x86_64__) && defined(__ILP32__)) || \ - (defined(__arm__) && !defined(__ARMEB__)) || \ - (defined(__sh__) && defined(__LITTLE_ENDIAN__)) || defined(__bfin__) || \ - (defined(__MIPSEL__) && !defined(__mips64)) || defined(__cris__) || \ - defined(__loongarch32) || (defined(__riscv) && __riscv_xlen == 32) || \ - (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ - __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 4) -#define AIO_PADDED(x, y) \ - x; \ - unsigned y -#define AIO_PADDEDptr(x, y) \ - x; \ - unsigned y -#define AIO_PADDEDul(x, y) \ - unsigned long x; \ - unsigned y - -/* little endian, 64 bits */ -#elif defined(__ia64__) || defined(__x86_64__) || defined(__alpha__) || \ - (defined(__mips64) && defined(__MIPSEL__)) || \ - (defined(__aarch64__) && defined(__AARCH64EL__)) || \ - defined(__loongarch64) || (defined(__riscv) && __riscv_xlen == 64) || \ - (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ - __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 8) -#define AIO_PADDED(x, y) x, y -#define AIO_PADDEDptr(x, y) x -#define AIO_PADDEDul(x, y) unsigned long x - -/* big endian, 64 bits */ -#elif defined(__powerpc64__) || defined(__s390x__) || \ - (defined(__hppa__) && defined(__arch64__)) || \ - (defined(__sparc__) && defined(__arch64__)) || \ - (defined(__mips64) && defined(__MIPSEB__)) || \ - (defined(__aarch64__) && defined(__AARCH64EB__)) || \ - (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ - __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 8) -#define AIO_PADDED(x, y) \ - unsigned y; \ - x -#define AIO_PADDEDptr(x, y) x -#define AIO_PADDEDul(x, y) unsigned long x - -/* big endian, 32 bits */ -#elif defined(__PPC__) || defined(__s390__) || \ - (defined(__arm__) && defined(__ARMEB__)) || \ - (defined(__sh__) && defined(__BIG_ENDIAN__)) || defined(__sparc__) || \ - defined(__MIPSEB__) || defined(__m68k__) || defined(__hppa__) || \ - defined(__frv__) || defined(__avr32__) || \ - (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ - __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 4) -#define AIO_PADDED(x, y) \ - unsigned y; \ - x -#define AIO_PADDEDptr(x, y) \ - unsigned y; \ - x -#define AIO_PADDEDul(x, y) \ - unsigned y; \ - unsigned long x - -#else -#error endianness? -#endif - -struct io_iocb_poll { - AIO_PADDED(int events, __pad1); -}; - -struct io_iocb_sockaddr { - AIO_PADDEDptr(struct sockaddr *addr, __pad1); - AIO_PADDEDul(len, __pad2); -}; - -struct io_iocb_common { - AIO_PADDEDptr(void *buf, __pad1); - AIO_PADDEDul(nbytes, __pad2); - long long offset; - long long __pad3; - unsigned flags; - unsigned resfd; -}; - -struct io_iocb_vector { - AIO_PADDEDptr(const struct iovec *vec, __pad1); - AIO_PADDEDul(nr, __pad2); - long long offset; -}; - -struct iocb { - AIO_PADDEDptr(void *data, __pad1); - AIO_PADDED(unsigned key, aio_rw_flags); - short aio_lio_opcode; - short aio_reqprio; - int aio_fildes; - union { - struct io_iocb_common c; - struct io_iocb_vector v; - struct io_iocb_poll poll; - struct io_iocb_sockaddr saddr; - } u; -}; - -struct io_event { - AIO_PADDEDptr(void *data, __pad1); - AIO_PADDEDptr(struct iocb *obj, __pad2); - AIO_PADDEDul(res, __pad3); - AIO_PADDEDul(res2, __pad4); -}; - -#undef AIO_PADDED -#undef AIO_PADDEDptr -#undef AIO_PADDEDul - -// Inline helper — copied from . Only manipulates struct fields. -static inline void io_prep_pread(struct iocb *iocb, int fd, void *buf, - size_t count, long long offset) { - memset(iocb, 0, sizeof(*iocb)); - iocb->aio_fildes = fd; - iocb->aio_lio_opcode = IO_CMD_PREAD; - iocb->aio_reqprio = 0; - iocb->u.c.buf = buf; - iocb->u.c.nbytes = count; - iocb->u.c.offset = offset; -} - -// --------------------------------------------------------------------------- -// End: type and struct definitions from -// --------------------------------------------------------------------------- - -#endif // __linux__ diff --git a/src/core/algorithm/diskann/libaio_loader.h b/src/core/algorithm/diskann/libaio_loader.h deleted file mode 100644 index dca400ff9..000000000 --- a/src/core/algorithm/diskann/libaio_loader.h +++ /dev/null @@ -1,144 +0,0 @@ -// 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. - -// dlopen-based wrapper for libaio. Instead of linking against -laio at build -// time, the DiskAnn plugin loads libaio at runtime via dlopen()/dlsym(). This -// removes libaio as a hard build- and load-time dependency: the plugin .so no -// longer carries a NEEDED entry for libaio.so.1, so it can be dlopen'ed on -// hosts that don't have libaio installed. Callers that actually exercise the -// async-I/O path will fall back to synchronous pread() when libaio is absent. -// -// All ABI-stable type and struct definitions (struct iocb, struct io_event, -// io_context_t, PADDED macros, io_prep_pread(), ...) now live in libaio_def.h — -// a private header that replaces so the project has zero build-time -// dependency on libaio-dev. - -#pragma once - -#if defined(__linux) || defined(__linux__) - -#include -#include -#include -#include -#include "libaio_def.h" // ABI-stable struct definitions (replaces ) - -// Function-pointer typedefs for the four libaio syscalls used by DiskAnn. -typedef int (*aio_setup_fn)(int maxevents, io_context_t *ctxp); -typedef int (*aio_destroy_fn)(io_context_t ctx); -typedef int (*aio_submit_fn)(io_context_t ctx, long nr, struct iocb *ios[]); -typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr, - struct io_event *events, - struct timespec *timeout); - -// Runtime loader for libaio. Thread-safe singleton that dlopen()'s libaio -// once and caches the function pointers. If libaio is not present on the -// host, all pointers remain nullptr and callers should fall back to -// synchronous I/O (pread). -// -// Usage: -// if (LibAioLoader::Instance().Load()) { -// LibAioLoader::Instance().io_setup(...); -// } -class LibAioLoader { - public: - static LibAioLoader &Instance() { - static LibAioLoader instance; - return instance; - } - - // Load (or confirm already loaded) libaio. Returns true on success. - // Thread-safe and idempotent. - bool Load() { - if (available_.load(std::memory_order_acquire)) { - return true; - } - std::call_once(once_, [this] { this->TryLoad(); }); - return available_.load(std::memory_order_relaxed); - } - - bool IsAvailable() const { - return available_.load(std::memory_order_acquire); - } - - // Function pointers — nullptr until Load() succeeds. - aio_setup_fn io_setup; - aio_destroy_fn io_destroy; - aio_submit_fn io_submit; - aio_getevents_fn io_getevents; - - private: - LibAioLoader() - : io_setup(nullptr), - io_destroy(nullptr), - io_submit(nullptr), - io_getevents(nullptr) {} - - ~LibAioLoader() { - if (handle_ != nullptr) { - dlclose(handle_); - } - } - - LibAioLoader(const LibAioLoader &) = delete; - LibAioLoader &operator=(const LibAioLoader &) = delete; - - void TryLoad() { - // On Ubuntu 24.04 the libaio package was renamed with the t64 suffix - // (64-bit time_t transition), so probe both spellings. - static constexpr const char *kSonames[] = { - "libaio.so.1", - "libaio.so.1t64", - }; - - for (const char *soname : kSonames) { - void *h = dlopen(soname, RTLD_LAZY); - if (h == nullptr) { - continue; - } - - io_setup = reinterpret_cast(dlsym(h, "io_setup")); - io_destroy = reinterpret_cast(dlsym(h, "io_destroy")); - io_submit = reinterpret_cast(dlsym(h, "io_submit")); - - // io_getevents may be redirected to io_getevents_time64 on 32-bit - // platforms compiled with _TIME_BITS=64. - io_getevents = - reinterpret_cast(dlsym(h, "io_getevents")); - if (io_getevents == nullptr) { - io_getevents = - reinterpret_cast(dlsym(h, "io_getevents_time64")); - } - - if (io_setup && io_destroy && io_submit && io_getevents) { - handle_ = h; - available_.store(true, std::memory_order_release); - return; - } - - // Some symbols missing — try the next soname. - dlclose(h); - io_setup = nullptr; - io_destroy = nullptr; - io_submit = nullptr; - io_getevents = nullptr; - } - } - - std::once_flag once_; - std::atomic available_{false}; - void *handle_{nullptr}; -}; - -#endif // __linux__ From 1c7ff606c286b038578069a1d21e15c65cb2f2b4 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 16:29:12 +0800 Subject: [PATCH 20/46] fix: add io backend --- python/zvec/__init__.py | 2 + src/ailego/io/io_backend.h | 114 +++++++++++ src/ailego/io/libaio_def.h | 190 ++++++++++++++++++ src/ailego/io/libaio_loader.h | 144 +++++++++++++ .../python/model/common/python_config.cc | 21 ++ src/core/algorithm/diskann/diskann_builder.cc | 21 +- .../algorithm/diskann/diskann_file_reader.cc | 41 ++-- .../algorithm/diskann/diskann_searcher.cc | 21 +- .../algorithm/diskann/diskann_streamer.cc | 21 +- src/core/interface/indexes/diskann_index.cc | 13 ++ src/include/zvec/core/interface/index.h | 5 + 11 files changed, 548 insertions(+), 45 deletions(-) create mode 100644 src/ailego/io/io_backend.h create mode 100644 src/ailego/io/libaio_def.h create mode 100644 src/ailego/io/libaio_loader.h diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index 9e02c7575..d11cfea69 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -30,6 +30,7 @@ from zvec._zvec import ( get_default_jieba_dict_dir, + io_backend_type, set_default_jieba_dict_dir, ) @@ -127,6 +128,7 @@ "open", "set_default_jieba_dict_dir", "get_default_jieba_dict_dir", + "io_backend_type", # Core classes "Collection", "Doc", diff --git a/src/ailego/io/io_backend.h b/src/ailego/io/io_backend.h new file mode 100644 index 000000000..647fb653c --- /dev/null +++ b/src/ailego/io/io_backend.h @@ -0,0 +1,114 @@ +// 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. + +// Abstract I/O backend selector. +// +// 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. +// +// When no async backend is available, the caller should fall back to +// synchronous pread(). +// +// Usage: +// auto& backend = ailego::IOBackend::Instance(); +// if (backend.available() != ailego::IOBackendType::kSyncPread) { ... } +// LOG_INFO("I/O backend: %s", backend.name()); + +#pragma once + +#include + +namespace zvec { +namespace ailego { + +// Supported I/O backend types. +enum class IOBackendType { + kSyncPread, // Synchronous pread() — no async I/O + kLibAio, // libaio loaded at runtime via dlopen() +}; + +// Returns a human-readable name for the given backend type. +inline const char *IOBackendTypeName(IOBackendType type) { + switch (type) { + case IOBackendType::kLibAio: + return "libaio"; + case IOBackendType::kSyncPread: + return "sync_pread"; + } + return "unknown"; +} + +// Singleton that loads and queries an I/O backend on demand. +// +// 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. +class IOBackend { + public: + static IOBackend &Instance() { + static IOBackend instance; + return instance; + } + + // Try to load the best available backend (libaio > sync_pread). + // Returns the loaded backend type. + // Idempotent — if already loaded, returns immediately. + IOBackendType available() { + if (type_ != IOBackendType::kSyncPread) { + 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 kSyncPread). + // Idempotent — if the same backend is already loaded, returns immediately. + IOBackendType available(IOBackendType requested) { + if (type_ == requested && type_ != IOBackendType::kSyncPread) { + return type_; + } +#if defined(__linux) || defined(__linux__) + if (requested == IOBackendType::kLibAio) { + if (LibAioLoader::Instance().load() && + LibAioLoader::Instance().is_available()) { + type_ = IOBackendType::kLibAio; + return type_; + } + } +#endif + type_ = IOBackendType::kSyncPread; + return type_; + } + + // Returns the loaded backend type. + IOBackendType type() const { + return type_; + } + + // Human-readable name for the selected backend. + const char *name() const { + return IOBackendTypeName(type_); + } + + private: + IOBackend() = default; + + IOBackendType type_{IOBackendType::kSyncPread}; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/ailego/io/libaio_def.h b/src/ailego/io/libaio_def.h new file mode 100644 index 000000000..ce95bc982 --- /dev/null +++ b/src/ailego/io/libaio_def.h @@ -0,0 +1,190 @@ +// 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 replacement for . +// +// This header declares *only* the types, constants, and inline helpers that +// zvec needs from libaio. By doing so the project is completely decoupled +// from the system libaio-dev header: there is no `#include ` anywhere +// in the source tree, which means libaio-dev does not need to be installed at +// build time and the code is portable to cross-compilation environments that +// lack the header. +// +// The struct layouts (struct iocb, struct io_event, ...) are part of the Linux +// kernel ABI. They are copied verbatim from the upstream , including +// the PADDED macros that handle architecture-specific padding. The inline +// helper io_prep_pread() is likewise copied — it only manipulates struct fields +// and does not call into the library. + +#pragma once + +#include // struct timespec (used by io_getevents signature) +#include // memset() — used by io_prep_pread() inline helper + +#if defined(__linux) || defined(__linux__) + +struct sockaddr; +struct iovec; + +// --------------------------------------------------------------------------- +// Type and struct definitions copied from +// --------------------------------------------------------------------------- + +typedef struct io_context *io_context_t; + +typedef enum io_iocb_cmd { + IO_CMD_PREAD = 0, + IO_CMD_PWRITE = 1, + IO_CMD_FSYNC = 2, + IO_CMD_FDSYNC = 3, + IO_CMD_POLL = 5, + IO_CMD_NOOP = 6, + IO_CMD_PREADV = 7, + IO_CMD_PWRITEV = 8, +} io_iocb_cmd_t; + +// PADDED macros — copied verbatim from to guarantee ABI-compatible +// struct layout on every supported architecture. + +/* little endian, 32 bits */ +#if defined(__i386__) || (defined(__x86_64__) && defined(__ILP32__)) || \ + (defined(__arm__) && !defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__LITTLE_ENDIAN__)) || defined(__bfin__) || \ + (defined(__MIPSEL__) && !defined(__mips64)) || defined(__cris__) || \ + defined(__loongarch32) || (defined(__riscv) && __riscv_xlen == 32) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDptr(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDul(x, y) \ + unsigned long x; \ + unsigned y + +/* little endian, 64 bits */ +#elif defined(__ia64__) || defined(__x86_64__) || defined(__alpha__) || \ + (defined(__mips64) && defined(__MIPSEL__)) || \ + (defined(__aarch64__) && defined(__AARCH64EL__)) || \ + defined(__loongarch64) || (defined(__riscv) && __riscv_xlen == 64) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) x, y +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +/* big endian, 64 bits */ +#elif defined(__powerpc64__) || defined(__s390x__) || \ + (defined(__hppa__) && defined(__arch64__)) || \ + (defined(__sparc__) && defined(__arch64__)) || \ + (defined(__mips64) && defined(__MIPSEB__)) || \ + (defined(__aarch64__) && defined(__AARCH64EB__)) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +/* big endian, 32 bits */ +#elif defined(__PPC__) || defined(__s390__) || \ + (defined(__arm__) && defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__BIG_ENDIAN__)) || defined(__sparc__) || \ + defined(__MIPSEB__) || defined(__m68k__) || defined(__hppa__) || \ + defined(__frv__) || defined(__avr32__) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDul(x, y) \ + unsigned y; \ + unsigned long x + +#else +#error endianness? +#endif + +struct io_iocb_poll { + AIO_PADDED(int events, __pad1); +}; + +struct io_iocb_sockaddr { + AIO_PADDEDptr(struct sockaddr *addr, __pad1); + AIO_PADDEDul(len, __pad2); +}; + +struct io_iocb_common { + AIO_PADDEDptr(void *buf, __pad1); + AIO_PADDEDul(nbytes, __pad2); + long long offset; + long long __pad3; + unsigned flags; + unsigned resfd; +}; + +struct io_iocb_vector { + AIO_PADDEDptr(const struct iovec *vec, __pad1); + AIO_PADDEDul(nr, __pad2); + long long offset; +}; + +struct iocb { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDED(unsigned key, aio_rw_flags); + short aio_lio_opcode; + short aio_reqprio; + int aio_fildes; + union { + struct io_iocb_common c; + struct io_iocb_vector v; + struct io_iocb_poll poll; + struct io_iocb_sockaddr saddr; + } u; +}; + +struct io_event { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDEDptr(struct iocb *obj, __pad2); + AIO_PADDEDul(res, __pad3); + AIO_PADDEDul(res2, __pad4); +}; + +#undef AIO_PADDED +#undef AIO_PADDEDptr +#undef AIO_PADDEDul + +// Inline helper — copied from . Only manipulates struct fields. +static inline void io_prep_pread(struct iocb *iocb, int fd, void *buf, + size_t count, long long offset) { + memset(iocb, 0, sizeof(*iocb)); + iocb->aio_fildes = fd; + iocb->aio_lio_opcode = IO_CMD_PREAD; + iocb->aio_reqprio = 0; + iocb->u.c.buf = buf; + iocb->u.c.nbytes = count; + iocb->u.c.offset = offset; +} + +// --------------------------------------------------------------------------- +// End: type and struct definitions from +// --------------------------------------------------------------------------- + +#endif // __linux__ diff --git a/src/ailego/io/libaio_loader.h b/src/ailego/io/libaio_loader.h new file mode 100644 index 000000000..7fbbbbfa1 --- /dev/null +++ b/src/ailego/io/libaio_loader.h @@ -0,0 +1,144 @@ +// 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. + +// dlopen-based wrapper for libaio. Instead of linking against -laio at build +// time, the DiskAnn plugin loads libaio at runtime via dlopen()/dlsym(). This +// removes libaio as a hard build- and load-time dependency: the plugin .so no +// longer carries a NEEDED entry for libaio.so.1, so it can be dlopen'ed on +// hosts that don't have libaio installed. Callers that actually exercise the +// async-I/O path will fall back to synchronous pread() when libaio is absent. +// +// All ABI-stable type and struct definitions (struct iocb, struct io_event, +// io_context_t, PADDED macros, io_prep_pread(), ...) now live in libaio_def.h — +// a private header that replaces so the project has zero build-time +// dependency on libaio-dev. + +#pragma once + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include +#include // ABI-stable struct definitions (replaces ) + +// Function-pointer typedefs for the four libaio syscalls used by DiskAnn. +typedef int (*aio_setup_fn)(int maxevents, io_context_t *ctxp); +typedef int (*aio_destroy_fn)(io_context_t ctx); +typedef int (*aio_submit_fn)(io_context_t ctx, long nr, struct iocb *ios[]); +typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr, + struct io_event *events, + struct timespec *timeout); + +// Runtime loader for libaio. Thread-safe singleton that dlopen()'s libaio +// once and caches the function pointers. If libaio is not present on the +// host, all pointers remain nullptr and callers should fall back to +// synchronous I/O (pread). +// +// Usage: +// if (LibAioLoader::Instance().load()) { +// LibAioLoader::Instance().io_setup(...); +// } +class LibAioLoader { + public: + static LibAioLoader &Instance() { + static LibAioLoader instance; + return instance; + } + + // Load (or confirm already loaded) libaio. Returns true on success. + // Thread-safe and idempotent. + bool load() { + if (available_.load(std::memory_order_acquire)) { + return true; + } + std::call_once(once_, [this] { this->try_load(); }); + return available_.load(std::memory_order_relaxed); + } + + bool is_available() const { + return available_.load(std::memory_order_acquire); + } + + // Function pointers — nullptr until load() succeeds. + aio_setup_fn io_setup; + aio_destroy_fn io_destroy; + aio_submit_fn io_submit; + aio_getevents_fn io_getevents; + + private: + LibAioLoader() + : io_setup(nullptr), + io_destroy(nullptr), + io_submit(nullptr), + io_getevents(nullptr) {} + + ~LibAioLoader() { + if (handle_ != nullptr) { + dlclose(handle_); + } + } + + LibAioLoader(const LibAioLoader &) = delete; + LibAioLoader &operator=(const LibAioLoader &) = delete; + + void try_load() { + // On Ubuntu 24.04 the libaio package was renamed with the t64 suffix + // (64-bit time_t transition), so probe both spellings. + static constexpr const char *kSonames[] = { + "libaio.so.1", + "libaio.so.1t64", + }; + + for (const char *soname : kSonames) { + void *h = dlopen(soname, RTLD_LAZY); + if (h == nullptr) { + continue; + } + + io_setup = reinterpret_cast(dlsym(h, "io_setup")); + io_destroy = reinterpret_cast(dlsym(h, "io_destroy")); + io_submit = reinterpret_cast(dlsym(h, "io_submit")); + + // io_getevents may be redirected to io_getevents_time64 on 32-bit + // platforms compiled with _TIME_BITS=64. + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents")); + if (io_getevents == nullptr) { + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents_time64")); + } + + if (io_setup && io_destroy && io_submit && io_getevents) { + handle_ = h; + available_.store(true, std::memory_order_release); + return; + } + + // Some symbols missing — try the next soname. + dlclose(h); + io_setup = nullptr; + io_destroy = nullptr; + io_submit = nullptr; + io_getevents = nullptr; + } + } + + std::once_flag once_; + std::atomic available_{false}; + void *handle_{nullptr}; +}; + +#endif // __linux__ diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index dade4bbc3..866f14c9a 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "python_config.h" +#include #include namespace zvec { @@ -217,6 +218,26 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { "get_default_jieba_dict_dir", []() -> std::string { return GlobalConfig::Instance().jieba_dict_dir(); }, "Read the currently registered default jieba dict directory."); + + // Returns the current I/O backend type for DiskAnn async disk reads. + // When only sync_pread is available, prints the install hint to stdout. + m.def( + "io_backend_type", + []() -> std::string { + auto type = ailego::IOBackend::Instance().available(); + std::string name = ailego::IOBackendTypeName(type); + if (type == ailego::IOBackendType::kSyncPread) { + py::print( + "No async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " + "and retry. DiskAnn will fall back to synchronous pread() — " + "performance will be degraded."); + } + return name; + }, + "Returns the current I/O backend type for DiskAnn async disk reads. " + "\"libaio\" if libaio is available, \"sync_pread\" otherwise. " + "When \"sync_pread\", prints the install hint to stdout."); } diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index dd43fb29d..0186dfdf3 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -33,17 +34,19 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { LOG_INFO("Begin DiskAnnBuilder::init"); #if defined(__linux__) || defined(__linux) - // Eagerly load libaio at init time so the user gets immediate feedback - // about whether async I/O is available. LibAioLoader::Load() is idempotent - // and thread-safe. - if (!LibAioLoader::Instance().Load()) { + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kSyncPread) { LOG_WARN( - "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " - "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " - "to synchronous pread() — performance will be degraded."); + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); } else { - LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); } #endif diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index bc19d8eb5..b1b1dd033 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -18,6 +18,7 @@ #include #include #include +#include #include #define MAX_EVENTS 1024 @@ -37,15 +38,16 @@ static std::once_flag g_io_backend_log_once; int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - LibAioLoader::Instance().Load(); - std::call_once(g_io_backend_log_once, [] { - if (LibAioLoader::Instance().IsAvailable()) { - LOG_INFO("DiskAnn I/O backend: libaio (async I/O enabled)"); + auto &backend = ailego::IOBackend::Instance(); + std::call_once(g_io_backend_log_once, [&backend] { + if (backend.available() != ailego::IOBackendType::kSyncPread) { + LOG_INFO("DiskAnn I/O backend: %s (async I/O enabled)", backend.name()); } else { - LOG_WARN("DiskAnn I/O backend: synchronous pread (libaio not available)"); + LOG_WARN( + "DiskAnn I/O backend: synchronous pread (no async I/O available)"); } }); - if (!LibAioLoader::Instance().IsAvailable()) { + if (backend.available() == ailego::IOBackendType::kSyncPread) { return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); @@ -58,7 +60,8 @@ int setup_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - if (!LibAioLoader::Instance().IsAvailable()) { + if (ailego::IOBackend::Instance().available() == + ailego::IOBackendType::kSyncPread) { return 0; } int ret = LibAioLoader::Instance().io_destroy(ctx); @@ -90,7 +93,8 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { int execute_io(IOContext ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) - if (!LibAioLoader::Instance().Load()) { + if (ailego::IOBackend::Instance().available() == + ailego::IOBackendType::kSyncPread) { return execute_io_pread(fd, read_reqs); } uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); @@ -211,28 +215,27 @@ void LinuxAlignedFileReader::register_thread() { IOContext ctx = nullptr; - LibAioLoader::Instance().Load(); - std::call_once(g_io_backend_log_once, [] { - if (LibAioLoader::Instance().IsAvailable()) { - LOG_INFO("DiskAnn I/O backend: libaio (async I/O enabled)"); + auto &backend = ailego::IOBackend::Instance(); + std::call_once(g_io_backend_log_once, [&backend] { + if (backend.available() != ailego::IOBackendType::kSyncPread) { + LOG_INFO("DiskAnn I/O backend: %s (async I/O enabled)", backend.name()); } else { - LOG_WARN("DiskAnn I/O backend: synchronous pread (libaio not available)"); + LOG_WARN( + "DiskAnn I/O backend: synchronous pread (no async I/O available)"); } }); - if (!LibAioLoader::Instance().IsAvailable()) { + if (backend.available() == ailego::IOBackendType::kSyncPread) { lk.unlock(); return; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); if (ret != 0) { - lk.unlock(); 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 { LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); @@ -261,7 +264,8 @@ void LinuxAlignedFileReader::deregister_thread() { } // io_destroy is a syscall; keep it outside the lock to avoid blocking others - if (LibAioLoader::Instance().IsAvailable()) { + if (ailego::IOBackend::Instance().available() != + ailego::IOBackendType::kSyncPread) { LibAioLoader::Instance().io_destroy(ctx); } LOG_INFO("returned ctx from thread"); @@ -271,7 +275,8 @@ void LinuxAlignedFileReader::deregister_thread() { void LinuxAlignedFileReader::deregister_all_threads() { #if (defined(__linux) || defined(__linux__)) std::unique_lock lk(ctx_mut); - bool aio_available = LibAioLoader::Instance().IsAvailable(); + bool aio_available = ailego::IOBackend::Instance().available() != + ailego::IOBackendType::kSyncPread; for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; if (aio_available) { diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index c56d23b8f..75c6d81d1 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_searcher.h" +#include #include "diskann_context.h" #include "diskann_indexer.h" #include "diskann_params.h" @@ -26,17 +27,19 @@ DiskAnnSearcher::~DiskAnnSearcher() {} int DiskAnnSearcher::init(const ailego::Params &search_params) { #if defined(__linux__) || defined(__linux) - // Eagerly load libaio at init time so the user gets immediate feedback - // about whether async I/O is available. LibAioLoader::Load() is idempotent - // and thread-safe. - if (!LibAioLoader::Instance().Load()) { + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kSyncPread) { LOG_WARN( - "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " - "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " - "to synchronous pread() — performance will be degraded."); + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); } else { - LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); } #endif diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 37de6ec81..6fbd22b1c 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_streamer.h" +#include #include "diskann_context.h" #include "diskann_index_provider.h" #include "diskann_indexer.h" @@ -30,17 +31,19 @@ int DiskAnnStreamer::init(const IndexMeta &meta, meta_ = meta; #if defined(__linux__) || defined(__linux) - // Eagerly load libaio at init time so the user gets immediate feedback - // about whether async I/O is available. LibAioLoader::Load() is idempotent - // and thread-safe. - if (!LibAioLoader::Instance().Load()) { + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kSyncPread) { LOG_WARN( - "DiskAnn: libaio could not be loaded (tried libaio.so.1 and " - "libaio.so.1t64). Install it (e.g. 'apt-get install libaio1', or " - "'libaio1t64' on Ubuntu 24.04+) and retry. DiskAnn will fall back " - "to synchronous pread() — performance will be degraded."); + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); } else { - LOG_INFO("DiskAnn: libaio loaded successfully — async I/O enabled."); + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); } #endif diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 25274f727..111bb8052 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include "algorithm/diskann/diskann_params.h" #include "holder_builder.h" @@ -266,4 +267,16 @@ int DiskAnnIndex::Merge(const std::vector &indexes, return 0; } +ailego::IOBackendType DiskAnnIndex::io_backend_type() const { + auto &backend = ailego::IOBackend::Instance(); + ailego::IOBackendType type = backend.type(); + if (type == ailego::IOBackendType::kSyncPread) { + LOG_WARN( + "DiskAnn: only synchronous pread() is available. Install libaio " + "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " + "for async I/O support — performance will be degraded without it."); + } + return type; +} + } // namespace zvec::core_interface \ No newline at end of file diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index e9ed39051..7b4d288e7 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -360,6 +361,10 @@ class DiskAnnIndex : public Index { public: DiskAnnIndex() = default; + // Returns the I/O backend type currently loaded for DiskAnn async disk reads. + // If only sync_pread is available, logs a hint to install libaio. + ailego::IOBackendType io_backend_type() const; + protected: virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; From 54eaa118d59a0083a620e9155689d9248796a1fe Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 16:33:41 +0800 Subject: [PATCH 21/46] fix: add io backend --- src/core/interface/indexes/diskann_index.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 111bb8052..e518898ab 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -272,7 +272,7 @@ ailego::IOBackendType DiskAnnIndex::io_backend_type() const { ailego::IOBackendType type = backend.type(); if (type == ailego::IOBackendType::kSyncPread) { LOG_WARN( - "DiskAnn: only synchronous pread() is available. Install libaio " + "Only synchronous pread() is available. Install libaio " "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " "for async I/O support — performance will be degraded without it."); } From d2fba31931be402804abae385f4a15a92a368c44 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 16:56:35 +0800 Subject: [PATCH 22/46] fix: fix memcpy --- src/core/algorithm/diskann/diskann_entity.h | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h index af302290d..cdaeaed33 100644 --- a/src/core/algorithm/diskann/diskann_entity.h +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -72,17 +72,6 @@ struct DiskAnnMetaHeader { clear(); } - DiskAnnMetaHeader(const DiskAnnMetaHeader &header) { - memcpy(this, &header, sizeof(header)); - } - - DiskAnnMetaHeader &operator=(const DiskAnnMetaHeader &header) { - if (this != &header) { - memcpy(this, &header, sizeof(header)); - } - return *this; - } - void reset() { doc_cnt = 0U; } @@ -104,15 +93,6 @@ struct DiskAnnPqMeta { clear(); } - DiskAnnPqMeta(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - } - - DiskAnnPqMeta &operator=(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - return *this; - } - void clear() { memset(this, 0, sizeof(DiskAnnPqMeta)); } From 716adb90f756e1c4c30410e2aa42051f3586b330 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 7 Jul 2026 17:06:59 +0800 Subject: [PATCH 23/46] fix: fix warning --- src/core/algorithm/diskann/diskann_searcher_entity.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc index c9e49deba..3e353a216 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.cc +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -398,8 +398,8 @@ const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { const void *vec; if (ailego_unlikely(vector_segment_->read(total_offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed, id: %u, offset: %lu", id, - total_offset); + LOG_ERROR("Read vector from segment failed, id: %u, offset: %llu", id, + (unsigned long long)total_offset); return nullptr; } From e1048c067078483fbdd0c99bd2525b85b9837a69 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 8 Jul 2026 11:08:28 +0800 Subject: [PATCH 24/46] fix: remove ci --- .github/workflows/01-ci-pipeline.yml | 8 -- .github/workflows/09-diskann-libaio-test.yml | 109 ------------------- 2 files changed, 117 deletions(-) delete mode 100644 .github/workflows/09-diskann-libaio-test.yml diff --git a/.github/workflows/01-ci-pipeline.yml b/.github/workflows/01-ci-pipeline.yml index b0c3c7e27..ce9f8e2ef 100644 --- a/.github/workflows/01-ci-pipeline.yml +++ b/.github/workflows/01-ci-pipeline.yml @@ -97,14 +97,6 @@ jobs: os: ubuntu-24.04 compiler: clang - # DiskAnn w/ libaio: explicitly install libaio and run tests to exercise - # the async I/O path. The existing CI (without libaio-dev) already covers - # the w/o libaio pread() fallback path. - diskann-libaio-test: - name: DiskAnn w/ libaio (linux-x64) - needs: [lint, clang-tidy] - uses: ./.github/workflows/09-diskann-libaio-test.yml - build-android: if: github.event_name != 'push' || github.ref != 'refs/heads/main' name: Build & Test (android) diff --git a/.github/workflows/09-diskann-libaio-test.yml b/.github/workflows/09-diskann-libaio-test.yml deleted file mode 100644 index 1024c9ffb..000000000 --- a/.github/workflows/09-diskann-libaio-test.yml +++ /dev/null @@ -1,109 +0,0 @@ -name: DiskAnn libaio Test - -# Tests DiskAnn on Linux x86_64 with libaio explicitly installed, exercising -# the async I/O path (io_setup/io_submit/io_getevents). -# -# The existing CI (03-macos-linux-build.yml without libaio-dev installed) and -# non-Linux platform jobs already cover the "w/o libaio" fallback path — -# DiskAnn gracefully degrades to synchronous pread() when libaio is absent. -# This workflow complements that by ensuring the async I/O path is also -# exercised when libaio is present. - -on: - workflow_call: - workflow_dispatch: - -permissions: - contents: read - -jobs: - libaio-test: - name: DiskAnn w/ libaio (linux-x64) - runs-on: ubuntu-24.04 - - steps: - - name: Checkout code - uses: actions/checkout@v7 - with: - submodules: recursive - - - name: Install libaio runtime - run: | - sudo apt-get update -y - # libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older - sudo apt-get install -y libaio1t64 || sudo apt-get install -y libaio1 - echo "=== libaio status (should be present) ===" - dpkg -l | grep -i libaio || true - ldconfig -p | grep libaio || true - shell: bash - - - name: Setup ccache - uses: hendrikmuhs/ccache-action@v1.2 - with: - key: linux-x64-libaio-test - max-size: 150M - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: '3.10' - cache: 'pip' - cache-dependency-path: 'pyproject.toml' - - - name: Set up environment variables - run: | - NPROC=$(nproc 2>/dev/null || echo 2) - echo "NPROC=$NPROC" >> $GITHUB_ENV - echo "$(python -c 'import site; print(site.USER_BASE)')/bin" >> $GITHUB_PATH - shell: bash - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - python -m pip install \ - pybind11==3.0 \ - cmake==3.30.0 \ - ninja==1.11.1 \ - pytest \ - pytest-xdist \ - scikit-build-core \ - setuptools_scm - shell: bash - - - name: Build from source - run: | - cd "$GITHUB_WORKSPACE" - CMAKE_GENERATOR="Ninja" \ - CMAKE_BUILD_PARALLEL_LEVEL="$NPROC" \ - python -m pip install -v . \ - --no-build-isolation \ - --config-settings='cmake.define.BUILD_TOOLS=ON' \ - --config-settings='cmake.define.ENABLE_WERROR=ON' \ - --config-settings='cmake.define.CMAKE_C_COMPILER_LAUNCHER=ccache' \ - --config-settings='cmake.define.CMAKE_CXX_COMPILER_LAUNCHER=ccache' - shell: bash - - - name: Run C++ Tests (w/ libaio) - run: | - cd "$GITHUB_WORKSPACE/build" - cmake --build . --target unittest --parallel $NPROC - shell: bash - - - name: Verify no libaio in test binary - run: | - # Confirm the test binary does NOT have a NEEDED entry for libaio. - # This verifies the dlopen decoupling — libaio is loaded at runtime, - # not linked at build time. - TEST_BIN="$GITHUB_WORKSPACE/build/bin/diskann_searcher_test" - if [ -f "$TEST_BIN" ]; then - ldd "$TEST_BIN" | grep 'libaio' \ - && { echo "ERROR: test binary links libaio directly!"; exit 1; } \ - || echo "OK: test binary does not link libaio (dlopen-based)" - fi - shell: bash - - - name: Run Python Tests (w/ libaio) - run: | - cd "$GITHUB_WORKSPACE" - python -m pytest python/tests/ - shell: bash From fde7813b1130a1ce14a2732142a99ecd4ffea107 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 8 Jul 2026 11:15:17 +0800 Subject: [PATCH 25/46] fix: add ci --- .github/workflows/03-macos-linux-build.yml | 28 ++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index 03fe2cafa..fdb11ce16 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -175,3 +175,31 @@ jobs: ./c_api_index_example ./c_api_optimized_example shell: bash + + # ------------------------------------------------------------------ # + # DiskAnn libaio round: install libaio and re-run tests + # ------------------------------------------------------------------ # + - name: Install libaio runtime + if: matrix.platform == 'linux-x64' + run: | + sudo apt-get update -y + # libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older + sudo apt-get install -y libaio1t64 || sudo apt-get install -y libaio1 + echo "=== libaio status (should be present) ===" + dpkg -l | grep -i libaio || true + ldconfig -p | grep libaio || true + shell: bash + + - name: Run C++ Tests (w/ libaio) + if: matrix.platform == 'linux-x64' + run: | + cd "$GITHUB_WORKSPACE/build" + cmake --build . --target unittest --parallel $NPROC + shell: bash + + - name: Run Python Tests (w/ libaio) + if: matrix.platform == 'linux-x64' + run: | + cd "$GITHUB_WORKSPACE" + python -m pytest python/tests/ + shell: bash From ca89c3043cd335b1721a5d613e6a2f6eb0e14875 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 8 Jul 2026 11:30:07 +0800 Subject: [PATCH 26/46] fix: debug output --- src/binding/python/model/python_collection.cc | 11 ++++++++++- src/db/collection.cc | 7 +++++++ src/include/zvec/db/collection.h | 5 +++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index b1311f119..66944a0df 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -310,7 +310,16 @@ void ZVecPyCollection::bind_dql_methods( "given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. " "Raises KeyError if no HNSW index exists on the column, or " "ValueError if the column's index is not an HNSW index. Intended " - "for introspection and testing only; not part of the stable API."); + "for introspection and testing only; not part of the stable API.") + .def( + "_debug_io_backend_type", + [](const Collection &self) { + const auto result = self.DebugGetIoBackendType(); + return unwrap_expected(result); + }, + "Debug-only: returns the I/O backend type used by DiskAnn. " + "One of 'libaio', 'sync_pread'. Intended for introspection and " + "testing only; not part of the stable API."); } } // namespace zvec \ No newline at end of file diff --git a/src/db/collection.cc b/src/db/collection.cc index 14cb2f4bc..c13096438 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -134,6 +134,8 @@ class CollectionImpl : public Collection { Result DebugGetHnswStorageMode( const std::string &column_name) const override; + Result DebugGetIoBackendType() const override; + private: void prepare_schema(); @@ -1893,6 +1895,11 @@ Result CollectionImpl::DebugGetHnswStorageMode( Status::NotFound("No HNSW index found for column '", column_name, "'")); } +Result CollectionImpl::DebugGetIoBackendType() const { + auto type = ailego::IOBackend::Instance().available(); + return std::string(ailego::IOBackendTypeName(type)); +} + Status CollectionImpl::recovery() { if (!FileHelper::DirectoryExists(path_.c_str())) { return Status::InvalidArgument("collection path{", path_, "} not exist."); diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 83a289c52..bc9f0d63b 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -117,6 +117,11 @@ class Collection { //! introspection and testing; not part of the stable public API. virtual Result DebugGetHnswStorageMode( const std::string &column_name) const = 0; + + //! Debug-only: retrieve the I/O backend type used by DiskAnn. Returns + //! "libaio" or "sync_pread". Intended for introspection and testing; not + //! part of the stable public API. + virtual Result DebugGetIoBackendType() const = 0; }; } // namespace zvec \ No newline at end of file From 01b702c1d5e9ed5c2612c080b561cff22542a3d8 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 8 Jul 2026 14:22:38 +0800 Subject: [PATCH 27/46] fix: run ci --- .github/workflows/03-macos-linux-build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index fdb11ce16..1f6161242 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -190,16 +190,16 @@ jobs: ldconfig -p | grep libaio || true shell: bash - - name: Run C++ Tests (w/ libaio) + - name: Run DiskAnn C++ Tests (w/ libaio) if: matrix.platform == 'linux-x64' run: | cd "$GITHUB_WORKSPACE/build" - cmake --build . --target unittest --parallel $NPROC + ctest -R diskann --output-on-failure --parallel $NPROC shell: bash - - name: Run Python Tests (w/ libaio) + - name: Run DiskAnn Python Tests (w/ libaio) if: matrix.platform == 'linux-x64' run: | cd "$GITHUB_WORKSPACE" - python -m pytest python/tests/ + python -m pytest python/tests/test_collection_diskann.py -v shell: bash From 8dca6a6874024fce34f6e3bcf2263f39736d7e63 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 9 Jul 2026 11:59:42 +0800 Subject: [PATCH 28/46] fix: link core_knn_diskann in index_group_by_test The index_group_by_test uses DiskAnn indexes (guarded by #if DISKANN_SUPPORTED) to verify that group_by search is rejected for unsupported index types. However, the CMakeLists.txt for tests/core/interface/ did not link core_knn_diskann, causing the DiskAnn factory to not be registered at runtime. This resulted in IndexFactory::CreateAndInitIndex returning nullptr for DiskAnn indexes, failing the test assertion. Add core_knn_diskann to the LIBS list, consistent with how other test directories (tests/db/, tests/db/index/, tests/core/algorithm/diskann/) handle DiskAnn linking. --- tests/core/interface/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/interface/CMakeLists.txt b/tests/core/interface/CMakeLists.txt index 83bf8d95e..65c80ffbb 100644 --- a/tests/core/interface/CMakeLists.txt +++ b/tests/core/interface/CMakeLists.txt @@ -8,7 +8,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) NAME ${CC_TARGET} STRICT LIBS zvec_ailego core_framework core_metric core_interface core_knn_flat core_utility core_quantizer sparsehash core_knn_hnsw core_mix_reducer - core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana + core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana core_knn_diskann SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ) From b7d0b0fcca29b0727cab94fd4a431424c025b804 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 9 Jul 2026 15:14:55 +0800 Subject: [PATCH 29/46] fix: fix make --- examples/c++/CMakeLists.txt | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index d58814d52..e3ccd1309 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -15,10 +15,13 @@ endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) +set(ZVEC_SRC_DIR ${ZVEC_ROOT_DIR}/src) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) -# Add include and library search paths -include_directories(${ZVEC_INCLUDE_DIR}) +# Add include and library search paths. +# ZVEC_SRC_DIR is needed because public headers under src/include/ reference +# ailego/ and other internal headers that live directly under src/. +include_directories(${ZVEC_INCLUDE_DIR} ${ZVEC_SRC_DIR}) set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) # Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) From 86806b4b47111d05adc9c3c9fe9b175094f3e1d6 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 13 Jul 2026 11:06:50 +0800 Subject: [PATCH 30/46] support O_DIRECT --- .../algorithm/diskann/diskann_file_reader.cc | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index bed72cbe3..e6e9033ea 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -20,6 +20,11 @@ #include #include #include +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#include +#endif #define MAX_EVENTS 1024 @@ -482,6 +487,57 @@ void LinuxAlignedFileReader::open(const std::string &fname) { ::strerror(errno)); } +#if defined(__APPLE__) || defined(__MACH__) + // macOS has no O_DIRECT. A disk-based index that is served from the page + // cache degrades into an in-memory search, defeating DiskAnn's purpose, so + // the reader always bypasses the cache to emulate Linux O_DIRECT semantics + // (every read hits the device), via three steps: + // 1) F_NOCACHE - reads through this fd are not retained in the UBC; + // 2) F_RDAHEAD=0 - disable read-ahead so random reads are not prefetched; + // 3) mmap + msync(MS_INVALIDATE) - F_NOCACHE only stops NEW pages from + // being cached; XNU's cluster_read still copies from pages already + // resident in the UBC (e.g. warmed while the index was built), so those + // reads would be served from memory. Dropping the file's resident pages + // once at open guarantees subsequent reads go to disk. + if (this->file_desc != -1) { + if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { + LOG_WARN("fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " + "the page cache", + fname.c_str(), errno, ::strerror(errno)); + } else { + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled (page-cache bypass) for %s", + fname.c_str()); + } + + if (::fcntl(this->file_desc, F_RDAHEAD, 0) == -1) { + LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", fname.c_str(), + errno, ::strerror(errno)); + } + + off_t fsize = ::lseek(this->file_desc, 0, SEEK_END); + ::lseek(this->file_desc, 0, SEEK_SET); + if (fsize > 0) { + void *addr = ::mmap(nullptr, static_cast(fsize), PROT_READ, + MAP_SHARED, this->file_desc, 0); + if (addr != MAP_FAILED) { + if (::msync(addr, static_cast(fsize), MS_INVALIDATE) == -1) { + LOG_WARN("msync(MS_INVALIDATE) failed for %s (errno=%d: %s); resident " + "pages may still be served from cache", + fname.c_str(), errno, ::strerror(errno)); + } else { + LOG_INFO( + "DiskAnn macOS: invalidated resident page-cache for %s (%.2f GB)", + fname.c_str(), static_cast(fsize) / 1073741824.0); + } + ::munmap(addr, static_cast(fsize)); + } else { + LOG_WARN("mmap for cache invalidation failed for %s (errno=%d: %s)", + fname.c_str(), errno, ::strerror(errno)); + } + } + } +#endif + LOG_INFO("Opened file : %s", fname.c_str()); } From e9dafa6d8a2c266b14323ea3accc2ec92142089f Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Thu, 16 Jul 2026 17:19:35 +0800 Subject: [PATCH 31/46] upd --- .../algorithm/diskann/diskann_file_reader.cc | 44 ++++--------------- 1 file changed, 9 insertions(+), 35 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index e6e9033ea..c3cbf7e60 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -22,7 +22,6 @@ #include #if defined(__APPLE__) || defined(__MACH__) #include -#include #include #endif @@ -488,24 +487,21 @@ void LinuxAlignedFileReader::open(const std::string &fname) { } #if defined(__APPLE__) || defined(__MACH__) - // macOS has no O_DIRECT. A disk-based index that is served from the page - // cache degrades into an in-memory search, defeating DiskAnn's purpose, so - // the reader always bypasses the cache to emulate Linux O_DIRECT semantics - // (every read hits the device), via three steps: - // 1) F_NOCACHE - reads through this fd are not retained in the UBC; - // 2) F_RDAHEAD=0 - disable read-ahead so random reads are not prefetched; - // 3) mmap + msync(MS_INVALIDATE) - F_NOCACHE only stops NEW pages from - // being cached; XNU's cluster_read still copies from pages already - // resident in the UBC (e.g. warmed while the index was built), so those - // reads would be served from memory. Dropping the file's resident pages - // once at open guarantees subsequent reads go to disk. + // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it + // asks the kernel to minimize caching for I/O through this descriptor. This + // is advisory rather than a guarantee that every read reaches the device. + // Disable read-ahead as well because DiskAnn performs random reads. + // + // Do not mmap the entire index and call msync(MS_INVALIDATE) here. That does + // not provide a reliable global cache eviction guarantee and makes open time + // and virtual-address usage scale with the size of the index. if (this->file_desc != -1) { if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { LOG_WARN("fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " "the page cache", fname.c_str(), errno, ::strerror(errno)); } else { - LOG_INFO("DiskAnn macOS: F_NOCACHE enabled (page-cache bypass) for %s", + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); } @@ -513,28 +509,6 @@ void LinuxAlignedFileReader::open(const std::string &fname) { LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", fname.c_str(), errno, ::strerror(errno)); } - - off_t fsize = ::lseek(this->file_desc, 0, SEEK_END); - ::lseek(this->file_desc, 0, SEEK_SET); - if (fsize > 0) { - void *addr = ::mmap(nullptr, static_cast(fsize), PROT_READ, - MAP_SHARED, this->file_desc, 0); - if (addr != MAP_FAILED) { - if (::msync(addr, static_cast(fsize), MS_INVALIDATE) == -1) { - LOG_WARN("msync(MS_INVALIDATE) failed for %s (errno=%d: %s); resident " - "pages may still be served from cache", - fname.c_str(), errno, ::strerror(errno)); - } else { - LOG_INFO( - "DiskAnn macOS: invalidated resident page-cache for %s (%.2f GB)", - fname.c_str(), static_cast(fsize) / 1073741824.0); - } - ::munmap(addr, static_cast(fsize)); - } else { - LOG_WARN("mmap for cache invalidation failed for %s (errno=%d: %s)", - fname.c_str(), errno, ::strerror(errno)); - } - } } #endif From a2071b7a0780af60291c0c4d37ef907f38b25b5f Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Fri, 17 Jul 2026 14:47:35 +0800 Subject: [PATCH 32/46] clang format --- src/core/algorithm/diskann/diskann_file_reader.cc | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 4b73488d6..20c028f68 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -497,17 +497,17 @@ void LinuxAlignedFileReader::open(const std::string &fname) { // and virtual-address usage scale with the size of the index. if (this->file_desc != -1) { if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { - LOG_WARN("fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " - "the page cache", - fname.c_str(), errno, ::strerror(errno)); + LOG_WARN( + "fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " + "the page cache", + fname.c_str(), errno, ::strerror(errno)); } else { - LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", - fname.c_str()); + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); } if (::fcntl(this->file_desc, F_RDAHEAD, 0) == -1) { - LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", fname.c_str(), - errno, ::strerror(errno)); + LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", + fname.c_str(), errno, ::strerror(errno)); } } #endif From b15065241f222e85a67bb40a97f04f2dfa74e720 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Sat, 18 Jul 2026 17:07:41 +0800 Subject: [PATCH 33/46] upd --- .../algorithm/diskann/diskann_file_reader.cc | 293 ++++++++++++------ .../algorithm/diskann/diskann_file_reader.h | 4 +- .../diskann/diskann_file_reader_test.cc | 122 ++++++++ 3 files changed, 329 insertions(+), 90 deletions(-) create mode 100644 tests/core/algorithm/diskann/diskann_file_reader_test.cc diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 20c028f68..6a1712df9 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -21,6 +21,7 @@ #include #include #if defined(__APPLE__) || defined(__MACH__) +#include #include #include #endif @@ -38,6 +39,8 @@ typedef struct iocb iocb_t; // regardless of which entry point (setup_io_ctx or register_thread) // triggers it first. static std::once_flag g_io_backend_log_once; +#elif defined(__APPLE__) || defined(__MACH__) +static std::once_flag g_io_backend_log_once; #endif int setup_io_ctx(IOContext &ctx) { @@ -57,8 +60,12 @@ int setup_io_ctx(IOContext &ctx) { int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); return ret; #elif defined(__APPLE__) || defined(__MACH__) - // Create a kqueue for this context. On macOS the kqueue is used to - // monitor file descriptor readiness for async-style I/O. + std::call_once(g_io_backend_log_once, [] { + LOG_INFO( + "DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); + }); + // Each macOS I/O context owns a kqueue which receives EVFILT_AIO + // completion notifications for POSIX AIO requests. int kq = ::kqueue(); if (kq == -1) { LOG_ERROR("kqueue() failed in setup_io_ctx; errno=%d, %s", errno, @@ -110,101 +117,207 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { } #if defined(__APPLE__) || defined(__MACH__) -// Execute batch I/O on macOS using kqueue to monitor file descriptor -// readiness and pread for actual data transfer. -// -// On macOS, regular file descriptors are almost always "readable", so -// kqueue's primary value here is providing the same async I/O interface -// as Linux's libaio. For each read request we: -// 1. Attempt a non-blocking pread (via O_NONBLOCK on the fd). -// 2. If EAGAIN, wait on kqueue for EVFILT_READ readiness, then retry. -// 3. Fall back to blocking pread if kqueue encounters an error. +// POSIX AIO on Darwin accepts a maximum of AIO_LISTIO_MAX requests in the +// interfaces used to wait for a batch. Keeping the same bound here also avoids +// exhausting the process-wide AIO request limit when several search threads +// run concurrently. +static constexpr size_t kMacAioBatchSize = AIO_LISTIO_MAX; + +static bool validate_aio_result(int aio_err, int64_t result, + const AlignedRead &req) { + if (aio_err == 0 && result == static_cast(req.len)) { + return true; + } + + LOG_WARN( + "macOS aio request failed; aio_error=%d, %s, result=%ld, " + "expected=%lu, offset=%lu", + aio_err, aio_err == 0 ? "success" : ::strerror(aio_err), (long)result, + (unsigned long)req.len, (unsigned long)req.offset); + return false; +} + +// Once the kqueue has been closed, its AIO knotes are detached and the normal +// POSIX aio_error()/aio_return() cleanup path is safe again. This keeps live +// requests and destination buffers from escaping the call if kevent64 fails. +static bool drain_aio_after_kqueue_failure( + std::vector &cbs, const std::vector &read_reqs, + size_t req_begin, const std::vector &submitted, + std::vector &completed, size_t submitted_count) { + size_t completed_count = 0; + bool all_ok = true; + for (size_t i = 0; i < cbs.size(); ++i) { + if (submitted[i] != 0 && completed[i] != 0) { + ++completed_count; + } + } + + while (completed_count < submitted_count) { + bool made_progress = false; + std::vector pending; + pending.reserve(submitted_count - completed_count); + + for (size_t i = 0; i < cbs.size(); ++i) { + if (submitted[i] == 0 || completed[i] != 0) { + continue; + } + + int aio_err = ::aio_error(&cbs[i]); + if (aio_err == EINPROGRESS) { + pending.push_back(&cbs[i]); + continue; + } + + ssize_t result = ::aio_return(&cbs[i]); + all_ok = validate_aio_result(aio_err, result, + read_reqs[req_begin + i]) && + all_ok; + completed[i] = 1; + ++completed_count; + made_progress = true; + } + + if (completed_count == submitted_count || made_progress) { + continue; + } + + int ret; + do { + ret = ::aio_suspend(pending.data(), static_cast(pending.size()), + nullptr); + } while (ret == -1 && errno == EINTR); + if (ret == -1) { + LOG_ERROR("aio_suspend failed while draining requests; errno=%d, %s", + errno, ::strerror(errno)); + struct timespec pause = {0, 1000000}; // 1 ms + ::nanosleep(&pause, nullptr); + all_ok = false; + } + } + + return all_ok; +} + +// Submit a batch through POSIX aio_read(). SIGEV_KEVENT makes XNU attach an +// EVFILT_AIO one-shot event to the supplied kqueue for each request. kqueue is +// therefore used for completion notification, not regular-file readiness. // -// The kqueue fd is passed in as the IOContext. If no valid kqueue is -// available, we fall back to plain blocking pread. -static int execute_io_kqueue(int kq, int fd, - std::vector &read_reqs) { - // If no kqueue available, fall back to blocking pread. +// Darwin's EVFILT_AIO filter is also the completion reaper: fetching the event +// removes the request from the process AIO done queue. The error and return +// values are carried in kevent64_s::ext[0] and ext[1], respectively. Calling +// aio_error()/aio_return() after fetching the event would incorrectly report +// EINVAL because the kernel has already released that request. +static int execute_io_aio_kqueue(int &kq, int fd, + std::vector &read_reqs, + uint64_t n_retries) { if (kq < 0) { return execute_io_pread(fd, read_reqs); } - // Register the file descriptor with the kqueue for read events. - // EV_CLEAR gives edge-triggered semantics so we only get notified - // when new data becomes available. - struct kevent ke; - EV_SET(&ke, fd, EVFILT_READ, EV_ADD | EV_CLEAR, 0, 0, nullptr); - - for (auto &req : read_reqs) { - while (true) { - ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); - - if (bytes_read > 0) { - // Successfully read data; verify full read. - if ((size_t)bytes_read != req.len) { - // Partial read — retry for the remaining bytes. - // Update offset and buffer to read the rest. - char *buf_ptr = static_cast(req.buf) + bytes_read; - uint64_t new_offset = req.offset + bytes_read; - size_t remaining = req.len - bytes_read; - while (remaining > 0) { - ssize_t n = ::pread(fd, buf_ptr, remaining, new_offset); - if (n < 0) { - if (errno == EINTR) continue; - LOG_ERROR("pread retry failed; errno=%d, %s", errno, - ::strerror(errno)); - return IndexError_Runtime; - } - if (n == 0) break; - buf_ptr += n; - new_offset += n; - remaining -= n; - } - if (remaining > 0) { - LOG_ERROR("pread short read after retry; remaining=%zu", remaining); - return IndexError_Runtime; - } + for (size_t req_begin = 0; req_begin < read_reqs.size(); + req_begin += kMacAioBatchSize) { + size_t n_ops = + std::min(kMacAioBatchSize, read_reqs.size() - req_begin); + std::vector cbs(n_ops); + std::vector submitted(n_ops, 0); + std::vector completed(n_ops, 0); + size_t submitted_count = 0; + bool submission_ok = true; + + for (size_t i = 0; i < n_ops; ++i) { + std::memset(&cbs[i], 0, sizeof(cbs[i])); + cbs[i].aio_fildes = fd; + cbs[i].aio_offset = static_cast(read_reqs[req_begin + i].offset); + cbs[i].aio_buf = read_reqs[req_begin + i].buf; + cbs[i].aio_nbytes = read_reqs[req_begin + i].len; + cbs[i].aio_reqprio = 0; + cbs[i].aio_lio_opcode = LIO_READ; + cbs[i].aio_sigevent.sigev_notify = SIGEV_KEVENT; + cbs[i].aio_sigevent.sigev_signo = kq; + cbs[i].aio_sigevent.sigev_value.sival_ptr = &cbs[i]; + + uint64_t tries = 0; + while (::aio_read(&cbs[i]) == -1) { + int submit_errno = errno; + if ((submit_errno == EINTR || submit_errno == EAGAIN) && + tries < n_retries) { + ++tries; + continue; } - break; // Success, move to next request. + LOG_WARN( + "aio_read submission failed; errno=%d, %s, offset=%lu, " + "len=%lu; falling back to pread after draining submitted I/O", + submit_errno, ::strerror(submit_errno), + (unsigned long)read_reqs[req_begin + i].offset, + (unsigned long)read_reqs[req_begin + i].len); + submission_ok = false; + break; } - - if (bytes_read == 0) { - // EOF — should not happen for a valid index file. - LOG_ERROR("pread returned 0 (EOF); offset=%lu, len=%lu", - (unsigned long)req.offset, (unsigned long)req.len); - return IndexError_Runtime; + if (!submission_ok) { + break; } + submitted[i] = 1; + ++submitted_count; + } - // bytes_read == -1, error - if (errno == EINTR) { - continue; // Retry on signal. + size_t completed_count = 0; + bool all_ok = true; + bool kqueue_ok = true; + while (completed_count < submitted_count) { + struct kevent64_s events[AIO_LISTIO_MAX]; + int n_events = ::kevent64(kq, nullptr, 0, events, + static_cast(n_ops), 0, nullptr); + if (n_events == -1) { + if (errno == EINTR) { + continue; + } + LOG_ERROR("kevent failed while waiting for AIO; errno=%d, %s", errno, + ::strerror(errno)); + kqueue_ok = false; + break; } - if (errno == EAGAIN || errno == EWOULDBLOCK) { - // Data not ready — wait on kqueue for readability. - struct kevent events[1]; - struct timespec ts; - ts.tv_sec = 5; // 5 second timeout as a safety net. - ts.tv_nsec = 0; - int n_ev = ::kevent(kq, &ke, 1, events, 1, &ts); - if (n_ev < 0) { - if (errno == EINTR) continue; - // kqueue error — fall back to blocking pread for this request. - LOG_WARN("kevent failed; errno=%d, %s, falling back to pread", errno, - ::strerror(errno)); - return execute_io_pread(fd, read_reqs); + + for (int event_index = 0; event_index < n_events; ++event_index) { + size_t request_index = n_ops; + for (size_t i = 0; i < n_ops; ++i) { + if (events[event_index].ident == + reinterpret_cast(&cbs[i])) { + request_index = i; + break; + } } - if (n_ev == 0) { - // Timeout — fall back to blocking pread. - LOG_WARN("kqueue timeout, falling back to pread"); - return execute_io_pread(fd, read_reqs); + + if (request_index == n_ops || submitted[request_index] == 0 || + completed[request_index] != 0 || + events[event_index].filter != EVFILT_AIO) { + LOG_WARN("ignoring an unknown or duplicate macOS AIO event"); + continue; } - // Event triggered — retry pread. - continue; + + int aio_err = static_cast(events[event_index].ext[0]); + int64_t result = static_cast(events[event_index].ext[1]); + const AlignedRead &req = read_reqs[req_begin + request_index]; + all_ok = validate_aio_result(aio_err, result, req) && all_ok; + completed[request_index] = 1; + ++completed_count; } + if (!kqueue_ok) { + break; + } + } - // Other error — fall back to blocking pread. - LOG_ERROR("pread failed; errno=%d, %s, falling back to pread", errno, - ::strerror(errno)); + if (!kqueue_ok) { + // Detach the EVFILT_AIO knotes before using aio_return() to reap the + // requests through the POSIX API, and disable AIO for this context. + ::close(kq); + kq = -1; + all_ok = drain_aio_after_kqueue_failure( + cbs, read_reqs, req_begin, submitted, completed, + submitted_count) && + all_ok; + } + + if (!submission_ok || !kqueue_ok || !all_ok) { return execute_io_pread(fd, read_reqs); } } @@ -213,7 +326,7 @@ static int execute_io_kqueue(int kq, int fd, } #endif // __APPLE__ -int execute_io(IOContext ctx, int fd, std::vector &read_reqs, +int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) if (ailego::IOBackend::Instance().available() == @@ -295,9 +408,9 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, return 0; #elif defined(__APPLE__) || defined(__MACH__) - // On macOS, use kqueue-based I/O. The IOContext (ctx) is a kqueue fd. - (void)n_retries; - return execute_io_kqueue(ctx, fd, read_reqs); + // On macOS, submit POSIX AIO and use the IOContext kqueue for EVFILT_AIO + // completion notifications. + return execute_io_aio_kqueue(ctx, fd, read_reqs, n_retries); #else (void)ctx; (void)n_retries; @@ -384,7 +497,11 @@ void LinuxAlignedFileReader::register_thread() { LOG_ERROR("kqueue() failed in register_thread; errno=%d, %s", errno, ::strerror(errno)); } else { - LOG_INFO("allocating kqueue ctx: %d", kq); + std::call_once(g_io_backend_log_once, [] { + LOG_INFO( + "DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); + }); + LOG_INFO("allocating POSIX AIO kqueue ctx: %d", kq); ctx = kq; ctx_map[thread_id] = ctx; } diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index c2702c399..1023fc40c 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -94,8 +94,8 @@ class AlignedFileReader { // Reader implementation used on all supported platforms. // On Linux (x86_64 and ARM64) it uses libaio for asynchronous batch I/O. -// On macOS (including ARM/Apple Silicon) it uses kqueue to monitor file -// descriptor readiness and pread for actual data transfer. +// On macOS (including ARM/Apple Silicon) it submits POSIX AIO requests and +// uses kqueue EVFILT_AIO events for completion notification. class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc new file mode 100644 index 000000000..627fed07a --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -0,0 +1,122 @@ +// 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. + +#include "diskann_file_reader.h" +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core; + +namespace { + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + if (fd_ >= 0) { + ::close(fd_); + } + ::unlink(path_); + } + + int fd() const { + return fd_; + } + + const char *path() const { + return path_; + } + + void close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + } + + private: + char path_[64] = "DiskAnnFileReaderTest.XXXXXX"; + int fd_; +}; + +} // namespace + +TEST(DiskAnnFileReaderTest, BatchAlignedReads) { + constexpr size_t kPageSize = 4096; + constexpr size_t kPageCount = 32; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize * kPageCount); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * kPageSize, + static_cast(page + 1), kPageSize); + } + + size_t written = 0; + while (written < source.size()) { + ssize_t ret = ::pwrite(file.fd(), source.data() + written, + source.size() - written, written); + ASSERT_GT(ret, 0); + written += static_cast(ret); + } + ASSERT_EQ(::fsync(file.fd()), 0); + file.close(); + + void *raw_buffer = nullptr; + ASSERT_EQ(::posix_memalign(&raw_buffer, kPageSize, source.size()), 0); + std::unique_ptr output(raw_buffer, &std::free); + + std::vector requests; + requests.reserve(kPageCount); + for (size_t i = 0; i < kPageCount; ++i) { + // Use a non-sequential page order and more than AIO_LISTIO_MAX requests + // so the macOS backend has to submit multiple batches. + size_t source_page = (i * 7) % kPageCount; + requests.emplace_back(source_page * kPageSize, kPageSize, + static_cast(output.get()) + + i * kPageSize); + } + + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx{ +#if defined(__APPLE__) || defined(__MACH__) + -1 +#else + 0 +#endif + }; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_EQ(reader.read(requests, ctx), 0); + + for (size_t i = 0; i < kPageCount; ++i) { + size_t source_page = (i * 7) % kPageCount; + const auto *page = + static_cast(output.get()) + i * kPageSize; + for (size_t byte = 0; byte < kPageSize; ++byte) { + ASSERT_EQ(page[byte], static_cast(source_page + 1)); + } + } + + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} From 3a1d16147060f91788ab0121ab159741fc0876bc Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 15:38:21 +0800 Subject: [PATCH 34/46] fix --- src/binding/python/model/python_collection.cc | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 7349c46db..4468a53ff 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -313,16 +313,7 @@ void ZVecPyCollection::bind_dql_methods( "given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. " "Raises KeyError if no HNSW index exists on the column, or " "ValueError if the column's index is not an HNSW index. Intended " - "for introspection and testing only; not part of the stable API.") - .def( - "_debug_io_backend_type", - [](const Collection &self) { - const auto result = self.DebugGetIoBackendType(); - return unwrap_expected(result); - }, - "Debug-only: returns the I/O backend type used by DiskAnn. " - "One of 'libaio', 'sync_pread'. Intended for introspection and " - "testing only; not part of the stable API."); + "for introspection and testing only; not part of the stable API."); } } // namespace zvec From a15f953dbd9b817ef5a910ae1af4fdc60b58555f Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 16:20:40 +0800 Subject: [PATCH 35/46] add Posix AIO --- python/tests/test_typing.py | 4 +- python/zvec/__init__.pyi | 1 + python/zvec/typing/__init__.pyi | 6 ++- src/ailego/io/io_backend_def.h | 35 ++++++++++++--- .../python/model/common/python_config.cc | 5 ++- src/binding/python/typing/python_type.cc | 6 ++- src/core/interface/indexes/diskann_index.cc | 5 ++- src/db/collection.cc | 10 ++++- src/include/zvec/ailego/io/io_backend.h | 8 ++-- src/include/zvec/c_api.h | 10 +++-- src/include/zvec/core/interface/index.h | 3 +- src/include/zvec/db/collection.h | 6 +-- tests/ailego/io/io_backend_test.cc | 43 +++++++++++++++++++ 13 files changed, 114 insertions(+), 28 deletions(-) create mode 100644 tests/ailego/io/io_backend_test.cc diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py index d566d5efc..ab821f66c 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -34,6 +34,7 @@ (DataType.FLOAT, "FLOAT"), (IndexType.HNSW, "HNSW"), (IOBackendType.PREAD, "PREAD"), + (IOBackendType.POSIX_AIO, "POSIX_AIO"), (MetricType.COSINE, "COSINE"), (QuantizeType.INT8, "INT8"), (StatusCode.OK, "OK"), @@ -49,6 +50,7 @@ def test_enum_names(member, name): (DataType.FLOAT, 8), (IndexType.HNSW, 1), (IOBackendType.PREAD, 0), + (IOBackendType.POSIX_AIO, 2), (MetricType.COSINE, 3), (QuantizeType.INT8, 2), (StatusCode.OK, 0), @@ -112,7 +114,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", "POSIX_AIO"]) def test_io_backend_type_has_member(member): assert member in IOBackendType.__members__ diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 09e828dae..e31b68da1 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -52,6 +52,7 @@ 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). + On macOS this is IOBackendType.POSIX_AIO. On Linux this is IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise.""" def io_backend_description() -> str: diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index fab03f4b9..b78776d7f 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -130,6 +130,7 @@ class IOBackendType: - PREAD: Synchronous pread() — no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). + - POSIX_AIO: macOS POSIX AIO with kqueue completion notifications. Examples: >>> from zvec.typing import IOBackendType @@ -142,13 +143,16 @@ class IOBackendType: PREAD LIBAIO + + POSIX_AIO """ LIBAIO: typing.ClassVar[IOBackendType] # value = + POSIX_AIO: typing.ClassVar[IOBackendType] # value = PREAD: typing.ClassVar[IOBackendType] # value = __members__: typing.ClassVar[ dict[str, IOBackendType] - ] # value = {'PREAD': , 'LIBAIO': } + ] # value = {'PREAD': , 'LIBAIO': , 'POSIX_AIO': } def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d795c3046..1b79f4c85 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -40,6 +40,8 @@ inline const char *IOBackendTypeName(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio"; + case IOBackendType::kPosixAio: + return "posix_aio"; case IOBackendType::kPread: return "pread"; } @@ -52,6 +54,8 @@ inline const char *IOBackendDescription(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; + case IOBackendType::kPosixAio: + return "macOS POSIX AIO backend with kqueue completion notifications."; case IOBackendType::kPread: return "No async I/O backend available. Install libaio (e.g. " "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " @@ -63,8 +67,8 @@ inline const char *IOBackendDescription(IOBackendType type) { // Singleton that loads and queries an I/O backend on demand. // -// available() (no arg) tries the best backend with priority (libaio > pread) -// and returns the loaded backend type. +// available() (no arg) returns the platform backend. On Linux it tries libaio +// before falling back to pread; on macOS POSIX AIO is provided by the system. // available(IOBackendType) tries a specific backend. // Use type() / name() to query the loaded backend without triggering a load. class IOBackend { @@ -74,8 +78,8 @@ class IOBackend { return instance; } - // Try to load the best available backend (libaio > pread). - // Returns the loaded backend type. + // Return the selected platform backend. Linux probes libaio before falling + // back to pread; macOS uses the system POSIX AIO backend. // Idempotent — if already loaded, returns immediately. IOBackendType available() { if (type_ != IOBackendType::kPread) { @@ -84,13 +88,20 @@ class IOBackend { 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). + // Try to load the requested backend. Returns the selected backend type, + // which may differ from the request when it is unavailable on this platform. // Idempotent — if the same backend is already loaded, returns immediately. IOBackendType available(IOBackendType requested) { if (type_ == requested && type_ != IOBackendType::kPread) { return type_; } +#if defined(__APPLE__) || defined(__MACH__) + // POSIX AIO and EVFILT_AIO are provided by Darwin; no user-installed + // runtime dependency needs to be probed. A failure to create a particular + // kqueue context is handled by the DiskAnn reader as an operational error. + type_ = IOBackendType::kPosixAio; + return type_; +#endif #if defined(__linux) || defined(__linux__) if (requested == IOBackendType::kLibAio) { if (LibAioLoader::Instance().load() && @@ -112,6 +123,10 @@ class IOBackend { return available() == IOBackendType::kLibAio; } + bool is_posix_aio() { + return available() == IOBackendType::kPosixAio; + } + // Returns the loaded backend type. IOBackendType type() const { return type_; @@ -130,7 +145,13 @@ class IOBackend { private: IOBackend() = default; - IOBackendType type_{IOBackendType::kPread}; + IOBackendType type_{ +#if defined(__APPLE__) || defined(__MACH__) + IOBackendType::kPosixAio +#else + IOBackendType::kPread +#endif + }; }; } // namespace ailego diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f48b..3f9f073fc 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -228,7 +228,8 @@ 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.LIBAIO if libaio is available, " + "On macOS this is IOBackendType.POSIX_AIO; on Linux this is " + "IOBackendType.LIBAIO when libaio is available and " "IOBackendType.PREAD otherwise."); // Returns a human-readable description of the I/O backend, including @@ -242,4 +243,4 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index d20cac755..87ae79489 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -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(). +- POSIX_AIO: macOS POSIX AIO with kqueue completion notifications. Examples: >>> from zvec.typing import IOBackendType @@ -153,7 +154,8 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. IOBackendType.LIBAIO )pbdoc") .value("PREAD", ailego::IOBackendType::kPread) - .value("LIBAIO", ailego::IOBackendType::kLibAio); + .value("LIBAIO", ailego::IOBackendType::kLibAio) + .value("POSIX_AIO", ailego::IOBackendType::kPosixAio); } void ZVecPyTyping::bind_status(py::module_ &m) { @@ -245,4 +247,4 @@ Construct a status with the given code and optional message. }); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 0553dd145..e719c8ae6 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -273,14 +273,15 @@ int DiskAnnIndex::Merge(const std::vector &indexes, } ailego::IOBackendType DiskAnnIndex::io_backend_type() const { - auto &backend = ailego::IOBackend::Instance(); - ailego::IOBackendType type = backend.type(); + ailego::IOBackendType type = ailego::current_io_backend_type(); +#if defined(__linux__) || defined(__linux) if (type == ailego::IOBackendType::kPread) { LOG_WARN( "Only synchronous pread() is available. Install libaio " "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " "for async I/O support — performance will be degraded without it."); } +#endif return type; } diff --git a/src/db/collection.cc b/src/db/collection.cc index 7d6c07c98..a6fc88834 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -1837,7 +1837,15 @@ Result CollectionImpl::DebugGetHnswStorageMode( Result CollectionImpl::DebugGetIoBackendType() const { const auto type = ailego::current_io_backend_type(); - return type == ailego::IOBackendType::kLibAio ? "libaio" : "sync_pread"; + switch (type) { + case ailego::IOBackendType::kLibAio: + return "libaio"; + case ailego::IOBackendType::kPosixAio: + return "posix_aio"; + case ailego::IOBackendType::kPread: + return "sync_pread"; + } + return "unknown"; } Status CollectionImpl::recovery() { diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 008b01514..1ef168095 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -30,12 +30,14 @@ namespace ailego { // Supported I/O backend types. 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, // Linux libaio loaded at runtime via dlopen() + kPosixAio = 2 // macOS POSIX AIO with kqueue completion }; // Returns the currently active I/O backend type. -// Triggers backend initialization on first call (libaio > pread). +// On Linux, triggers backend initialization on first call (libaio > pread). +// On macOS, returns the system POSIX AIO backend. IOBackendType current_io_backend_type(); // Returns a human-readable description of the currently active I/O backend. diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..00afc9876 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,14 +789,16 @@ 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_POSIX_AIO \ + 2 /**< macOS POSIX AIO with kqueue completion */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. * - * Pure introspection \u2014 no side effects, no install hints. + * On Linux, the first call may probe for libaio. No configuration is changed + * and no installation hints are emitted. * - * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD). + * @return zvec_io_backend_type_t The selected backend type. */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -805,7 +807,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". + * "libaio", "posix_aio", "pread", or "unknown". */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_type_name(zvec_io_backend_type_t type); diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index f5dd797cc..3fe89ff96 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -377,8 +377,7 @@ class DiskAnnIndex : public Index { public: DiskAnnIndex() = default; - // Returns the I/O backend type currently loaded for DiskAnn async disk reads. - // If only sync_pread is available, logs a hint to install libaio. + // Returns the I/O backend type selected for DiskAnn disk reads. ailego::IOBackendType io_backend_type() const; protected: diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index bc9f0d63b..11c04edb1 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -119,9 +119,9 @@ class Collection { const std::string &column_name) const = 0; //! Debug-only: retrieve the I/O backend type used by DiskAnn. Returns - //! "libaio" or "sync_pread". Intended for introspection and testing; not - //! part of the stable public API. + //! "libaio", "posix_aio", or "sync_pread". Intended for introspection and + //! testing; not part of the stable public API. virtual Result DebugGetIoBackendType() const = 0; }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/tests/ailego/io/io_backend_test.cc b/tests/ailego/io/io_backend_test.cc new file mode 100644 index 000000000..66eccffb5 --- /dev/null +++ b/tests/ailego/io/io_backend_test.cc @@ -0,0 +1,43 @@ +// 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. + +#include +#include +#include +#include + +namespace zvec { +namespace ailego { + +TEST(IOBackendTest, StableBackendValues) { + EXPECT_EQ(static_cast(IOBackendType::kPread), 0U); + EXPECT_EQ(static_cast(IOBackendType::kLibAio), 1U); + EXPECT_EQ(static_cast(IOBackendType::kPosixAio), 2U); +} + +TEST(IOBackendTest, BackendNames) { + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kPread), "pread"); + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kLibAio), "libaio"); + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kPosixAio), "posix_aio"); +} + +#if defined(__APPLE__) || defined(__MACH__) +TEST(IOBackendTest, MacOSUsesPosixAio) { + EXPECT_EQ(current_io_backend_type(), IOBackendType::kPosixAio); + EXPECT_NE(current_io_backend_description().find("kqueue"), std::string::npos); +} +#endif + +} // namespace ailego +} // namespace zvec From 397e8702df8eacb1bad9657c02927cd3edeb43cb Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 16:28:13 +0800 Subject: [PATCH 36/46] fix --- tests/c/c_api_test.c | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..a790e2383 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -130,6 +130,30 @@ void test_version_functions(void) { TEST_END(); } +void test_io_backend_functions(void) { + TEST_START(); + + zvec_io_backend_type_t type = zvec_get_io_backend_type(); + const char *name = zvec_get_io_backend_type_name(type); + const char *description = zvec_get_io_backend_description(); + + TEST_ASSERT(name != NULL); + TEST_ASSERT(description != NULL); +#if defined(__APPLE__) || defined(__MACH__) + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_POSIX_AIO); + TEST_ASSERT(strcmp(name, "posix_aio") == 0); +#elif defined(__linux__) || defined(__linux) + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_LIBAIO || + type == ZVEC_IO_BACKEND_TYPE_PREAD); +#else + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_PREAD); +#endif + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name(UINT32_MAX), "unknown") == + 0); + + TEST_END(); +} + void test_error_handling_functions(void) { TEST_START(); @@ -6375,6 +6399,7 @@ int main(void) { printf("Cleanup completed.\n\n"); test_version_functions(); + test_io_backend_functions(); test_error_handling_functions(); test_zvec_config(); test_zvec_initialize(); From 960f588dd03152b94e92d03606fcb9e42907335d Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 17:33:58 +0800 Subject: [PATCH 37/46] fix --- src/ailego/io/io_backend_def.h | 18 +++- src/core/algorithm/diskann/diskann_context.h | 2 +- .../algorithm/diskann/diskann_file_reader.cc | 92 +++++++++++-------- .../algorithm/diskann/diskann_file_reader.h | 16 +++- src/core/algorithm/diskann/diskann_indexer.h | 2 +- src/core/interface/indexes/diskann_index.cc | 14 --- src/db/collection.cc | 15 --- src/include/zvec/core/interface/index.h | 4 - src/include/zvec/db/collection.h | 5 - tests/ailego/io/io_backend_test.cc | 12 ++- tests/c/c_api_test.c | 6 +- .../diskann/diskann_file_reader_test.cc | 41 +++++++-- .../db/crash_recovery/write_recovery_test.cc | 4 +- 13 files changed, 132 insertions(+), 99 deletions(-) diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index 1b79f4c85..19118800a 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -32,6 +32,10 @@ #include #include +#if defined(__APPLE__) +#include +#endif + namespace zvec { namespace ailego { @@ -49,7 +53,7 @@ inline const char *IOBackendTypeName(IOBackendType type) { } // Returns a human-readable description for the given backend type. -// When the backend is kPread, includes installation guidance for libaio. +// On Linux, the kPread description includes installation guidance for libaio. inline const char *IOBackendDescription(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: @@ -57,10 +61,15 @@ inline const char *IOBackendDescription(IOBackendType type) { case IOBackendType::kPosixAio: return "macOS POSIX AIO backend with kqueue completion notifications."; case IOBackendType::kPread: +#if defined(__linux) || defined(__linux__) return "No async I/O backend available. Install libaio (e.g. " "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " "and retry. DiskAnn will fall back to synchronous pread() \u2014 " "performance will be degraded."; +#else + return "Synchronous pread() backend; no async I/O backend is " + "available."; +#endif } return "Unknown I/O backend."; } @@ -95,10 +104,11 @@ class IOBackend { if (type_ == requested && type_ != IOBackendType::kPread) { return type_; } -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX // POSIX AIO and EVFILT_AIO are provided by Darwin; no user-installed // runtime dependency needs to be probed. A failure to create a particular - // kqueue context is handled by the DiskAnn reader as an operational error. + // kqueue context is handled by the DiskAnn reader by falling back to + // synchronous pread(). type_ = IOBackendType::kPosixAio; return type_; #endif @@ -146,7 +156,7 @@ class IOBackend { IOBackend() = default; IOBackendType type_{ -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX IOBackendType::kPosixAio #else IOBackendType::kPread diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index e5729a19f..1f4bc75ac 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -350,7 +350,7 @@ class DiskAnnContext : public IndexContext, std::map group_topk_heaps_{}; IOContext io_ctx_{ -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX -1 #else 0 diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 6a1712df9..16a880506 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -20,7 +20,7 @@ #include #include #include -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX #include #include #include @@ -39,7 +39,7 @@ typedef struct iocb iocb_t; // regardless of which entry point (setup_io_ctx or register_thread) // triggers it first. static std::once_flag g_io_backend_log_once; -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX static std::once_flag g_io_backend_log_once; #endif @@ -59,18 +59,20 @@ int setup_io_ctx(IOContext &ctx) { } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); return ret; -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX std::call_once(g_io_backend_log_once, [] { - LOG_INFO( - "DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); }); // Each macOS I/O context owns a kqueue which receives EVFILT_AIO // completion notifications for POSIX AIO requests. + ctx = -1; int kq = ::kqueue(); if (kq == -1) { - LOG_ERROR("kqueue() failed in setup_io_ctx; errno=%d, %s", errno, - ::strerror(errno)); - return IndexError_Runtime; + LOG_WARN( + "kqueue() failed in setup_io_ctx; errno=%d, %s; falling back to " + "synchronous pread", + errno, ::strerror(errno)); + return 0; } ctx = kq; return 0; @@ -87,7 +89,7 @@ int destroy_io_ctx(IOContext &ctx) { } int ret = LibAioLoader::Instance().io_destroy(ctx); return ret; -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX if (ctx >= 0) { ::close(ctx); ctx = -1; @@ -98,7 +100,11 @@ int destroy_io_ctx(IOContext &ctx) { #endif } -static int execute_io_pread(int fd, std::vector &read_reqs) { +static int execute_io_pread(int fd, std::vector &read_reqs, + ailego::IOBackendType *used_backend = nullptr) { + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kPread; + } for (auto &req : read_reqs) { ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); if (bytes_read < 0) { @@ -116,7 +122,7 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX // POSIX AIO on Darwin accepts a maximum of AIO_LISTIO_MAX requests in the // interfaces used to wait for a batch. Keeping the same bound here also avoids // exhausting the process-wide AIO request limit when several search threads @@ -169,8 +175,7 @@ static bool drain_aio_after_kqueue_failure( } ssize_t result = ::aio_return(&cbs[i]); - all_ok = validate_aio_result(aio_err, result, - read_reqs[req_begin + i]) && + all_ok = validate_aio_result(aio_err, result, read_reqs[req_begin + i]) && all_ok; completed[i] = 1; ++completed_count; @@ -209,15 +214,15 @@ static bool drain_aio_after_kqueue_failure( // EINVAL because the kernel has already released that request. static int execute_io_aio_kqueue(int &kq, int fd, std::vector &read_reqs, - uint64_t n_retries) { + uint64_t n_retries, + ailego::IOBackendType *used_backend) { if (kq < 0) { - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } for (size_t req_begin = 0; req_begin < read_reqs.size(); req_begin += kMacAioBatchSize) { - size_t n_ops = - std::min(kMacAioBatchSize, read_reqs.size() - req_begin); + size_t n_ops = std::min(kMacAioBatchSize, read_reqs.size() - req_begin); std::vector cbs(n_ops); std::vector submitted(n_ops, 0); std::vector completed(n_ops, 0); @@ -265,8 +270,8 @@ static int execute_io_aio_kqueue(int &kq, int fd, bool kqueue_ok = true; while (completed_count < submitted_count) { struct kevent64_s events[AIO_LISTIO_MAX]; - int n_events = ::kevent64(kq, nullptr, 0, events, - static_cast(n_ops), 0, nullptr); + int n_events = ::kevent64(kq, nullptr, 0, events, static_cast(n_ops), + 0, nullptr); if (n_events == -1) { if (errno == EINTR) { continue; @@ -311,27 +316,31 @@ static int execute_io_aio_kqueue(int &kq, int fd, // requests through the POSIX API, and disable AIO for this context. ::close(kq); kq = -1; - all_ok = drain_aio_after_kqueue_failure( - cbs, read_reqs, req_begin, submitted, completed, - submitted_count) && - all_ok; + all_ok = + drain_aio_after_kqueue_failure(cbs, read_reqs, req_begin, submitted, + completed, submitted_count) && + all_ok; } if (!submission_ok || !kqueue_ok || !all_ok) { - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } } + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kPosixAio; + } return 0; } #endif // __APPLE__ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, - uint64_t n_retries = 0) { + uint64_t n_retries = 0, + ailego::IOBackendType *used_backend = nullptr) { #if (defined(__linux) || defined(__linux__)) if (ailego::IOBackend::Instance().available() == ailego::IOBackendType::kPread) { - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); @@ -368,7 +377,7 @@ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, "io_submit failed; returned: %d, expected=%lu. falling back to " "pread", ret, n_ops); - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } // Phase 2: io_getevents with retry (never re-submits). @@ -387,7 +396,7 @@ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, "io_getevents failed; returned: %d, expected=%lu, errno=%d, %s, " "falling back to pread", ret, n_ops, errno, ::strerror(-ret)); - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } // Phase 3: verify each completed read (res must equal requested length). @@ -402,19 +411,22 @@ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, } } if (!all_ok) { - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } } + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kLibAio; + } return 0; -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX // On macOS, submit POSIX AIO and use the IOContext kqueue for EVFILT_AIO // completion notifications. - return execute_io_aio_kqueue(ctx, fd, read_reqs, n_retries); + return execute_io_aio_kqueue(ctx, fd, read_reqs, n_retries, used_backend); #else (void)ctx; (void)n_retries; - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); #endif } @@ -483,7 +495,7 @@ void LinuxAlignedFileReader::register_thread() { ctx_map[thread_id] = ctx; } lk.unlock(); -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX auto thread_id = std::this_thread::get_id(); std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { @@ -498,8 +510,7 @@ void LinuxAlignedFileReader::register_thread() { ::strerror(errno)); } else { std::call_once(g_io_backend_log_once, [] { - LOG_INFO( - "DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); }); LOG_INFO("allocating POSIX AIO kqueue ctx: %d", kq); ctx = kq; @@ -531,7 +542,7 @@ void LinuxAlignedFileReader::deregister_thread() { LibAioLoader::Instance().io_destroy(ctx); } LOG_INFO("returned ctx from thread"); -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX auto thread_id = std::this_thread::get_id(); IOContext ctx; @@ -565,7 +576,7 @@ void LinuxAlignedFileReader::deregister_all_threads() { } } ctx_map.clear(); -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX std::unique_lock lk(ctx_mut); for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; @@ -603,7 +614,7 @@ void LinuxAlignedFileReader::open(const std::string &fname) { ::strerror(errno)); } -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it // asks the kernel to minimize caching for I/O through this descriptor. This // is advisory rather than a guarantee that every read reaches the device. @@ -640,7 +651,8 @@ void LinuxAlignedFileReader::close() { } int LinuxAlignedFileReader::read(std::vector &read_reqs, - IOContext &ctx, bool async) { + IOContext &ctx, bool async, + ailego::IOBackendType *used_backend) { if (async == true) { LOG_WARN("Async currently not supported"); } @@ -650,7 +662,7 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return IndexError_Runtime; } - int ret = execute_io(ctx, this->file_desc, read_reqs); + int ret = execute_io(ctx, this->file_desc, read_reqs, 0, used_backend); return ret; } diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 1023fc40c..258b4c1aa 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -17,11 +17,15 @@ #include +#if defined(__APPLE__) +#include +#endif + #if (defined(__linux) || defined(__linux__)) #include // dlopen-based libaio wrapper #endif -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX #include #include #include @@ -33,6 +37,7 @@ #include #include #include +#include #include #include "diskann_util.h" @@ -44,7 +49,7 @@ namespace core { // On other platforms, IOContext is a uint32_t placeholder. #if (defined(__linux) || defined(__linux__)) typedef io_context_t IOContext; -#elif defined(__APPLE__) || defined(__MACH__) +#elif defined(__APPLE__) && TARGET_OS_OSX typedef int IOContext; #else typedef uint32_t IOContext; @@ -88,8 +93,11 @@ class AlignedFileReader { virtual void open(const std::string &fname) = 0; virtual void close() = 0; + // used_backend, when provided, receives the backend that completed this + // call, including a synchronous fallback from an async backend. virtual int read(std::vector &read_reqs, IOContext &ctx, - bool async = false) = 0; + bool async = false, + ailego::IOBackendType *used_backend = nullptr) = 0; }; // Reader implementation used on all supported platforms. @@ -117,7 +125,7 @@ class LinuxAlignedFileReader : public AlignedFileReader { void close(); int read(std::vector &read_reqs, IOContext &ctx, - bool async = false); + bool async = false, ailego::IOBackendType *used_backend = nullptr); }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index 5eea9c7c3..abd8039d6 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -92,7 +92,7 @@ class DiskAnnIndexer { PQTable::Pointer pq_table_; IOContext init_ctx_{ -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX -1 #else 0 diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index e719c8ae6..f1553b6ff 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -15,7 +15,6 @@ #include #include #include -#include #include #include "algorithm/diskann/diskann_params.h" #include "holder_builder.h" @@ -272,17 +271,4 @@ int DiskAnnIndex::Merge(const std::vector &indexes, return 0; } -ailego::IOBackendType DiskAnnIndex::io_backend_type() const { - ailego::IOBackendType type = ailego::current_io_backend_type(); -#if defined(__linux__) || defined(__linux) - if (type == ailego::IOBackendType::kPread) { - LOG_WARN( - "Only synchronous pread() is available. Install libaio " - "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " - "for async I/O support — performance will be degraded without it."); - } -#endif - return type; -} - } // namespace zvec::core_interface diff --git a/src/db/collection.cc b/src/db/collection.cc index a6fc88834..33750f3b1 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -134,8 +134,6 @@ class CollectionImpl : public Collection { Result DebugGetHnswStorageMode( const std::string &column_name) const override; - Result DebugGetIoBackendType() const override; - private: void prepare_schema(); @@ -1835,19 +1833,6 @@ Result CollectionImpl::DebugGetHnswStorageMode( Status::NotFound("No HNSW index found for column '", column_name, "'")); } -Result CollectionImpl::DebugGetIoBackendType() const { - const auto type = ailego::current_io_backend_type(); - switch (type) { - case ailego::IOBackendType::kLibAio: - return "libaio"; - case ailego::IOBackendType::kPosixAio: - return "posix_aio"; - case ailego::IOBackendType::kPread: - return "sync_pread"; - } - return "unknown"; -} - Status CollectionImpl::recovery() { if (!FileHelper::DirectoryExists(path_.c_str())) { return Status::InvalidArgument("collection path{", path_, "} not exist."); diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 3fe89ff96..118d06d44 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include @@ -377,9 +376,6 @@ class DiskAnnIndex : public Index { public: DiskAnnIndex() = default; - // Returns the I/O backend type selected for DiskAnn disk reads. - ailego::IOBackendType io_backend_type() const; - protected: virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 11c04edb1..581a9449b 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -117,11 +117,6 @@ class Collection { //! introspection and testing; not part of the stable public API. virtual Result DebugGetHnswStorageMode( const std::string &column_name) const = 0; - - //! Debug-only: retrieve the I/O backend type used by DiskAnn. Returns - //! "libaio", "posix_aio", or "sync_pread". Intended for introspection and - //! testing; not part of the stable public API. - virtual Result DebugGetIoBackendType() const = 0; }; } // namespace zvec diff --git a/tests/ailego/io/io_backend_test.cc b/tests/ailego/io/io_backend_test.cc index 66eccffb5..4ae331361 100644 --- a/tests/ailego/io/io_backend_test.cc +++ b/tests/ailego/io/io_backend_test.cc @@ -12,11 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include -#include #include +#if defined(__APPLE__) +#include +#endif + namespace zvec { namespace ailego { @@ -32,11 +36,15 @@ TEST(IOBackendTest, BackendNames) { EXPECT_STREQ(IOBackendTypeName(IOBackendType::kPosixAio), "posix_aio"); } -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX TEST(IOBackendTest, MacOSUsesPosixAio) { EXPECT_EQ(current_io_backend_type(), IOBackendType::kPosixAio); EXPECT_NE(current_io_backend_description().find("kqueue"), std::string::npos); } +#elif defined(__APPLE__) +TEST(IOBackendTest, NonMacOSAppleUsesPread) { + EXPECT_EQ(current_io_backend_type(), IOBackendType::kPread); +} #endif } // namespace ailego diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index a790e2383..f63a47482 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -20,6 +20,10 @@ #include #include +#if defined(__APPLE__) +#include +#endif + // Platform-specific headers #ifdef _WIN32 #include @@ -139,7 +143,7 @@ void test_io_backend_functions(void) { TEST_ASSERT(name != NULL); TEST_ASSERT(description != NULL); -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_POSIX_AIO); TEST_ASSERT(strcmp(name, "posix_aio") == 0); #elif defined(__linux__) || defined(__linux) diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc index 627fed07a..95969c13e 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -67,8 +67,8 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { std::vector source(kPageSize * kPageCount); for (size_t page = 0; page < kPageCount; ++page) { - std::memset(source.data() + page * kPageSize, - static_cast(page + 1), kPageSize); + std::memset(source.data() + page * kPageSize, static_cast(page + 1), + kPageSize); } size_t written = 0; @@ -92,21 +92,25 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { // so the macOS backend has to submit multiple batches. size_t source_page = (i * 7) % kPageCount; requests.emplace_back(source_page * kPageSize, kPageSize, - static_cast(output.get()) + - i * kPageSize); + static_cast(output.get()) + i * kPageSize); } LinuxAlignedFileReader reader; reader.open(file.path()); IOContext ctx{ -#if defined(__APPLE__) || defined(__MACH__) +#if defined(__APPLE__) && TARGET_OS_OSX -1 #else 0 #endif }; ASSERT_EQ(setup_io_ctx(ctx), 0); - ASSERT_EQ(reader.read(requests, ctx), 0); + zvec::ailego::IOBackendType used_backend = + zvec::ailego::IOBackendType::kPread; + ASSERT_EQ(reader.read(requests, ctx, false, &used_backend), 0); +#if defined(__APPLE__) && TARGET_OS_OSX + EXPECT_EQ(used_backend, zvec::ailego::IOBackendType::kPosixAio); +#endif for (size_t i = 0; i < kPageCount; ++i) { size_t source_page = (i * 7) % kPageCount; @@ -120,3 +124,28 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { EXPECT_EQ(destroy_io_ctx(ctx), 0); reader.close(); } + +#if defined(__APPLE__) && TARGET_OS_OSX +TEST(DiskAnnFileReaderTest, InvalidKqueueFallsBackToPread) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + constexpr char kSource[] = "DiskAnn pread fallback"; + ASSERT_EQ(::pwrite(file.fd(), kSource, sizeof(kSource), 0), + static_cast(sizeof(kSource))); + file.close(); + + char output[sizeof(kSource)] = {}; + std::vector requests{{0, sizeof(kSource), output}}; + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = -1; + zvec::ailego::IOBackendType used_backend = + zvec::ailego::IOBackendType::kPosixAio; + + ASSERT_EQ(reader.read(requests, ctx, false, &used_backend), 0); + EXPECT_EQ(used_backend, zvec::ailego::IOBackendType::kPread); + EXPECT_EQ(std::memcmp(output, kSource, sizeof(kSource)), 0); + reader.close(); +} +#endif diff --git a/tests/db/crash_recovery/write_recovery_test.cc b/tests/db/crash_recovery/write_recovery_test.cc index 9f878b1eb..702e9e498 100644 --- a/tests/db/crash_recovery/write_recovery_test.cc +++ b/tests/db/crash_recovery/write_recovery_test.cc @@ -188,7 +188,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { collection.reset(); } - RunGeneratorAndCrash("0", "10000", "insert", "0", 5); + RunGeneratorAndCrash("0", "10000", "insert", "0", 3); auto result = Collection::Open(dir_path_, options_); ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " @@ -196,7 +196,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { auto collection = result.value(); uint64_t doc_count{collection->Stats().value().doc_count}; ASSERT_GT(doc_count, 800) - << "Document count is too low after 5s of insertion and recovery"; + << "Document count is too low after 3s of insertion and recovery"; for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) { const auto expected_doc = CreateTestDoc(doc_id, 0); From 099f49bb4308d75c9309359194289561891a7b9a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 19:38:01 +0800 Subject: [PATCH 38/46] upd --- CMakeLists.txt | 7 ++++--- examples/c/diskann_example.c | 8 +++----- python/tests/detail/fixture_helper.py | 13 +++++++------ python/tests/test_collection_diskann.py | 22 ++++++++++++---------- src/core/algorithm/CMakeLists.txt | 2 +- src/db/index/common/schema.cc | 10 +++++----- src/include/zvec/c_api.h | 3 ++- 7 files changed, 34 insertions(+), 31 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 1565ab79e..c1cf038f3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -123,8 +123,9 @@ endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") # DiskAnn support: -# - Linux (x86_64, i686, i386, aarch64, arm64) with libaio (loaded via dlopen) -# - macOS (x86_64, ARM64/Apple Silicon) with kqueue +# - Linux (x86_64, i686, i386, aarch64, arm64), using libaio when available +# and synchronous pread otherwise +# - macOS (x86_64, ARM64/Apple Silicon), using POSIX AIO with kqueue if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS) OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND NOT IOS)) set(DISKANN_SUPPORTED ON) @@ -132,7 +133,7 @@ if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_6 else() set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled - supported on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)") + message(STATUS "DiskAnn support disabled - supported on Linux x86/x86_64/ARM64 and macOS") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") diff --git a/examples/c/diskann_example.c b/examples/c/diskann_example.c index 5ef80562f..a2fe47034 100644 --- a/examples/c/diskann_example.c +++ b/examples/c/diskann_example.c @@ -21,8 +21,9 @@ * a Vamana graph structure combined with product quantization (PQ) to * achieve high recall with efficient disk I/O. * - * NOTE: DiskANN requires Linux x86_64 with libaio. On other platforms the - * example will compile but the runtime plugin will fail to load. + * NOTE: DiskANN supports Linux x86/x86_64/ARM64 (libaio when available, + * synchronous pread fallback otherwise) and macOS (POSIX AIO with kqueue). + * Schema validation rejects DiskANN on other platforms. * * Workflow demonstrated: * 1. Create collection schema with DiskANN-indexed vector field @@ -265,9 +266,6 @@ int main(void) { &results, &result_count); if (error != ZVEC_OK) { handle_error(error, "executing DiskANN query"); - printf( - " (This is expected on non-Linux platforms — DiskANN requires " - "libaio)\n"); } else { printf(" Query returned %zu results:\n", result_count); for (size_t r = 0; r < result_count && r < 5; r++) { diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index f5e6c08e4..2aed7bf66 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,11 +2,10 @@ import logging import platform -DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( - "x86_64", - "AMD64", - "i686", - "i386", +_MACHINE = platform.machine().lower() +DISKANN_SUPPORTED = platform.system() == "Darwin" or ( + platform.system() == "Linux" + and _MACHINE in ("x86_64", "amd64", "i686", "i386", "aarch64", "arm64") ) from typing import Any, Generator @@ -27,7 +26,9 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_DONE = True if not DISKANN_SUPPORTED: - _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" + _DISKANN_PRELOAD_REASON = ( + "DiskAnn is supported on Linux x86/x86_64/ARM64 and macOS" + ) return _DISKANN_PRELOAD_REASON _DISKANN_PRELOAD_REASON = None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 8daf8c770..98b14060b 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -15,16 +15,13 @@ Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index. -Two platform-level prerequisites are enforced at module import time: +Platform behavior: -1. DiskAnn is currently built only for Linux x86_64 — other platforms are - skipped wholesale. -2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() / - DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to - synchronous pread() — the tests still run but with degraded performance. +1. Linux x86/x86_64/ARM64 probes libaio at runtime and falls back to + synchronous pread() when libaio is unavailable. +2. macOS uses the system POSIX AIO implementation with kqueue notifications. -If either prerequisite fails the whole module is skipped so the rest of the -test-suite is not affected. +The module is skipped on platforms where DiskAnn is not built. """ from __future__ import annotations @@ -38,9 +35,14 @@ # --------------------------------------------------------------------------- # # Platform gating (must happen BEFORE we touch zvec). # --------------------------------------------------------------------------- # +_machine = platform.machine().lower() +_diskann_supported = sys.platform == "darwin" or ( + sys.platform == "linux" + and _machine in ("x86_64", "amd64", "i686", "i386", "aarch64", "arm64") +) pytestmark = pytest.mark.skipif( - not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), - reason="DiskAnn plugin is only supported on Linux x86_64", + not _diskann_supported, + reason="DiskAnn is supported on Linux x86/x86_64/ARM64 and macOS", ) import zvec # noqa: E402 diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index 23aea7773..5d5654a9a 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -17,7 +17,7 @@ else() # Empty stub library for unsupported platforms file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc "// Stub implementation for unsupported platforms\n" - "// DiskAnn supports Linux (x86_64/ARM64 with libaio) and macOS (kqueue)\n" + "// DiskAnn supports Linux x86/x86_64/ARM64 and macOS\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..b75e72c1c 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,19 +198,19 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable + // DiskAnn supports Linux x86/x86_64/ARM64 and macOS. The CMake variable // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the // single source of truth for platform eligibility — it is also used by // index_factory.cc to conditionally compile the DiskAnn index // registration. Using the same macro here ensures that schema // validation and index registration agree on supported platforms. // - // libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() - // and DiskAnnStreamer::init(); if libaio is missing, DiskAnn falls - // back to synchronous pread() with degraded performance. + // Linux probes libaio at runtime and falls back to synchronous pread() + // when it is unavailable. macOS uses system POSIX AIO with kqueue + // completion notifications. #if !DISKANN_SUPPORTED return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 only)"); + "DiskAnn is supported only on Linux x86/x86_64/ARM64 and macOS"); #endif } diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 00afc9876..d0313ac22 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -815,7 +815,8 @@ zvec_get_io_backend_type_name(zvec_io_backend_type_t type); /** * @brief Get a human-readable description of the current I/O backend. * - * When only pread is available, includes installation guidance for libaio. + * On Linux, the pread description includes installation guidance for libaio. + * Other platforms return a generic synchronous-backend description. * * @return Thread-local string valid until the next call on this thread. */ From 0a1b76fc472b438faa3bec18017d7a4d7b530b0e Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 20 Jul 2026 21:43:07 +0800 Subject: [PATCH 39/46] fix --- .../algorithm/diskann/diskann_streamer.cc | 1 + .../diskann/diskann_searcher_test.cc | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index f62395bb7..5aa6213d5 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -354,6 +354,7 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } ctx->set_list_size(list_size_); + ctx->set_magic(magic_); return Context::Pointer(ctx); } diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index ff432886f..c2b0dbf35 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -220,6 +220,27 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { EXPECT_GT(recall, 0.90f); EXPECT_GT(topk1Recall, 0.80f); EXPECT_GT(cost, 2.0f); + + // A context created by the streamer must carry the streamer's magic so it + // can be reused instead of being recreated on every search. + IndexStreamer::Pointer streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + ASSERT_EQ(0, streamer->init(*_index_meta_ptr, search_params)); + + auto streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_EQ(0, streamer_storage->open(path, false)); + ASSERT_EQ(0, streamer->open(streamer_storage)); + + auto streamer_ctx = streamer->create_context(); + ASSERT_NE(streamer_ctx, nullptr); + streamer_ctx->set_topk(topk); + auto *original_ctx = streamer_ctx.get(); + + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); } TEST_F(DiskAnnSearcherTest, TestNodeCache) { From 12ea468e881676a2b1e2baaae75f5810caa729dd Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 10:45:41 +0800 Subject: [PATCH 40/46] clang format --- tests/core/algorithm/diskann/diskann_file_reader_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc index 95969c13e..58da74a3f 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -101,7 +101,7 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { #if defined(__APPLE__) && TARGET_OS_OSX -1 #else - 0 + nullptr #endif }; ASSERT_EQ(setup_io_ctx(ctx), 0); From f4a3947270952214480e2933b4cbfc2c766aeb5a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 11:40:28 +0800 Subject: [PATCH 41/46] rm kqueue --- examples/c/diskann_example.c | 3 +- python/tests/test_collection_diskann.py | 2 +- python/zvec/typing/__init__.pyi | 2 +- src/ailego/io/io_backend_def.h | 8 +- src/binding/python/typing/python_type.cc | 2 +- src/core/CMakeLists.txt | 2 +- .../algorithm/diskann/diskann_file_reader.cc | 174 ++++-------------- .../algorithm/diskann/diskann_file_reader.h | 16 +- src/db/index/common/schema.cc | 4 +- src/include/zvec/ailego/io/io_backend.h | 2 +- src/include/zvec/c_api.h | 2 +- tests/ailego/io/io_backend_test.cc | 3 +- .../diskann/diskann_file_reader_test.cc | 13 +- 13 files changed, 56 insertions(+), 177 deletions(-) diff --git a/examples/c/diskann_example.c b/examples/c/diskann_example.c index a2fe47034..facabbeb1 100644 --- a/examples/c/diskann_example.c +++ b/examples/c/diskann_example.c @@ -22,7 +22,8 @@ * achieve high recall with efficient disk I/O. * * NOTE: DiskANN supports Linux x86/x86_64/ARM64 (libaio when available, - * synchronous pread fallback otherwise) and macOS (POSIX AIO with kqueue). + * synchronous pread fallback otherwise) and macOS (POSIX AIO with + * aio_suspend()). * Schema validation rejects DiskANN on other platforms. * * Workflow demonstrated: diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 98b14060b..b948bb27f 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -19,7 +19,7 @@ 1. Linux x86/x86_64/ARM64 probes libaio at runtime and falls back to synchronous pread() when libaio is unavailable. -2. macOS uses the system POSIX AIO implementation with kqueue notifications. +2. macOS uses the system POSIX AIO implementation with aio_suspend() waits. The module is skipped on platforms where DiskAnn is not built. """ diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index b78776d7f..0e969a7a7 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -130,7 +130,7 @@ class IOBackendType: - PREAD: Synchronous pread() — no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). - - POSIX_AIO: macOS POSIX AIO with kqueue completion notifications. + - POSIX_AIO: macOS POSIX AIO with aio_suspend() completion waits. Examples: >>> from zvec.typing import IOBackendType diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index 19118800a..4155066e5 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -59,7 +59,7 @@ inline const char *IOBackendDescription(IOBackendType type) { case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; case IOBackendType::kPosixAio: - return "macOS POSIX AIO backend with kqueue completion notifications."; + return "macOS POSIX AIO backend with aio_suspend() completion waits."; case IOBackendType::kPread: #if defined(__linux) || defined(__linux__) return "No async I/O backend available. Install libaio (e.g. " @@ -105,10 +105,8 @@ class IOBackend { return type_; } #if defined(__APPLE__) && TARGET_OS_OSX - // POSIX AIO and EVFILT_AIO are provided by Darwin; no user-installed - // runtime dependency needs to be probed. A failure to create a particular - // kqueue context is handled by the DiskAnn reader by falling back to - // synchronous pread(). + // POSIX AIO is provided by Darwin; no user-installed runtime dependency + // needs to be probed. type_ = IOBackendType::kPosixAio; return type_; #endif diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index 87ae79489..58a260b72 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -146,7 +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(). -- POSIX_AIO: macOS POSIX AIO with kqueue completion notifications. +- POSIX_AIO: macOS POSIX AIO with aio_suspend() completion waits. Examples: >>> from zvec.typing import IOBackendType diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index 84dd6899e..f29596ce3 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -75,7 +75,7 @@ endif() set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) # The DiskAnn runtime loader (LibAioLoader) uses dlopen/dlsym, so link libdl -# on Linux. On macOS, kqueue is part of the system and does not require dlopen. +# on Linux. On macOS, POSIX AIO is provided by the system. if(CMAKE_SYSTEM_NAME STREQUAL "Linux") list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) endif() diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 16a880506..69dcf08e3 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -13,17 +13,14 @@ // limitations under the License. #include "diskann_file_reader.h" -#include #include -#include #include -#include +#include +#include #include #include #if defined(__APPLE__) && TARGET_OS_OSX #include -#include -#include #endif #define MAX_EVENTS 1024 @@ -61,20 +58,11 @@ int setup_io_ctx(IOContext &ctx) { return ret; #elif defined(__APPLE__) && TARGET_OS_OSX std::call_once(g_io_backend_log_once, [] { - LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with aio_suspend"); }); - // Each macOS I/O context owns a kqueue which receives EVFILT_AIO - // completion notifications for POSIX AIO requests. - ctx = -1; - int kq = ::kqueue(); - if (kq == -1) { - LOG_WARN( - "kqueue() failed in setup_io_ctx; errno=%d, %s; falling back to " - "synchronous pread", - errno, ::strerror(errno)); - return 0; - } - ctx = kq; + // POSIX AIO does not require a persistent per-context kernel object on + // macOS. Keep a placeholder value for the cross-platform IOContext API. + ctx = 0; return 0; #else return 0; @@ -90,10 +78,7 @@ int destroy_io_ctx(IOContext &ctx) { int ret = LibAioLoader::Instance().io_destroy(ctx); return ret; #elif defined(__APPLE__) && TARGET_OS_OSX - if (ctx >= 0) { - ::close(ctx); - ctx = -1; - } + ctx = 0; return 0; #else return 0; @@ -143,28 +128,23 @@ static bool validate_aio_result(int aio_err, int64_t result, return false; } -// Once the kqueue has been closed, its AIO knotes are detached and the normal -// POSIX aio_error()/aio_return() cleanup path is safe again. This keeps live -// requests and destination buffers from escaping the call if kevent64 fails. -static bool drain_aio_after_kqueue_failure( +// Wait for and reap every submitted request. Requests and their destination +// buffers must remain alive until aio_error() no longer reports EINPROGRESS, +// and each completed request must be reaped exactly once with aio_return(). +static bool drain_aio_requests( std::vector &cbs, const std::vector &read_reqs, - size_t req_begin, const std::vector &submitted, - std::vector &completed, size_t submitted_count) { + size_t req_begin, size_t submitted_count) { size_t completed_count = 0; bool all_ok = true; - for (size_t i = 0; i < cbs.size(); ++i) { - if (submitted[i] != 0 && completed[i] != 0) { - ++completed_count; - } - } + std::vector completed(submitted_count, 0); while (completed_count < submitted_count) { bool made_progress = false; std::vector pending; pending.reserve(submitted_count - completed_count); - for (size_t i = 0; i < cbs.size(); ++i) { - if (submitted[i] == 0 || completed[i] != 0) { + for (size_t i = 0; i < submitted_count; ++i) { + if (completed[i] != 0) { continue; } @@ -203,29 +183,15 @@ static bool drain_aio_after_kqueue_failure( return all_ok; } -// Submit a batch through POSIX aio_read(). SIGEV_KEVENT makes XNU attach an -// EVFILT_AIO one-shot event to the supplied kqueue for each request. kqueue is -// therefore used for completion notification, not regular-file readiness. -// -// Darwin's EVFILT_AIO filter is also the completion reaper: fetching the event -// removes the request from the process AIO done queue. The error and return -// values are carried in kevent64_s::ext[0] and ext[1], respectively. Calling -// aio_error()/aio_return() after fetching the event would incorrectly report -// EINVAL because the kernel has already released that request. -static int execute_io_aio_kqueue(int &kq, int fd, - std::vector &read_reqs, - uint64_t n_retries, - ailego::IOBackendType *used_backend) { - if (kq < 0) { - return execute_io_pread(fd, read_reqs, used_backend); - } - +// Submit a batch through POSIX aio_read() and wait for completion with the +// portable aio_suspend() API supported by all target macOS SDK versions. +static int execute_io_aio_suspend(int fd, std::vector &read_reqs, + uint64_t n_retries, + ailego::IOBackendType *used_backend) { for (size_t req_begin = 0; req_begin < read_reqs.size(); req_begin += kMacAioBatchSize) { size_t n_ops = std::min(kMacAioBatchSize, read_reqs.size() - req_begin); std::vector cbs(n_ops); - std::vector submitted(n_ops, 0); - std::vector completed(n_ops, 0); size_t submitted_count = 0; bool submission_ok = true; @@ -237,9 +203,7 @@ static int execute_io_aio_kqueue(int &kq, int fd, cbs[i].aio_nbytes = read_reqs[req_begin + i].len; cbs[i].aio_reqprio = 0; cbs[i].aio_lio_opcode = LIO_READ; - cbs[i].aio_sigevent.sigev_notify = SIGEV_KEVENT; - cbs[i].aio_sigevent.sigev_signo = kq; - cbs[i].aio_sigevent.sigev_value.sival_ptr = &cbs[i]; + cbs[i].aio_sigevent.sigev_notify = SIGEV_NONE; uint64_t tries = 0; while (::aio_read(&cbs[i]) == -1) { @@ -261,68 +225,13 @@ static int execute_io_aio_kqueue(int &kq, int fd, if (!submission_ok) { break; } - submitted[i] = 1; ++submitted_count; } - size_t completed_count = 0; - bool all_ok = true; - bool kqueue_ok = true; - while (completed_count < submitted_count) { - struct kevent64_s events[AIO_LISTIO_MAX]; - int n_events = ::kevent64(kq, nullptr, 0, events, static_cast(n_ops), - 0, nullptr); - if (n_events == -1) { - if (errno == EINTR) { - continue; - } - LOG_ERROR("kevent failed while waiting for AIO; errno=%d, %s", errno, - ::strerror(errno)); - kqueue_ok = false; - break; - } - - for (int event_index = 0; event_index < n_events; ++event_index) { - size_t request_index = n_ops; - for (size_t i = 0; i < n_ops; ++i) { - if (events[event_index].ident == - reinterpret_cast(&cbs[i])) { - request_index = i; - break; - } - } + bool all_ok = + drain_aio_requests(cbs, read_reqs, req_begin, submitted_count); - if (request_index == n_ops || submitted[request_index] == 0 || - completed[request_index] != 0 || - events[event_index].filter != EVFILT_AIO) { - LOG_WARN("ignoring an unknown or duplicate macOS AIO event"); - continue; - } - - int aio_err = static_cast(events[event_index].ext[0]); - int64_t result = static_cast(events[event_index].ext[1]); - const AlignedRead &req = read_reqs[req_begin + request_index]; - all_ok = validate_aio_result(aio_err, result, req) && all_ok; - completed[request_index] = 1; - ++completed_count; - } - if (!kqueue_ok) { - break; - } - } - - if (!kqueue_ok) { - // Detach the EVFILT_AIO knotes before using aio_return() to reap the - // requests through the POSIX API, and disable AIO for this context. - ::close(kq); - kq = -1; - all_ok = - drain_aio_after_kqueue_failure(cbs, read_reqs, req_begin, submitted, - completed, submitted_count) && - all_ok; - } - - if (!submission_ok || !kqueue_ok || !all_ok) { + if (!submission_ok || !all_ok) { return execute_io_pread(fd, read_reqs, used_backend); } } @@ -420,9 +329,8 @@ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, } return 0; #elif defined(__APPLE__) && TARGET_OS_OSX - // On macOS, submit POSIX AIO and use the IOContext kqueue for EVFILT_AIO - // completion notifications. - return execute_io_aio_kqueue(ctx, fd, read_reqs, n_retries, used_backend); + (void)ctx; + return execute_io_aio_suspend(fd, read_reqs, n_retries, used_backend); #else (void)ctx; (void)n_retries; @@ -503,19 +411,10 @@ void LinuxAlignedFileReader::register_thread() { return; } - IOContext ctx = -1; - int kq = ::kqueue(); - if (kq == -1) { - LOG_ERROR("kqueue() failed in register_thread; errno=%d, %s", errno, - ::strerror(errno)); - } else { - std::call_once(g_io_backend_log_once, [] { - LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with kqueue completion"); - }); - LOG_INFO("allocating POSIX AIO kqueue ctx: %d", kq); - ctx = kq; - ctx_map[thread_id] = ctx; - } + std::call_once(g_io_backend_log_once, [] { + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with aio_suspend"); + }); + ctx_map[thread_id] = 0; lk.unlock(); #endif } @@ -544,7 +443,6 @@ void LinuxAlignedFileReader::deregister_thread() { LOG_INFO("returned ctx from thread"); #elif defined(__APPLE__) && TARGET_OS_OSX auto thread_id = std::this_thread::get_id(); - IOContext ctx; { std::lock_guard lk(ctx_mut); @@ -553,14 +451,10 @@ void LinuxAlignedFileReader::deregister_thread() { LOG_ERROR("deregister_thread: thread not registered"); return; } - ctx = it->second; ctx_map.erase(it); } - if (ctx >= 0) { - ::close(ctx); - } - LOG_INFO("returned kqueue ctx from thread"); + LOG_INFO("deregistered POSIX AIO thread"); #endif } @@ -578,12 +472,6 @@ void LinuxAlignedFileReader::deregister_all_threads() { ctx_map.clear(); #elif defined(__APPLE__) && TARGET_OS_OSX std::unique_lock lk(ctx_mut); - for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { - IOContext ctx = x->second; - if (ctx >= 0) { - ::close(ctx); - } - } ctx_map.clear(); #endif } diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 258b4c1aa..16d089ef5 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -15,8 +15,6 @@ #define MAX_IO_DEPTH 128 -#include - #if defined(__APPLE__) #include #endif @@ -25,14 +23,7 @@ #include // dlopen-based libaio wrapper #endif -#if defined(__APPLE__) && TARGET_OS_OSX -#include -#include -#include -#endif - -#include -#include +#include #include #include #include @@ -45,7 +36,8 @@ namespace zvec { namespace core { // On Linux, IOContext is the libaio io_context_t. -// On macOS, IOContext is an int holding a kqueue file descriptor. +// On macOS, IOContext is an unused int placeholder because POSIX AIO does not +// require a persistent per-context resource. // On other platforms, IOContext is a uint32_t placeholder. #if (defined(__linux) || defined(__linux__)) typedef io_context_t IOContext; @@ -103,7 +95,7 @@ class AlignedFileReader { // Reader implementation used on all supported platforms. // On Linux (x86_64 and ARM64) it uses libaio for asynchronous batch I/O. // On macOS (including ARM/Apple Silicon) it submits POSIX AIO requests and -// uses kqueue EVFILT_AIO events for completion notification. +// waits for completion with aio_suspend(). class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index b75e72c1c..875520f63 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -206,8 +206,8 @@ Status FieldSchema::validate() const { // validation and index registration agree on supported platforms. // // Linux probes libaio at runtime and falls back to synchronous pread() - // when it is unavailable. macOS uses system POSIX AIO with kqueue - // completion notifications. + // when it is unavailable. macOS uses system POSIX AIO with + // aio_suspend() completion waits. #if !DISKANN_SUPPORTED return Status::NotSupported( "DiskAnn is supported only on Linux x86/x86_64/ARM64 and macOS"); diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 1ef168095..c177c424a 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -32,7 +32,7 @@ namespace ailego { enum class IOBackendType { kPread = 0, // Synchronous pread() — no async I/O kLibAio = 1, // Linux libaio loaded at runtime via dlopen() - kPosixAio = 2 // macOS POSIX AIO with kqueue completion + kPosixAio = 2 // macOS POSIX AIO with aio_suspend() completion waits }; // Returns the currently active I/O backend type. diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index d0313ac22..9b8d4a410 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -790,7 +790,7 @@ typedef uint32_t zvec_io_backend_type_t; #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ #define ZVEC_IO_BACKEND_TYPE_POSIX_AIO \ - 2 /**< macOS POSIX AIO with kqueue completion */ + 2 /**< macOS POSIX AIO with aio_suspend() completion waits */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. diff --git a/tests/ailego/io/io_backend_test.cc b/tests/ailego/io/io_backend_test.cc index 4ae331361..800696bde 100644 --- a/tests/ailego/io/io_backend_test.cc +++ b/tests/ailego/io/io_backend_test.cc @@ -39,7 +39,8 @@ TEST(IOBackendTest, BackendNames) { #if defined(__APPLE__) && TARGET_OS_OSX TEST(IOBackendTest, MacOSUsesPosixAio) { EXPECT_EQ(current_io_backend_type(), IOBackendType::kPosixAio); - EXPECT_NE(current_io_backend_description().find("kqueue"), std::string::npos); + EXPECT_NE(current_io_backend_description().find("aio_suspend"), + std::string::npos); } #elif defined(__APPLE__) TEST(IOBackendTest, NonMacOSAppleUsesPread) { diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc index 58da74a3f..2eb2d76c7 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -99,7 +99,7 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { reader.open(file.path()); IOContext ctx{ #if defined(__APPLE__) && TARGET_OS_OSX - -1 + 0 #else nullptr #endif @@ -126,7 +126,7 @@ TEST(DiskAnnFileReaderTest, BatchAlignedReads) { } #if defined(__APPLE__) && TARGET_OS_OSX -TEST(DiskAnnFileReaderTest, InvalidKqueueFallsBackToPread) { +TEST(DiskAnnFileReaderTest, ShortAioReadFallsBackToPread) { TemporaryFile file; ASSERT_GE(file.fd(), 0); @@ -135,17 +135,16 @@ TEST(DiskAnnFileReaderTest, InvalidKqueueFallsBackToPread) { static_cast(sizeof(kSource))); file.close(); - char output[sizeof(kSource)] = {}; - std::vector requests{{0, sizeof(kSource), output}}; + char output[sizeof(kSource) * 2] = {}; + std::vector requests{{0, sizeof(output), output}}; LinuxAlignedFileReader reader; reader.open(file.path()); - IOContext ctx = -1; + IOContext ctx = 0; zvec::ailego::IOBackendType used_backend = zvec::ailego::IOBackendType::kPosixAio; - ASSERT_EQ(reader.read(requests, ctx, false, &used_backend), 0); + EXPECT_NE(reader.read(requests, ctx, false, &used_backend), 0); EXPECT_EQ(used_backend, zvec::ailego::IOBackendType::kPread); - EXPECT_EQ(std::memcmp(output, kSource, sizeof(kSource)), 0); reader.close(); } #endif From 216e2fa478e41248ecaf531807a941283f4b87a6 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 11:41:51 +0800 Subject: [PATCH 42/46] clang format --- CMakeLists.txt | 2 +- src/core/algorithm/diskann/diskann_file_reader.cc | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index c1cf038f3..a58a334b6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -125,7 +125,7 @@ message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") # DiskAnn support: # - Linux (x86_64, i686, i386, aarch64, arm64), using libaio when available # and synchronous pread otherwise -# - macOS (x86_64, ARM64/Apple Silicon), using POSIX AIO with kqueue +# - macOS (x86_64, ARM64/Apple Silicon), using POSIX AIO with aio_suspend if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS) OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND NOT IOS)) set(DISKANN_SUPPORTED ON) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 69dcf08e3..813f645fb 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -13,10 +13,10 @@ // limitations under the License. #include "diskann_file_reader.h" -#include -#include #include #include +#include +#include #include #include #if defined(__APPLE__) && TARGET_OS_OSX @@ -131,9 +131,9 @@ static bool validate_aio_result(int aio_err, int64_t result, // Wait for and reap every submitted request. Requests and their destination // buffers must remain alive until aio_error() no longer reports EINPROGRESS, // and each completed request must be reaped exactly once with aio_return(). -static bool drain_aio_requests( - std::vector &cbs, const std::vector &read_reqs, - size_t req_begin, size_t submitted_count) { +static bool drain_aio_requests(std::vector &cbs, + const std::vector &read_reqs, + size_t req_begin, size_t submitted_count) { size_t completed_count = 0; bool all_ok = true; std::vector completed(submitted_count, 0); From a3213b867c8f79a0f8f6f453b0e91a318f17dee3 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 14:52:17 +0800 Subject: [PATCH 43/46] upd --- src/core/algorithm/diskann/diskann_builder.cc | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index 26afea93f..4ab837b16 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -287,24 +287,25 @@ int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } LOG_INFO("Built cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + // A worker may fail before processing its entire shard. Always wait for the + // remaining workers before returning so they cannot access this builder, + // its mutex, or the stack-owned finished counter after their lifetimes end. + task_group->wait_finish(); + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to build index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } @@ -324,24 +325,24 @@ int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to prune index while waiting finish"); - return errcode_; - } LOG_INFO("Prune cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + // See build_internal(): error paths must not let worker closures outlive the + // builder or the stack-owned completion counter. + task_group->wait_finish(); + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to prune index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } From 2f3069cf5c9e2776a5b6ec5e9cc192205e98e29a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 15:55:22 +0800 Subject: [PATCH 44/46] fix --- src/core/algorithm/diskann/diskann_builder_entity.cc | 2 +- src/core/algorithm/diskann/diskann_pq_trainer.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc index 4381293a5..992372a9a 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.cc +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -53,7 +53,7 @@ int DiskAnnBuilderEntity::add_vector(diskann_key_t key, const void *vec) { keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); uint32_t neighbor_cnt = 0; - std::vector neighbor{max_build_degree_, 0}; + std::vector neighbor(max_build_degree_, 0); neighbors_buffer_.append(reinterpret_cast(&neighbor_cnt), sizeof(uint32_t)); diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc index 73ca01656..c84744cb8 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.cc +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -153,7 +153,7 @@ int DiskAnnPqTrainer::convert_pivot_data( cluster * dim + chunk_offsets[chunk]; const T *feature_ptr = reinterpret_cast(centroids[idx].feature()); - for (size_t d = 0; d <= chunk_dims[chunk]; ++d) { + for (size_t d = 0; d < chunk_dims[chunk]; ++d) { pivot_data_ptr[d] = feature_ptr[d]; } } From be4a7e0eeb600c74d75fc93521d76768a14fdf15 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 16:13:22 +0800 Subject: [PATCH 45/46] fix --- src/core/algorithm/diskann/diskann_indexer.cc | 9 ++++++++- src/core/algorithm/diskann/diskann_searcher.cc | 8 ++++---- src/core/algorithm/diskann/diskann_streamer.cc | 6 +++++- tests/core/algorithm/diskann/diskann_searcher_test.cc | 6 ++++++ 4 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..36ef567ac 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -1086,8 +1086,15 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { io_timer.reset(); - reader_->read(frontier_read_reqs, io_ctx); // synchronous IO linux + int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR( + "cached_beam_search_by_group: reader_->read failed, ret=%d", + read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } } for (auto &cached_neighbor : cached_neighbors) { diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a8bfed2e9..13d5d1b3a 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -177,10 +177,10 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; + int ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0 || ctx->error())) { + LOG_ERROR("DiskAnn knn search failed, ret=%d", ret); + return ret != 0 ? ret : IndexError_Runtime; } ctx->topk_to_result(i); diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 5aa6213d5..ca312e4b5 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -176,7 +176,11 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); + int ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0 || ctx->error())) { + LOG_ERROR("DiskAnn knn search failed, ret=%d", ret); + return ret != 0 ? ret : IndexError_Runtime; + } ctx->topk_to_result(i); diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index c2b0dbf35..8804e4407 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -241,6 +242,11 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { EXPECT_EQ(original_ctx, streamer_ctx.get()); ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); EXPECT_EQ(original_ctx, streamer_ctx.get()); + + // I/O failures from the indexer must be propagated by the streamer instead + // of being converted into a successful search with incomplete results. + ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + EXPECT_NE(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); } TEST_F(DiskAnnSearcherTest, TestNodeCache) { From ff9256bf16ae757e677e38f581a1a07c967e459a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 21 Jul 2026 16:23:37 +0800 Subject: [PATCH 46/46] clang format --- src/core/algorithm/diskann/diskann_indexer.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 36ef567ac..936779050 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -1089,9 +1089,8 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); if (read_ret != 0) { - LOG_ERROR( - "cached_beam_search_by_group: reader_->read failed, ret=%d", - read_ret); + LOG_ERROR("cached_beam_search_by_group: reader_->read failed, ret=%d", + read_ret); ctx->set_error(true); return IndexError_Runtime; }