diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index 16fbb1cca..2d7cc8992 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -180,7 +180,7 @@ jobs: # DiskAnn libaio round: install libaio and re-run tests # ------------------------------------------------------------------ # - name: Install libaio runtime - if: matrix.platform == 'linux-x64' + if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64' run: | sudo apt-get update -y # libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older @@ -190,20 +190,35 @@ jobs: ldconfig -p | grep libaio || true shell: bash + - name: Verify DiskAnn uses libaio + if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64' + run: | + python - <<'PY' + import zvec + from zvec.typing import IOBackendType + + backend = zvec.io_backend_type() + print(f"DiskAnn I/O backend: {backend.name}") + assert backend == IOBackendType.LIBAIO, ( + f"expected LIBAIO after installing libaio, got {backend.name}" + ) + PY + shell: bash + - name: Run DiskAnn C++ Tests (w/ libaio) - if: matrix.platform == 'linux-x64' + if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64' run: | cd "$GITHUB_WORKSPACE/build" - ctest -R diskann --output-on-failure --parallel $NPROC + ctest -R diskann --no-tests=error --output-on-failure \ + --parallel $NPROC shell: bash - name: Run DiskAnn Python Tests (w/ libaio) - if: matrix.platform == 'linux-x64' + if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64' run: | cd "$GITHUB_WORKSPACE" python -m pytest python/tests/test_collection_diskann.py -v shell: bash - # Verify installing libaio does not affect existing non-DiskAnn tests. - name: Run HNSW Tests (w/ libaio) if: matrix.platform == 'linux-x64' diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..1565ab79e 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; 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) +# DiskAnn support: +# - Linux (x86_64, i686, i386, aarch64, arm64) with libaio (loaded via dlopen) +# - 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}") diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 4e5c703e1..e9531236f 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -15,6 +15,7 @@ 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) include_directories(${ZVEC_INCLUDE_DIR}) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index f5e6c08e4..538c65df3 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 @@ -27,9 +25,8 @@ 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 - _DISKANN_PRELOAD_REASON = None return None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 8daf8c770..458c93cad 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. 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. @@ -39,8 +39,14 @@ # 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 is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)", ) import zvec # noqa: E402 diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py index d566d5efc..895571eba 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -49,6 +49,7 @@ def test_enum_names(member, name): (DataType.FLOAT, 8), (IndexType.HNSW, 1), (IOBackendType.PREAD, 0), + (IOBackendType.THREAD_POOL_PREAD, 2), (MetricType.COSINE, 3), (QuantizeType.INT8, 2), (StatusCode.OK, 0), @@ -112,7 +113,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", "THREAD_POOL_PREAD"]) 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..62ed86598 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -52,12 +52,13 @@ from .zvec import create_and_open, init, open def io_backend_type() -> IOBackendType: """Returns the current I/O backend type for DiskAnn async disk reads as an IOBackendType enum (zvec.typing.IOBackendType). - IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise.""" + On Linux this is IOBackendType.LIBAIO when libaio is available; on macOS + this is IOBackendType.THREAD_POOL_PREAD.""" def io_backend_description() -> str: """Returns a human-readable description of the current I/O backend. - When only pread is available, includes instructions for installing - libaio to enable async I/O.""" + On Linux, when only pread is available, includes instructions for + installing libaio to enable async I/O.""" def set_default_jieba_dict_dir(dir: str) -> None: """Register the process-wide default jieba dict directory.""" diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index fab03f4b9..9c4d671dc 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -130,6 +130,8 @@ class IOBackendType: - PREAD: Synchronous pread() — no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). + - THREAD_POOL_PREAD: Blocking pread() executed by a shared worker pool, + with completion delivered through kqueue. Examples: >>> from zvec.typing import IOBackendType @@ -142,13 +144,18 @@ class IOBackendType: PREAD LIBAIO + + THREAD_POOL_PREAD """ LIBAIO: typing.ClassVar[IOBackendType] # value = PREAD: typing.ClassVar[IOBackendType] # value = + THREAD_POOL_PREAD: typing.ClassVar[ + IOBackendType + ] # value = __members__: typing.ClassVar[ dict[str, IOBackendType] - ] # value = {'PREAD': , 'LIBAIO': } + ] # value = {'PREAD': , 'LIBAIO': , 'THREAD_POOL_PREAD': } 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..30b268749 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -19,7 +19,7 @@ // 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 +// When no asynchronous backend is available, the caller should fall back to // synchronous pread(). // // Usage: @@ -29,6 +29,10 @@ #pragma once +#if defined(__APPLE__) && defined(__MACH__) +#include +#endif + #include #include @@ -40,6 +44,8 @@ inline const char *IOBackendTypeName(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio"; + case IOBackendType::kThreadPoolPread: + return "thread_pool_pread"; case IOBackendType::kPread: return "pread"; } @@ -52,6 +58,9 @@ inline const char *IOBackendDescription(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; + case IOBackendType::kThreadPoolPread: + return "Blocking pread() calls executed by a shared worker pool, with " + "completion delivered through kqueue."; 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 +72,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) selects libaio on Linux when available, worker-pool +// pread() on macOS, and synchronous pread() otherwise. // available(IOBackendType) tries a specific backend. // Use type() / name() to query the loaded backend without triggering a load. class IOBackend { @@ -74,13 +83,18 @@ class IOBackend { return instance; } - // Try to load the best available backend (libaio > pread). + // Try to load the best available backend for the current platform. // Returns the loaded backend type. // Idempotent — if already loaded, returns immediately. IOBackendType available() { if (type_ != IOBackendType::kPread) { return type_; } +#if defined(__APPLE__) && defined(__MACH__) +#if TARGET_OS_OSX + return available(IOBackendType::kThreadPoolPread); +#endif +#endif return available(IOBackendType::kLibAio); } @@ -91,6 +105,14 @@ class IOBackend { if (type_ == requested && type_ != IOBackendType::kPread) { return type_; } +#if defined(__APPLE__) && defined(__MACH__) +#if TARGET_OS_OSX + if (requested == IOBackendType::kThreadPoolPread) { + type_ = IOBackendType::kThreadPoolPread; + return type_; + } +#endif +#endif #if defined(__linux) || defined(__linux__) if (requested == IOBackendType::kLibAio) { if (LibAioLoader::Instance().load() && @@ -112,6 +134,10 @@ class IOBackend { return available() == IOBackendType::kLibAio; } + bool is_thread_pool_pread() { + return available() == IOBackendType::kThreadPoolPread; + } + // Returns the loaded backend type. IOBackendType type() const { return type_; diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f48b..507aa5188 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -228,18 +228,18 @@ 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, " - "IOBackendType.PREAD otherwise."); + "On Linux this is IOBackendType.LIBAIO when libaio is available; " + "on macOS this is IOBackendType.THREAD_POOL_PREAD."); // Returns a human-readable description of the I/O backend, including - // installation guidance for libaio when only pread is available. + // installation guidance for libaio on Linux when only pread is available. m.def( "io_backend_description", []() -> std::string { return ailego::current_io_backend_description(); }, "Returns a human-readable description of the current I/O backend. " - "When only pread is available, includes instructions for installing " - "libaio to enable async I/O."); + "On Linux, when only pread is available, includes instructions for " + "installing libaio to enable async I/O."); } -} // 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..70ba35123 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -146,6 +146,8 @@ 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(). +- THREAD_POOL_PREAD: Blocking pread() executed by a shared worker pool, with + completion delivered through kqueue. Examples: >>> from zvec.typing import IOBackendType @@ -153,7 +155,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("THREAD_POOL_PREAD", ailego::IOBackendType::kThreadPoolPread); } void ZVecPyTyping::bind_status(py::module_ &m) { @@ -245,4 +248,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/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index f874eba62..23aea7773 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 5f1b7bf7f..a0723056e 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -14,7 +14,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. 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") +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() @@ -31,4 +31,4 @@ cc_library( INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" -) \ 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 dd824ff23..e5729a19f 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -349,7 +349,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_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)); } diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 85be8c334..49fff0c2a 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -13,14 +13,23 @@ // limitations under the License. #include "diskann_file_reader.h" +#include #include #include +#include #include #include #include +#include +#include #include #include +#include #include +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#endif #define MAX_EVENTS 1024 @@ -60,8 +69,23 @@ int setup_io_ctx(IOContext &ctx) { return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - return ret; +#elif defined(__APPLE__) || defined(__MACH__) + // Each context owns one kqueue used for worker-completion notifications. + int kq = ::kqueue(); + if (kq == -1) { + // kqueue is an optimization, not a correctness requirement. Keep the + // context valid in synchronous-fallback mode when descriptor allocation + // fails. + LOG_WARN( + "kqueue() failed in setup_io_ctx; errno=%d, %s. " + "Falling back to synchronous pread", + errno, ::strerror(errno)); + ctx = -1; + return 0; + } + ctx = kq; + return 0; #else return 0; #endif @@ -73,31 +97,165 @@ int destroy_io_ctx(IOContext &ctx) { return 0; } int ret = LibAioLoader::Instance().io_destroy(ctx); - return ret; +#elif defined(__APPLE__) || defined(__MACH__) + if (ctx >= 0) { + ::close(ctx); + ctx = -1; + } + return 0; #else return 0; #endif } -static int execute_io_pread(int fd, std::vector &read_reqs) { - for (auto &req : read_reqs) { - ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); - if (bytes_read < 0) { - LOG_ERROR("pread failed; errno=%d, %s, offset=%lu, len=%lu", errno, - ::strerror(errno), (unsigned long)req.offset, - (unsigned long)req.len); - return IndexError_Runtime; +static int execute_one_pread(int fd, const AlignedRead &req) { + auto *buf = static_cast(req.buf); + uint64_t offset = req.offset; + uint64_t remaining = req.len; + + while (remaining > 0) { + ssize_t bytes_read = + ::pread(fd, buf, static_cast(remaining), offset); + if (bytes_read > 0) { + buf += bytes_read; + offset += static_cast(bytes_read); + remaining -= static_cast(bytes_read); + continue; } - if ((size_t)bytes_read != req.len) { - LOG_ERROR("pread short read; got=%zd, expected=%lu", bytes_read, - (unsigned long)req.len); + if (bytes_read == 0) { + LOG_ERROR("pread returned 0 (EOF); offset=%llu, remaining=%llu", + (unsigned long long)offset, (unsigned long long)remaining); return IndexError_Runtime; } + if (errno == EINTR) { + continue; + } + + LOG_ERROR("pread failed; errno=%d, %s, offset=%llu, len=%llu", errno, + ::strerror(errno), (unsigned long long)offset, + (unsigned long long)remaining); + return IndexError_Runtime; + } + + return 0; +} + +static int execute_io_pread(int fd, std::vector &read_reqs) { + for (const auto &req : read_reqs) { + int ret = execute_one_pread(fd, req); + if (ret != 0) { + return ret; + } } return 0; } +#if defined(__APPLE__) || defined(__MACH__) +namespace { + +struct MacIOBatch { + explicit MacIOBatch(size_t count) : remaining(count) {} + + std::atomic remaining; + std::atomic result{0}; +}; + +ailego::ThreadPool &mac_io_thread_pool() { + static ailego::ThreadPool pool( + std::max(1u, std::min(std::thread::hardware_concurrency(), 16u)), false); + return pool; +} + +void trigger_kqueue_completion(int kq, uintptr_t ident) { + struct kevent trigger; + EV_SET(&trigger, ident, EVFILT_USER, 0, NOTE_TRIGGER, 0, nullptr); + if (::kevent(kq, &trigger, 1, nullptr, 0, nullptr) == -1) { + LOG_WARN("Failed to trigger kqueue completion; errno=%d, %s", errno, + ::strerror(errno)); + } +} + +} // namespace + +// Execute batch I/O on macOS using a bounded worker pool for the blocking +// pread calls. The calling thread submits all requests and waits for +// EVFILT_USER completion notifications on its IOContext kqueue. This keeps +// multiple reads in flight without pretending that regular-file readiness is +// disk completion. +static int execute_io_kqueue(int kq, int fd, + std::vector &read_reqs) { + if (read_reqs.empty()) { + return 0; + } + + if (kq < 0) { + return execute_io_pread(fd, read_reqs); + } + + auto batch = std::make_shared(read_reqs.size()); + uintptr_t ident = reinterpret_cast(batch.get()); + + struct kevent registration; + EV_SET(®istration, ident, EVFILT_USER, EV_ADD | EV_CLEAR, 0, 0, nullptr); + if (::kevent(kq, ®istration, 1, nullptr, 0, nullptr) == -1) { + LOG_WARN( + "Failed to register kqueue completion event; errno=%d, %s. " + "Falling back to synchronous pread", + errno, ::strerror(errno)); + return execute_io_pread(fd, read_reqs); + } + + auto task_group = mac_io_thread_pool().make_group(); + for (const auto &req : read_reqs) { + task_group->execute([batch, fd, ident, kq, req]() { + int ret = execute_one_pread(fd, req); + if (ret != 0) { + int expected = 0; + batch->result.compare_exchange_strong(expected, ret); + } + + if (batch->remaining.fetch_sub(1, std::memory_order_acq_rel) == 1) { + trigger_kqueue_completion(kq, ident); + } + }); + } + + while (batch->remaining.load(std::memory_order_acquire) != 0) { + struct kevent completion; + struct timespec timeout { + 1, 0 + }; + int event_count = ::kevent(kq, nullptr, 0, &completion, 1, &timeout); + if (event_count < 0) { + if (errno == EINTR) { + continue; + } + LOG_WARN( + "Failed to wait for kqueue completion; errno=%d, %s. " + "Waiting for submitted worker tasks directly", + errno, ::strerror(errno)); + break; + } + } + + // The counter reaches zero before the final task returns from its closure. + // Waiting on the group guarantees no worker still references this batch or + // its kqueue registration before the event is deleted. + task_group->wait_finish(); + + struct kevent deletion; + EV_SET(&deletion, ident, EVFILT_USER, EV_DELETE, 0, 0, nullptr); + if (::kevent(kq, &deletion, 1, nullptr, 0, nullptr) == -1 && + errno != ENOENT) { + LOG_WARN("Failed to delete kqueue completion event; errno=%d, %s", errno, + ::strerror(errno)); + } + + return batch->result.load(std::memory_order_acquire); +} +#endif // __APPLE__ + int execute_io(IOContext ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) @@ -178,7 +336,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 } @@ -216,7 +380,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; } @@ -238,10 +401,30 @@ void LinuxAlignedFileReader::register_thread() { } } 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_WARN( + "kqueue() failed in register_thread; errno=%d, %s. " + "Falling back to synchronous pread", + errno, ::strerror(errno)); + ctx_map[thread_id] = ctx; + } else { + LOG_INFO("allocating kqueue ctx: %d", kq); + ctx = kq; + ctx_map[thread_id] = ctx; + } lk.unlock(); #endif } @@ -263,19 +446,36 @@ void LinuxAlignedFileReader::deregister_thread() { } // io_destroy is a syscall; keep it outside the lock to avoid blocking others - if (ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread) { + if (!ailego::IOBackend::Instance().is_pread()) { LibAioLoader::Instance().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 } void LinuxAlignedFileReader::deregister_all_threads() { #if (defined(__linux) || defined(__linux__)) std::unique_lock lk(ctx_mut); - bool aio_available = ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread; + bool aio_available = !ailego::IOBackend::Instance().is_pread(); for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; if (aio_available) { @@ -283,6 +483,15 @@ void LinuxAlignedFileReader::deregister_all_threads() { } } 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 } @@ -312,6 +521,32 @@ void LinuxAlignedFileReader::open(const std::string &fname) { ::strerror(errno)); } +#if defined(__APPLE__) || defined(__MACH__) + // 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 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)); + } + } +#endif + LOG_INFO("Opened file : %s", fname.c_str()); } @@ -338,6 +573,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 a1cb7c91a..a17461a1c 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -21,8 +21,17 @@ #include // dlopen-based libaio wrapper #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 @@ -52,9 +66,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 } }; @@ -79,6 +96,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) a shared worker pool performs pread +// operations and kqueue delivers completion notifications to the caller. class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..3c7a73a2d 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -709,8 +709,13 @@ int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, io_timer.reset(); - reader_->read(frontier_read_reqs, io_ctx); + int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("get_vector: reader_->read failed, ret=%d", read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( node_per_sector_, max_node_size_, frontier_neighbor.second, id); @@ -982,7 +987,7 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { group_topk_heap.limit(ctx->group_topk()); } - topk_heap.emplace(id, info); + group_topk_heap.emplace(id, info); } // stage 2, expand to reach group num as possible @@ -1086,8 +1091,14 @@ 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_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_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; } diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..51cbc875c 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -160,7 +160,10 @@ 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)) { + return ret; + } ctx->topk_to_result(i); @@ -338,6 +341,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/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 5af446733..dcb0236b7 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -271,4 +271,12 @@ int DiskAnnIndex::Merge(const std::vector &indexes, return 0; } -} // namespace zvec::core_interface \ No newline at end of file +ailego::IOBackendType DiskAnnIndex::io_backend_type() const { + ailego::IOBackendType type = ailego::current_io_backend_type(); + if (type == ailego::IOBackendType::kPread) { + LOG_WARN("%s", ailego::current_io_backend_description().c_str()); + } + return type; +} + +} // namespace zvec::core_interface diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..82f4ac7d7 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,19 +198,20 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable + // 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. + // On Linux, libaio is loaded eagerly (via dlopen); if it is missing, + // DiskAnn falls back to synchronous pread() with degraded performance. + // On macOS, DiskAnn uses kqueue. #if !DISKANN_SUPPORTED 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)."); #endif } diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 008b01514..47beff0fb 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -30,16 +30,19 @@ 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, // libaio loaded at runtime via dlopen() + kThreadPoolPread = 2 // Worker-pool pread() with completion notification }; // Returns the currently active I/O backend type. -// Triggers backend initialization on first call (libaio > pread). +// Triggers backend initialization on first call. Linux prefers libaio and +// macOS uses worker-pool pread() with kqueue completion notification. IOBackendType current_io_backend_type(); // Returns a human-readable description of the currently active I/O backend. -// When only pread is available, includes installation guidance for libaio. +// On Linux, when only pread is available, includes installation guidance for +// libaio. std::string current_io_backend_description(); } // namespace ailego diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..cde9d5216 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,6 +789,8 @@ typedef uint32_t zvec_io_backend_type_t; 0 /**< Synchronous pread() \u2014 no async I/O */ #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ +#define ZVEC_IO_BACKEND_TYPE_THREAD_POOL_PREAD \ + 2 /**< Worker-pool pread() with kqueue completion notification */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. @@ -796,7 +798,9 @@ typedef uint32_t zvec_io_backend_type_t; * Pure introspection \u2014 no side effects, no install hints. * * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD). + * (ZVEC_IO_BACKEND_TYPE_LIBAIO, + * ZVEC_IO_BACKEND_TYPE_THREAD_POOL_PREAD, or + * ZVEC_IO_BACKEND_TYPE_PREAD). */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -805,7 +809,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", "thread_pool_pread", "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 118d06d44..1c79b26df 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 @@ -376,6 +377,10 @@ class DiskAnnIndex : public Index { public: DiskAnnIndex() = default; + // Returns the I/O backend type currently loaded for DiskAnn async disk reads. + // If only pread is available, logs a hint to install libaio. + 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 83a289c52..581a9449b 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -119,4 +119,4 @@ class Collection { const std::string &column_name) const = 0; }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..fb2cf908e 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -20,6 +20,10 @@ #include #include +#if defined(__APPLE__) && defined(__MACH__) +#include +#endif + // Platform-specific headers #ifdef _WIN32 #include @@ -130,6 +134,33 @@ void test_version_functions(void) { TEST_END(); } +void test_io_backend_functions(void) { + TEST_START(); + + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name(ZVEC_IO_BACKEND_TYPE_PREAD), + "pread") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name(ZVEC_IO_BACKEND_TYPE_LIBAIO), + "libaio") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_THREAD_POOL_PREAD), + "thread_pool_pread") == 0); + + zvec_io_backend_type_t type = zvec_get_io_backend_type(); + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_PREAD || + type == ZVEC_IO_BACKEND_TYPE_LIBAIO || + type == ZVEC_IO_BACKEND_TYPE_THREAD_POOL_PREAD); +#if defined(__APPLE__) && defined(__MACH__) +#if TARGET_OS_OSX + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_THREAD_POOL_PREAD); +#else + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_PREAD); +#endif +#endif + TEST_ASSERT(zvec_get_io_backend_description() != NULL); + + TEST_END(); +} + void test_error_handling_functions(void) { TEST_START(); @@ -6375,6 +6406,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(); 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..a701137f4 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -0,0 +1,261 @@ +// 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 +#if defined(__APPLE__) && defined(__MACH__) +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core; + +namespace { + +constexpr size_t kPageSize = 4096; + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + close(); + ::unlink(path_); + } + + int fd() const { + return fd_; + } + + const char *path() const { + return path_; + } + + void close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + } + + bool write_all(const void *data, size_t size) { + const auto *bytes = static_cast(data); + size_t written = 0; + while (written < size) { + ssize_t ret = ::pwrite(fd_, bytes + written, size - written, written); + if (ret < 0 && errno == EINTR) { + continue; + } + if (ret <= 0) { + return false; + } + written += static_cast(ret); + } + return ::fsync(fd_) == 0; + } + + private: + char path_[64] = "/tmp/zvec-diskann-reader-XXXXXX"; + int fd_; +}; + +using AlignedBuffer = std::unique_ptr; + +AlignedBuffer make_aligned_buffer(size_t size) { + void *buffer = nullptr; + if (::posix_memalign(&buffer, kPageSize, size) != 0) { + return AlignedBuffer(nullptr, &std::free); + } + std::memset(buffer, 0, size); + return AlignedBuffer(buffer, &std::free); +} + +} // namespace + +#if defined(__APPLE__) && defined(__MACH__) +#if TARGET_OS_OSX +TEST(DiskAnnFileReaderTest, ReportsThreadPoolPreadBackend) { + EXPECT_EQ(zvec::ailego::IOBackendType::kThreadPoolPread, + zvec::ailego::current_io_backend_type()); +} +#else +TEST(DiskAnnFileReaderTest, ReportsPreadBackendOnNonMacApplePlatform) { + EXPECT_EQ(zvec::ailego::IOBackendType::kPread, + zvec::ailego::current_io_backend_type()); +} +#endif +#endif + +TEST(DiskAnnFileReaderTest, BatchAlignedReadsPreserveRequestOrder) { + 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); + } + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + AlignedBuffer output = make_aligned_buffer(source.size()); + ASSERT_NE(output, nullptr); + + std::vector requests; + requests.reserve(kPageCount); + for (size_t i = 0; i < kPageCount; ++i) { + const 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{}; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_EQ(reader.read(requests, ctx, false), 0); + + for (size_t i = 0; i < kPageCount; ++i) { + const 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(); +} + +TEST(DiskAnnFileReaderTest, ConcurrentContextsReadSameFile) { + constexpr size_t kPageCount = 16; + constexpr size_t kThreadCount = 4; + + 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); + } + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + LinuxAlignedFileReader reader; + reader.open(file.path()); + + std::atomic all_ok{true}; + std::vector workers; + workers.reserve(kThreadCount); + for (size_t thread_index = 0; thread_index < kThreadCount; ++thread_index) { + workers.emplace_back([&reader, &all_ok, thread_index]() { + AlignedBuffer output = make_aligned_buffer(kPageSize * kPageCount); + if (output == nullptr) { + all_ok.store(false, std::memory_order_relaxed); + return; + } + + std::vector requests; + requests.reserve(kPageCount); + for (size_t i = 0; i < kPageCount; ++i) { + size_t source_page = (i * 5 + thread_index) % kPageCount; + requests.emplace_back( + source_page * kPageSize, kPageSize, + static_cast(output.get()) + i * kPageSize); + } + + IOContext ctx{}; + if (setup_io_ctx(ctx) != 0) { + all_ok.store(false, std::memory_order_relaxed); + return; + } + + bool thread_ok = reader.read(requests, ctx, false) == 0; + for (size_t i = 0; thread_ok && i < kPageCount; ++i) { + size_t source_page = (i * 5 + thread_index) % kPageCount; + const auto *page = + static_cast(output.get()) + i * kPageSize; + for (size_t byte = 0; byte < kPageSize; ++byte) { + if (page[byte] != static_cast(source_page + 1)) { + thread_ok = false; + break; + } + } + } + if (destroy_io_ctx(ctx) != 0) { + thread_ok = false; + } + if (!thread_ok) { + all_ok.store(false, std::memory_order_relaxed); + } + }); + } + + for (auto &worker : workers) { + worker.join(); + } + + EXPECT_TRUE(all_ok.load(std::memory_order_relaxed)); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, ShortReadReturnsError) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize, 0x5a); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + AlignedBuffer output = make_aligned_buffer(kPageSize * 2); + ASSERT_NE(output, nullptr); + std::vector requests{ + {0, kPageSize * 2, output.get()}, + }; + + LinuxAlignedFileReader reader; + reader.open(file.path()); + + IOContext ctx{}; + ASSERT_EQ(setup_io_ctx(ctx), 0); + EXPECT_NE(reader.read(requests, ctx, false), 0); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, ReadBeforeOpenReturnsError) { + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{ + {0, kPageSize, output.get()}, + }; + + LinuxAlignedFileReader reader; + IOContext ctx{}; + EXPECT_NE(reader.read(requests, ctx, false), 0); +} diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index ff432886f..07e5fa431 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -220,6 +222,32 @@ 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()); + + // 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) { @@ -566,15 +594,20 @@ TEST_F(DiskAnnSearcherTest, TestGroup) { total_time += t2 - t1; auto &group_result = ctx->group_result(); + ASSERT_EQ(group_num, group_result.size()); + std::set seen_group_ids; for (uint32_t i = 0; i < group_result.size(); ++i) { const std::string &group_id = group_result[i].group_id(); auto &result = group_result[i].docs(); + ASSERT_TRUE(seen_group_ids.insert(group_id).second); ASSERT_GT(result.size(), 0); + ASSERT_LE(result.size(), group_topk); std::cout << "Group ID: " << group_id << std::endl; for (uint32_t j = 0; j < result.size(); ++j) { + EXPECT_EQ(group_id, groupbyFunc(result[j].key())); std::cout << "\tKey: " << result[j].key() << std::fixed << std::setprecision(3) << ", Score: " << result[j].score() << std::endl; @@ -726,6 +759,12 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { float vector_value = *((float *)(knnResult[0].vector_string().data())); ASSERT_EQ(vector_value, i); } + + ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + std::string vector_after_truncate; + EXPECT_EQ(IndexError_Runtime, searcher->get_vector(doc_cnt - 1, linearCtx, + vector_after_truncate)); + EXPECT_TRUE(vector_after_truncate.empty()); } TEST_F(DiskAnnSearcherTest, TestRnnSearch) { diff --git a/tests/db/crash_recovery/write_recovery_test.cc b/tests/db/crash_recovery/write_recovery_test.cc index 702e9e498..9f878b1eb 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);