From 26d515a76c0317d6d10676ebea656bd0beed6881 Mon Sep 17 00:00:00 2001 From: Ron Marcus Date: Wed, 31 Dec 2025 11:27:12 +0200 Subject: [PATCH 1/3] feat: [NCS] Add Near-Compute Storage (NCS) core infrastructure This commit introduces the core infrastructure for Near-Compute Storage (NCS), a shared hot storage tier designed to decouple compute and storage in Milvus. Key changes: - Added `NcsSingleton`, `NcsFactory`, and `NcsConnector` interfaces to abstract storage backends. - Implemented `InMemoryNcs` and `InMemNcsConnector` backed by `InMemoryKV` for testing. - Implemented `RedisNcs` and `RedisNcsConnector` as a reference remote storage backend. - Exposed C-compatible API (`ncs_c.h`) to allow integration with the Milvus Go control plane. - Added unit tests for both In-Memory and Redis implementations. This infrastructure enables upper layers (Knowhere, Milvus) to interact with disaggregated storage through a unified Key-Value API. issue: milvus-io/milvus#45178 --- CMakeLists.txt | 4 + Dockerfile.builder | 50 +++++++++ README.md | 24 ++++ conanfile.py | 3 +- include/common/EasyAssert.h | 15 +-- include/common/SpanBytes.h | 19 ++++ include/ncs/InMemNcsConnector.h | 33 ++++++ include/ncs/InMemoryKV.h | 32 ++++++ include/ncs/InMemoryNcs.h | 30 +++++ include/ncs/RedisNcs.h | 37 +++++++ include/ncs/RedisNcsConnector.h | 34 ++++++ include/ncs/ncs.h | 124 +++++++++++++++++++++ src/ncs/InMemNcsConnector.cpp | 84 ++++++++++++++ src/ncs/InMemoryKV.cpp | 60 ++++++++++ src/ncs/InMemoryNcs.cpp | 46 ++++++++ src/ncs/RedisNcs.cpp | 173 +++++++++++++++++++++++++++++ src/ncs/RedisNcsConnector.cpp | 184 +++++++++++++++++++++++++++++++ src/ncs/ncs.cpp | 110 ++++++++++++++++++ test/CMakeLists.txt | 1 + test/test_ncs/CMakeLists.txt | 55 +++++++++ test/test_ncs/test_inmem_ncs.cpp | 93 ++++++++++++++++ test/test_ncs/test_redis_ncs.cpp | 140 +++++++++++++++++++++++ 22 files changed, 1343 insertions(+), 8 deletions(-) create mode 100644 Dockerfile.builder create mode 100644 include/common/SpanBytes.h create mode 100644 include/ncs/InMemNcsConnector.h create mode 100644 include/ncs/InMemoryKV.h create mode 100644 include/ncs/InMemoryNcs.h create mode 100644 include/ncs/RedisNcs.h create mode 100644 include/ncs/RedisNcsConnector.h create mode 100644 include/ncs/ncs.h create mode 100644 src/ncs/InMemNcsConnector.cpp create mode 100644 src/ncs/InMemoryKV.cpp create mode 100644 src/ncs/InMemoryNcs.cpp create mode 100644 src/ncs/RedisNcs.cpp create mode 100644 src/ncs/RedisNcsConnector.cpp create mode 100644 src/ncs/ncs.cpp create mode 100644 test/test_ncs/CMakeLists.txt create mode 100644 test/test_ncs/test_inmem_ncs.cpp create mode 100644 test/test_ncs/test_redis_ncs.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e4018b..5fcb8ed 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -25,6 +25,7 @@ set( CMAKE_EXPORT_COMPILE_COMMANDS ON ) set(MILVUS_COMMON_WORKSPACE ${CMAKE_CURRENT_SOURCE_DIR}) include_directories(${MILVUS_COMMON_WORKSPACE}/include) + set(CMAKE_CXX_FLAGS "-Wall -fPIC ${CMAKE_CXX_FLAGS}") if (WITH_ASAN) @@ -40,12 +41,14 @@ find_package(gflags REQUIRED) find_package(glog REQUIRED) find_package(fmt REQUIRED) find_package(prometheus-cpp REQUIRED) +find_package(hiredis REQUIRED) list(APPEND COMMON_LINKER_LIBS glog::glog) list(APPEND COMMON_LINKER_LIBS prometheus-cpp::core prometheus-cpp::push) list(APPEND COMMON_LINKER_LIBS fmt::fmt-header-only) list(APPEND COMMON_LINKER_LIBS Folly::folly) list(APPEND COMMON_LINKER_LIBS gflags::gflags) +list(APPEND COMMON_LINKER_LIBS hiredis::hiredis) list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_trace) list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_exporter_ostream_span) @@ -54,6 +57,7 @@ list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_exporter_otlp_ht list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_exporter_jaeger_trace) file(GLOB_RECURSE SRC_FILES src/*.cpp src/*.cc) + if(__X86_64) set_source_files_properties(src/knowhere/thread_pool.cc PROPERTIES COMPILE_OPTIONS "-msse4.2" diff --git a/Dockerfile.builder b/Dockerfile.builder new file mode 100644 index 0000000..235865a --- /dev/null +++ b/Dockerfile.builder @@ -0,0 +1,50 @@ +# Builder image for milvus-common unit tests (based on ubuntu-22.04) +# Usage examples: +# docker build -f Dockerfile.builder -t milvus-common-builder:latest . +# docker run --rm -it -v$(pwd):/workspace -w /workspace milvus-common-builder:latest # get an interactive shell + +FROM ubuntu:22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV LANG=C.UTF-8 +ENV LC_ALL=C.UTF-8 + +# Install toolchain + dependencies +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + software-properties-common \ + ca-certificates \ + curl \ + gnupg2 \ + cmake \ + libopenblas-dev \ + libaio-dev \ + python3 \ + python3-pip \ + build-essential \ + libpci3 \ + redis-server \ + && add-apt-repository ppa:ubuntu-toolchain-r/test -y \ + && apt-get update \ + && apt-get install -y --no-install-recommends gcc-12 g++-12 \ + && rm -rf /var/lib/apt/lists/* + +# Make gcc-12 / g++-12 available via CC/CXX env vars +ENV CC=gcc-12 +ENV CXX=g++-12 + +# Install a specific conan +RUN pip3 install --no-cache-dir conan==1.61.0 + +# Add default conan remote if reachable (best-effort) +RUN conan remote add default-conan-local https://milvus01.jfrog.io/artifactory/api/conan/default-conan-local || true + +WORKDIR /workspace + +# Create entrypoint script that starts Redis and then runs bash +RUN echo '#!/bin/bash\nredis-server --daemonize yes\nexec "$@"' > /entrypoint.sh \ + && chmod +x /entrypoint.sh + +# Default entrypoint starts Redis, then runs bash for interactive use +ENTRYPOINT ["/entrypoint.sh"] +CMD ["/bin/bash"] diff --git a/README.md b/README.md index cb17972..94db86f 100644 --- a/README.md +++ b/README.md @@ -25,3 +25,27 @@ conan build .. # run ut ./test/test_cachinglayer/cachinglayer_test ``` + +## Build using the provided Docker builder image (alternative) + +If you don't want to install the toolchain and conan locally, you can use the included +`Dockerfile.builder` image which mirrors the CI environment (Ubuntu 22.04, gcc-12, +conan 1.61). + +```bash +# Build the builder image (run from repository root) +docker build -f Dockerfile.builder -t milvus-common-builder:latest . + +# Start an interactive shell with the repository mounted at /workspace +docker run --rm -it -v "$(pwd)":/workspace -v "${HOME}/.conan":/root/.conan --add-host=host.docker.internal:host-gateway -w /workspace milvus-common-builder:latest + +# Inside the container run the same build commands as CI: +mkdir -p build && cd build +conan install .. --build=missing -o with_ut=True -o with_asan=True -s compiler.libcxx=libstdc++11 -s compiler.version=12 -s build_type=Release +conan build .. + +# Run tests inside the container (example) +./test/test_cachinglayer/cachinglayer_test +``` + + diff --git a/conanfile.py b/conanfile.py index 3696edb..21acd16 100644 --- a/conanfile.py +++ b/conanfile.py @@ -26,7 +26,8 @@ class MilvusCommonConan(ConanFile): "libevent/2.1.12#4fd19d10d3bed63b3a8952c923454bc0", "openssl/3.1.2#02594c4c0a6e2b4feb3cd15119993597", "folly/2023.10.30.10@milvus/dev", - "boost/1.82.0" + "boost/1.82.0", + "hiredis/1.2.0" ) options = { diff --git a/include/common/EasyAssert.h b/include/common/EasyAssert.h index 122b1e2..3666cfd 100644 --- a/include/common/EasyAssert.h +++ b/include/common/EasyAssert.h @@ -63,14 +63,15 @@ enum ErrorCode { MemAllocateFailed = 2034, MemAllocateSizeNotMatch = 2035, MmapError = 2036, + NcsUploadError = 2037, // timeout or cancel related - FollyOtherException = 2037, - FollyCancel = 2038, - OutOfRange = 2039, - GcpNativeError = 2040, - TextIndexNotFound = 2041, - InvalidParameter = 2042, - InsufficientResource = 2043, + FollyOtherException = 2038, + FollyCancel = 2039, + OutOfRange = 2040, + GcpNativeError = 2041, + TextIndexNotFound = 2042, + InvalidParameter = 2043, + InsufficientResource = 2044, KnowhereError = 2099 }; diff --git a/include/common/SpanBytes.h b/include/common/SpanBytes.h new file mode 100644 index 0000000..05c1493 --- /dev/null +++ b/include/common/SpanBytes.h @@ -0,0 +1,19 @@ +#pragma once + +#include + +namespace milvus { + +class SpanBytes { +public: + SpanBytes(void* data, size_t size) : data_(data), size_(size) {} + + void* data() const { return data_; } + size_t size() const { return size_; } + +private: + void* data_; + size_t size_; +}; + +} // namespace milvus \ No newline at end of file diff --git a/include/ncs/InMemNcsConnector.h b/include/ncs/InMemNcsConnector.h new file mode 100644 index 0000000..fc05d3c --- /dev/null +++ b/include/ncs/InMemNcsConnector.h @@ -0,0 +1,33 @@ +#pragma once + +#include "ncs/ncs.h" +#include "ncs/InMemoryKV.h" +#include + +namespace milvus { + +using std::make_unique; + +class InMemNcsConnector : public NcsConnector { +public: + friend class InMemoryNcsConnectorCreator; + + // Interface implementations + std::vector multiGet(const std::vector& keys, const std::vector& buffs) override; + std::vector multiPut(const std::vector& keys, const std::vector& buffs) override; + std::vector multiDelete(const std::vector& keys) override; + +private: + explicit InMemNcsConnector(uint64_t bucketId); // Private constructor +}; + +class InMemoryNcsConnectorCreator : public NcsConnectorCreator { +public: + NcsConnector* factoryMethod(const NcsDescriptor* descriptor) override; + const std::string& getKind() const override; + +private: + static const std::string KIND; +}; + +} // namespace milvus \ No newline at end of file diff --git a/include/ncs/InMemoryKV.h b/include/ncs/InMemoryKV.h new file mode 100644 index 0000000..d351a14 --- /dev/null +++ b/include/ncs/InMemoryKV.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include +#include + +#include "common/SpanBytes.h" + +namespace milvus { + +class InMemoryKV { +public: + static InMemoryKV* Instance(); + + // Put a value into an existing bucket. Returns true on success, false if the + // bucket does not exist. + bool put(uint64_t bucketId, uint32_t key, const SpanBytes& value); + bool createBucket(uint64_t bucketId); + bool deleteBucket(uint64_t bucketId); + bool hasBucket(uint64_t bucketId) const; + bool get(uint64_t bucketId, uint32_t key, const SpanBytes& buff) const; + bool deleteKey(uint64_t bucketId, uint32_t key); +private: + InMemoryKV() = default; + using BucketMap = std::unordered_map>; + using DataMap = std::unordered_map; + + DataMap data_; +}; + +} // namespace milvus diff --git a/include/ncs/InMemoryNcs.h b/include/ncs/InMemoryNcs.h new file mode 100644 index 0000000..becaa2d --- /dev/null +++ b/include/ncs/InMemoryNcs.h @@ -0,0 +1,30 @@ +#pragma once + +#include "ncs/ncs.h" +#include "ncs/InMemoryKV.h" +#include +#include + + +namespace milvus { + +class InMemoryNcs : public Ncs { +public: + NcsStatus createBucket(uint64_t bucketId) override; + NcsStatus deleteBucket(uint64_t bucketId) override; + NcsBucketStatus getBucketNcsStatus(uint64_t bucketId) override; + bool isBucketExist(uint64_t bucketId) override; + ~InMemoryNcs() override = default; +private: + InMemoryNcs() = default; + friend class InMemoryNcsFactory; +}; + +class InMemoryNcsFactory : public NcsFactory { +public: + static const std::string KIND; + std::unique_ptr createNcs(const json& params = json{}) override; + const std::string& getKind() const override; +}; + +} // namespace milvus diff --git a/include/ncs/RedisNcs.h b/include/ncs/RedisNcs.h new file mode 100644 index 0000000..f866797 --- /dev/null +++ b/include/ncs/RedisNcs.h @@ -0,0 +1,37 @@ +#pragma once + +#include "ncs/ncs.h" +#include "log/Log.h" +#include +#include +#include +#include + +namespace milvus { + +class RedisNcs : public Ncs { +public: + NcsStatus createBucket(uint64_t bucketId) override; + NcsStatus deleteBucket(uint64_t bucketId) override; + NcsBucketStatus getBucketNcsStatus(uint64_t bucketId) override; + bool isBucketExist(uint64_t bucketId) override; + ~RedisNcs() override; + +private: + explicit RedisNcs(const std::string& host, int port); + redisContext* context_; + std::string host_; + int port_; + std::mutex mutex_; + + friend class RedisNcsFactory; +}; + +class RedisNcsFactory : public NcsFactory { +public: + static const std::string KIND; + std::unique_ptr createNcs(const json& params = json{}) override; + const std::string& getKind() const override; +}; + +} // namespace milvus diff --git a/include/ncs/RedisNcsConnector.h b/include/ncs/RedisNcsConnector.h new file mode 100644 index 0000000..8dfcd85 --- /dev/null +++ b/include/ncs/RedisNcsConnector.h @@ -0,0 +1,34 @@ +#pragma once + +#include "ncs/ncs.h" +#include "log/Log.h" +#include +#include +#include + +namespace milvus { + +class RedisNcsConnector : public NcsConnector { +public: + ~RedisNcsConnector() override; + std::vector multiGet(const std::vector& keys, const std::vector& buffs) override; + std::vector multiPut(const std::vector& keys, const std::vector& buffs) override; + std::vector multiDelete(const std::vector& keys) override; + +private: + explicit RedisNcsConnector(uint64_t bucketId, const std::string& host, int port); + redisContext* context_; + std::string host_; + int port_; + + friend class RedisNcsConnectorCreator; +}; + +class RedisNcsConnectorCreator : public NcsConnectorCreator { +public: + static const std::string KIND; + NcsConnector* factoryMethod(const NcsDescriptor* descriptor) override; + const std::string& getKind() const override; +}; + +} // namespace milvus diff --git a/include/ncs/ncs.h b/include/ncs/ncs.h new file mode 100644 index 0000000..bef34f0 --- /dev/null +++ b/include/ncs/ncs.h @@ -0,0 +1,124 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include "nlohmann/json.hpp" + +#include "common/SpanBytes.h" + + +namespace milvus { + +using std::string; +using std::vector; +using std::unique_ptr; +using std::unordered_map; +using json = nlohmann::json; + +enum class NcsStatus { + OK, + ERROR +}; + +class NcsDescriptor { +public: + NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const json& extras); + virtual ~NcsDescriptor() = default; + const std::string& getKind() const; + uint64_t getbucketId() const; + const json& getExtras() const { return extras_; } + + NLOHMANN_DEFINE_TYPE_INTRUSIVE(NcsDescriptor, ncsKind_, bucketId_, extras_) + NcsDescriptor() = default; +private: + std::string ncsKind_; + uint64_t bucketId_; + json extras_; +}; + +class NcsBucketStatus{ + // TBD (capacity, occupancy, etc) +}; + + +class Ncs{ +public: + virtual NcsStatus createBucket(uint64_t bucketId) = 0; + virtual NcsStatus deleteBucket(uint64_t bucketId) = 0; + virtual NcsBucketStatus getBucketNcsStatus(uint64_t bucketId) = 0; + virtual bool isBucketExist(uint64_t bucketId) = 0; + virtual ~Ncs() = default; +protected: + Ncs() = default; +}; + +class NcsFactory { +public: + virtual std::unique_ptr createNcs(const json& params = json{}) = 0; + virtual const std::string& getKind() const = 0; + virtual ~NcsFactory() = default; +}; + +class NcsFactoryRegistry { +public: + static NcsFactoryRegistry& Instance(); + void registerFactory(std::unique_ptr factory); + std::unique_ptr createNcs(const std::string& kind, const json& params = json{}); + bool hasKind(const std::string& kind) const; + ~NcsFactoryRegistry() = default; + +private: + NcsFactoryRegistry() = default; + std::unordered_map> registry_; +}; + +class NcsSingleton final{ +public: + static void initNcs(const std::string& kind, const json& extras = json{}); + static Ncs* Instance(); + +private: + inline static std::string kind_; + inline static json extras_; + inline static std::unique_ptr instance_ = nullptr; +}; + + +class NcsConnector { +public: + virtual ~NcsConnector() = default; + virtual std::vector multiGet(const std::vector& keys, const std::vector& buffs) = 0; + virtual std::vector multiPut(const std::vector& keys, const std::vector& buffs) = 0; + virtual std::vector multiDelete(const std::vector& keys) = 0; + +protected: + explicit NcsConnector(uint64_t bucketId) : bucketId_(bucketId) {} + const uint64_t bucketId_; +}; + +class NcsConnectorCreator { +public: + virtual ~NcsConnectorCreator() = default; + virtual NcsConnector* factoryMethod(const NcsDescriptor* descriptor) = 0; + virtual const std::string& getKind() const = 0; +}; + + + +class NcsConnectorFactory { +public: + static NcsConnectorFactory& Instance(); + NcsConnector* createConnector(const NcsDescriptor* descriptor); + void registerCreator(std::unique_ptr creator); + ~NcsConnectorFactory() = default; + +private: + NcsConnectorFactory() = default; + std::unordered_map> registry_; +}; + +} // namespace milvus \ No newline at end of file diff --git a/src/ncs/InMemNcsConnector.cpp b/src/ncs/InMemNcsConnector.cpp new file mode 100644 index 0000000..2c0a608 --- /dev/null +++ b/src/ncs/InMemNcsConnector.cpp @@ -0,0 +1,84 @@ +#include "ncs/InMemNcsConnector.h" +#include "log/Log.h" +#include +#include + +namespace milvus { + +using std::make_unique; + + +InMemNcsConnector::InMemNcsConnector(uint64_t bucketId) +: NcsConnector(bucketId) { + // Constructor performs no blocking checks; factory will validate bucket. +} + + +std::vector InMemNcsConnector::multiGet(const std::vector& keys, + const std::vector& buffs) { + std::vector results; + results.reserve(keys.size()); + + for (size_t i = 0; i < keys.size(); ++i) { + bool success = InMemoryKV::Instance()->get(bucketId_, keys[i], buffs[i]); + results.push_back(success ? NcsStatus::OK : NcsStatus::ERROR); + } + + return results; +} + +std::vector InMemNcsConnector::multiPut(const std::vector& keys, + const std::vector& buffs) { + std::vector results; + results.reserve(keys.size()); + + for (size_t i = 0; i < keys.size(); ++i) { + bool ok = InMemoryKV::Instance()->put(bucketId_, keys[i], buffs[i]); + results.push_back(ok ? NcsStatus::OK : NcsStatus::ERROR); + } + + return results; +} + +std::vector InMemNcsConnector::multiDelete(const std::vector& keys) { + std::vector results; + results.reserve(keys.size()); + + for (const auto& key : keys) { + bool success = InMemoryKV::Instance()->deleteKey(bucketId_, key); + results.push_back(success ? NcsStatus::OK : NcsStatus::ERROR); + } + + return results; +} + +const std::string InMemoryNcsConnectorCreator::KIND = "in_memory"; + +NcsConnector* InMemoryNcsConnectorCreator::factoryMethod(const NcsDescriptor* descriptor) { + if (descriptor->getKind() != KIND) { + LOG_ERROR("[NCS] InMemoryNcsConnectorCreator received incompatible descriptor kind: {}", descriptor->getKind()); + return nullptr; + } + // If bucket doesn't exist in backing KV, fail early and return nullptr. + if (!InMemoryKV::Instance()->hasBucket(descriptor->getbucketId())) { + LOG_ERROR("[NCS] InMemNcsConnector creation failed: bucket {} does not exist.", descriptor->getbucketId()); + return nullptr; + } + return new InMemNcsConnector(descriptor->getbucketId()); +} + +const std::string& InMemoryNcsConnectorCreator::getKind() const { + return KIND; +} + +// Register the connector type on startup +namespace { + struct RegisterInMemConnector { + RegisterInMemConnector() { + NcsConnectorFactory::Instance().registerCreator( + std::make_unique()); + } + } registerInMemConnector; +} + +} // namespace milvus \ No newline at end of file diff --git a/src/ncs/InMemoryKV.cpp b/src/ncs/InMemoryKV.cpp new file mode 100644 index 0000000..ddd469b --- /dev/null +++ b/src/ncs/InMemoryKV.cpp @@ -0,0 +1,60 @@ +#include "ncs/InMemoryKV.h" +#include + +namespace milvus { + +InMemoryKV* InMemoryKV::Instance() { + static InMemoryKV instance; + return &instance; +} + +bool InMemoryKV::put(uint64_t bucketId, uint32_t key, const SpanBytes& value) { + auto bucket_it = data_.find(bucketId); + if (bucket_it == data_.end()) { + // bucket not created + return false; + } + + auto& bucket = bucket_it->second; + std::vector value_copy(static_cast(value.data()), + static_cast(value.data()) + value.size()); + bucket[key] = std::move(value_copy); + return true; +} + +bool InMemoryKV::createBucket(uint64_t bucketId) { + data_.emplace(bucketId, BucketMap{}); + // If bucket already exists, emplace does nothing; treat as success. + return true; +} + +bool InMemoryKV::deleteBucket(uint64_t bucketId) { + auto erased = data_.erase(bucketId); + return erased > 0; +} + +bool InMemoryKV::hasBucket(uint64_t bucketId) const { + return data_.find(bucketId) != data_.end(); +} + +bool InMemoryKV::get(uint64_t bucketId, uint32_t key, const SpanBytes& buff) const { + auto bucket_it = data_.find(bucketId); + if (bucket_it == data_.end()) return false; + + auto value_it = bucket_it->second.find(key); + if (value_it == bucket_it->second.end()) return false; + + if (value_it->second.size() > buff.size()) return false; + + std::memcpy(buff.data(), value_it->second.data(), value_it->second.size()); + return true; +} + +bool InMemoryKV::deleteKey(uint64_t bucketId, uint32_t key) { + auto bucket_it = data_.find(bucketId); + if (bucket_it == data_.end()) return false; + auto erased = bucket_it->second.erase(key); + return erased > 0; +} + +} // namespace milvus diff --git a/src/ncs/InMemoryNcs.cpp b/src/ncs/InMemoryNcs.cpp new file mode 100644 index 0000000..5b68627 --- /dev/null +++ b/src/ncs/InMemoryNcs.cpp @@ -0,0 +1,46 @@ +#include "ncs/InMemoryNcs.h" +#include + +namespace milvus { + +// InMemoryNcs implementation +NcsStatus InMemoryNcs::createBucket(uint64_t bucketId) { + bool ok = InMemoryKV::Instance()->createBucket(bucketId); + return ok ? NcsStatus::OK : NcsStatus::ERROR; +} + +NcsStatus InMemoryNcs::deleteBucket(uint64_t bucketId) { + bool ok = InMemoryKV::Instance()->deleteBucket(bucketId); + return ok ? NcsStatus::OK : NcsStatus::ERROR; +} + +NcsBucketStatus InMemoryNcs::getBucketNcsStatus(uint64_t /*bucketId*/) { + return NcsBucketStatus(); +} + +bool InMemoryNcs::isBucketExist(uint64_t bucketId) { + return InMemoryKV::Instance()->hasBucket(bucketId); +} + +// InMemoryNcsFactory implementation +const std::string InMemoryNcsFactory::KIND = "in_memory"; + +std::unique_ptr InMemoryNcsFactory::createNcs(const json& params) { + return std::unique_ptr(new InMemoryNcs()); +} + +const std::string& InMemoryNcsFactory::getKind() const { + return KIND; +} + +// Register the factory on startup +namespace { + struct RegisterInMemoryNcsFactory { + RegisterInMemoryNcsFactory() { + NcsFactoryRegistry::Instance().registerFactory( + std::make_unique()); + } + } registerInMemoryNcsFactory; +} + +} // namespace milvus diff --git a/src/ncs/RedisNcs.cpp b/src/ncs/RedisNcs.cpp new file mode 100644 index 0000000..788d2e8 --- /dev/null +++ b/src/ncs/RedisNcs.cpp @@ -0,0 +1,173 @@ +#include "ncs/RedisNcs.h" +#include "log/Log.h" +#include +#include +#include +#include + +namespace milvus { + +// RedisNcs implementation +RedisNcs::RedisNcs(const std::string& host, int port) + : context_(nullptr), host_(host), port_(port) { + context_ = redisConnect(host.c_str(), port); + if (context_ == nullptr || context_->err) { + if (context_) { + LOG_ERROR("[RedisNcs] Redis connection error: {}", context_->errstr); + redisFree(context_); + context_ = nullptr; + } else { + LOG_ERROR("[RedisNcs] Redis connection error: can't allocate redis context"); + } + throw std::runtime_error("Failed to connect to Redis at " + host + ":" + std::to_string(port)); + } + LOG_INFO("[RedisNcs] Connected to Redis at {}:{}", host, port); +} + +RedisNcs::~RedisNcs() { + if (context_) { + redisFree(context_); + context_ = nullptr; + } +} + +NcsStatus RedisNcs::createBucket(uint64_t bucketId) { + std::lock_guard lock(mutex_); + if (!context_) { + LOG_ERROR("[RedisNcs] Redis context is null"); + return NcsStatus::ERROR; + } + + std::string key = "bucket_" + std::to_string(bucketId) + "_valid"; + redisReply* reply = (redisReply*)redisCommand(context_, "SET %s true", key.c_str()); + + if (reply == nullptr) { + LOG_ERROR("[RedisNcs] Failed to create bucket {}: {}", bucketId, context_->errstr); + return NcsStatus::ERROR; + } + + if (reply->type == REDIS_REPLY_STATUS && std::string(reply->str) == "OK") { + freeReplyObject(reply); + LOG_INFO("[RedisNcs] Created bucket {}", bucketId); + return NcsStatus::OK; + } + + LOG_ERROR("[RedisNcs] Failed to create bucket {}. Reply type: {}, str: {}", + bucketId, reply->type, (reply->str ? reply->str : "null")); + freeReplyObject(reply); + return NcsStatus::ERROR; +} + +NcsStatus RedisNcs::deleteBucket(uint64_t bucketId) { + std::lock_guard lock(mutex_); + if (!context_) { + LOG_ERROR("[RedisNcs] Redis context is null"); + return NcsStatus::ERROR; + } + + std::string pattern = "bucket_" + std::to_string(bucketId) + "_*"; + std::vector keysToDelete; + + // Use SCAN to find all keys matching the pattern + int cursor = 0; + do { + redisReply* reply = (redisReply*)redisCommand(context_, + "SCAN %d MATCH %s COUNT 100", cursor, pattern.c_str()); + + if (reply == nullptr || reply->type != REDIS_REPLY_ARRAY) { + LOG_ERROR("[RedisNcs] Failed to scan keys for bucket {}", bucketId); + if (reply) freeReplyObject(reply); + return NcsStatus::ERROR; + } + + // Parse cursor + cursor = std::atoi(reply->element[0]->str); + + // Parse keys + redisReply* keysArray = reply->element[1]; + for (size_t i = 0; i < keysArray->elements; ++i) { + keysToDelete.push_back(keysArray->element[i]->str); + } + + freeReplyObject(reply); + } while (cursor != 0); + + // Delete all found keys using UNLINK (async delete) + if (!keysToDelete.empty()) { + // Build UNLINK command + std::string cmd = "UNLINK"; + for (const auto& key : keysToDelete) { + cmd += " " + key; + } + + redisReply* reply = (redisReply*)redisCommand(context_, cmd.c_str()); + if (reply == nullptr) { + LOG_ERROR("[RedisNcs] Failed to delete keys for bucket {}: {}", + bucketId, context_->errstr); + return NcsStatus::ERROR; + } + + LOG_INFO("[RedisNcs] Deleted {} keys for bucket {}", reply->integer, bucketId); + freeReplyObject(reply); + } + + return NcsStatus::OK; +} + +NcsBucketStatus RedisNcs::getBucketNcsStatus(uint64_t /*bucketId*/) { + return NcsBucketStatus(); +} + +bool RedisNcs::isBucketExist(uint64_t bucketId) { + std::lock_guard lock(mutex_); + if (!context_) { + LOG_ERROR("[RedisNcs] Redis context is null"); + return false; + } + + std::string key = "bucket_" + std::to_string(bucketId) + "_valid"; + redisReply* reply = (redisReply*)redisCommand(context_, "EXISTS %s", key.c_str()); + + if (reply == nullptr) { + LOG_ERROR("[RedisNcs] Failed to check bucket existence: {}", context_->errstr); + return false; + } + + bool exists = (reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); + freeReplyObject(reply); + + return exists; +} + +// RedisNcsFactory implementation +const std::string RedisNcsFactory::KIND = "redis"; + +std::unique_ptr RedisNcsFactory::createNcs(const json& params) { + if (!params.contains("redis_host")) { + throw std::runtime_error("RedisNcsFactory: 'redis_host' is required in params"); + } + if (!params.contains("redis_port")) { + throw std::runtime_error("RedisNcsFactory: 'redis_port' is required in params"); + } + + std::string host = params["redis_host"].get(); + int port = params["redis_port"].get(); + + return std::unique_ptr(new RedisNcs(host, port)); +} + +const std::string& RedisNcsFactory::getKind() const { + return KIND; +} + +// Register the factory on startup +namespace { + struct RegisterRedisNcsFactory { + RegisterRedisNcsFactory() { + NcsFactoryRegistry::Instance().registerFactory( + std::make_unique()); + } + } registerRedisNcsFactory; +} + +} // namespace milvus diff --git a/src/ncs/RedisNcsConnector.cpp b/src/ncs/RedisNcsConnector.cpp new file mode 100644 index 0000000..305e30b --- /dev/null +++ b/src/ncs/RedisNcsConnector.cpp @@ -0,0 +1,184 @@ +#include "ncs/RedisNcsConnector.h" +#include "log/Log.h" +#include +#include +#include +#include +#include + +namespace milvus { + +// RedisNcsConnector implementation +RedisNcsConnector::RedisNcsConnector(uint64_t bucketId, const std::string& host, int port) + : NcsConnector(bucketId), context_(nullptr), host_(host), port_(port) { + context_ = redisConnect(host.c_str(), port); + if (context_ == nullptr || context_->err) { + if (context_) { + LOG_ERROR("[RedisNcsConnector] Redis connection error: {}", context_->errstr); + redisFree(context_); + context_ = nullptr; + } else { + LOG_ERROR("[RedisNcsConnector] Redis connection error: can't allocate redis context"); + } + throw std::runtime_error("Failed to connect to Redis at " + host + ":" + std::to_string(port)); + } + LOG_DEBUG("[RedisNcsConnector] Connected to Redis at {}:{} for bucket {}", host, port, bucketId); +} + +RedisNcsConnector::~RedisNcsConnector() { + if (context_) { + redisFree(context_); + context_ = nullptr; + } +} + +std::vector RedisNcsConnector::multiGet( + const std::vector& keys, + const std::vector& buffs) { + + std::vector results(keys.size(), NcsStatus::ERROR); + + if (!context_) { + LOG_ERROR("[RedisNcsConnector] Redis context is null"); + return results; + } + + if (keys.size() != buffs.size()) { + LOG_ERROR("[RedisNcsConnector] Keys and buffers size mismatch"); + return results; + } + + for (size_t i = 0; i < keys.size(); ++i) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + + redisReply* reply = (redisReply*)redisCommand(context_, "GET %s", redisKey.c_str()); + + if (reply == nullptr) { + LOG_ERROR("[RedisNcsConnector] Failed to GET key {}: {}", redisKey, context_->errstr); + continue; + } + + if (reply->type == REDIS_REPLY_STRING) { + size_t dataSize = reply->len; + if (dataSize <= buffs[i].size()) { + std::memcpy(buffs[i].data(), reply->str, dataSize); + results[i] = NcsStatus::OK; + } else { + LOG_ERROR("[RedisNcsConnector] Buffer too small for key {}: need {}, have {}", + redisKey, dataSize, buffs[i].size()); + } + } else if (reply->type == REDIS_REPLY_NIL) { + LOG_WARN("[RedisNcsConnector] Key {} does not exist", redisKey); + } else { + LOG_ERROR("[RedisNcsConnector] Unexpected reply type for key {}", redisKey); + } + + freeReplyObject(reply); + } + + return results; +} + +std::vector RedisNcsConnector::multiPut( + const std::vector& keys, + const std::vector& buffs) { + + std::vector results(keys.size(), NcsStatus::ERROR); + + if (!context_) { + LOG_ERROR("[RedisNcsConnector] Redis context is null"); + return results; + } + + if (keys.size() != buffs.size()) { + LOG_ERROR("[RedisNcsConnector] Keys and buffers size mismatch"); + return results; + } + + for (size_t i = 0; i < keys.size(); ++i) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + + redisReply* reply = (redisReply*)redisCommand(context_, "SET %s %b", + redisKey.c_str(), buffs[i].data(), buffs[i].size()); + + if (reply == nullptr) { + LOG_ERROR("[RedisNcsConnector] Failed to SET key {}: {}", redisKey, context_->errstr); + continue; + } + + if (reply->type == REDIS_REPLY_STATUS && std::string(reply->str) == "OK") { + results[i] = NcsStatus::OK; + } else { + LOG_ERROR("[RedisNcsConnector] Failed to SET key {}", redisKey); + } + + freeReplyObject(reply); + } + + return results; +} + +std::vector RedisNcsConnector::multiDelete(const std::vector& keys) { + std::vector results(keys.size(), NcsStatus::ERROR); + + if (!context_) { + LOG_ERROR("[RedisNcsConnector] Redis context is null"); + return results; + } + + for (size_t i = 0; i < keys.size(); ++i) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + + redisReply* reply = (redisReply*)redisCommand(context_, "DEL %s", redisKey.c_str()); + + if (reply == nullptr) { + LOG_ERROR("[RedisNcsConnector] Failed to DEL key {}: {}", redisKey, context_->errstr); + continue; + } + + if (reply->type == REDIS_REPLY_INTEGER) { + results[i] = NcsStatus::OK; + } else { + LOG_ERROR("[RedisNcsConnector] Failed to DEL key {}", redisKey); + } + + freeReplyObject(reply); + } + + return results; +} + +// RedisNcsConnectorCreator implementation +const std::string RedisNcsConnectorCreator::KIND = "redis"; + +NcsConnector* RedisNcsConnectorCreator::factoryMethod(const NcsDescriptor* descriptor) { + const json& extras = descriptor->getExtras(); + + if (!extras.contains("redis_host")) { + throw std::runtime_error("RedisNcsConnectorCreator: 'redis_host' is required in descriptor extras"); + } + if (!extras.contains("redis_port")) { + throw std::runtime_error("RedisNcsConnectorCreator: 'redis_port' is required in descriptor extras"); + } + + std::string host = extras["redis_host"].get(); + int port = extras["redis_port"].get(); + + return new RedisNcsConnector(descriptor->getbucketId(), host, port); +} + +const std::string& RedisNcsConnectorCreator::getKind() const { + return KIND; +} + +// Register the connector creator on startup +namespace { + struct RegisterRedisNcsConnectorCreator { + RegisterRedisNcsConnectorCreator() { + NcsConnectorFactory::Instance().registerCreator( + std::make_unique()); + } + } registerRedisNcsConnectorCreator; +} + +} // namespace milvus diff --git a/src/ncs/ncs.cpp b/src/ncs/ncs.cpp new file mode 100644 index 0000000..1f942e5 --- /dev/null +++ b/src/ncs/ncs.cpp @@ -0,0 +1,110 @@ +#include "ncs/ncs.h" +#include +#include +#include "log/Log.h" + +namespace milvus { + +// NcsFactoryRegistry implementation +NcsFactoryRegistry& NcsFactoryRegistry::Instance() { + static NcsFactoryRegistry instance; + return instance; +} + +void NcsFactoryRegistry::registerFactory(std::unique_ptr factory) { + auto kind = factory->getKind(); + registry_[kind] = std::move(factory); +} + +std::unique_ptr NcsFactoryRegistry::createNcs(const std::string& kind, const json& params) { + LOG_DEBUG("[NCS] Creating NCS of kind: {} with params: {}", kind, params.dump()); + auto it = registry_.find(kind); + if (it != registry_.end()) { + return it->second->createNcs(params); + } + + std::string registered_kinds_str = + std::accumulate(registry_.begin(), registry_.end(), + std::string(), + [](const std::string& a, const auto& b) { + return a.empty() ? b.first : a + ", " + b.first; + }); + LOG_WARN("[NCS] No registered factory for NCS kind: {}. Registered kinds: {}", kind, registered_kinds_str); + return nullptr; +} + +bool NcsFactoryRegistry::hasKind(const std::string& kind) const { + return registry_.find(kind) != registry_.end(); +} + +// NcsSingleton implementation +void NcsSingleton::initNcs(const std::string& kind, const json& extras) { + if (!NcsFactoryRegistry::Instance().hasKind(kind)) { + throw std::runtime_error("NCS Factory kind '" + kind + "' is not registered."); + } + kind_ = kind; + extras_ = extras; +} + +Ncs* NcsSingleton::Instance() { + if(kind_.empty()){ + throw std::runtime_error("NCS Factory have not been set yet. Use initNcs() first."); + } + if(!instance_){ + instance_ = NcsFactoryRegistry::Instance().createNcs(kind_, extras_); + if (!instance_) { + throw std::runtime_error("NCS Factory kind '" + kind_ + "' is not registered."); + } + } + return instance_.get(); +} + +// NcsConnectorFactory implementation +NcsConnector* NcsConnectorFactory::createConnector(const NcsDescriptor* descriptor) { + LOG_DEBUG("[NCS] Creating NcsConnector of kind: {} with extra params: {}", descriptor->getKind(), descriptor->getExtras().dump()); + auto it = registry_.find(descriptor->getKind()); + if (it != registry_.end()) { + return it->second->factoryMethod(descriptor); + } + + std::string registered_kinds_str = + std::accumulate(registry_.begin(), registry_.end(), + std::string(), + [](const std::string& a, const auto& b) { + return a.empty() ? b.first : a + ", " + b.first; + }); + LOG_WARN("[NCS] No registered creator for NCS kind: {}. registered kinds: {}", descriptor->getKind(), registered_kinds_str); + return nullptr; +} + +void NcsConnectorFactory::registerCreator(std::unique_ptr creator) { + auto kind = creator->getKind(); + registry_[kind] = std::move(creator); +} + +// NcsConnectorFactory singleton implementation +NcsConnectorFactory& NcsConnectorFactory::Instance() { + static NcsConnectorFactory instance; + return instance; +} + +// Already defaulted in header + +// NcsDescriptor implementation +NcsDescriptor::NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const json& extras) + : ncsKind_(ncsKind), bucketId_(bucketId), extras_(extras) { +} + +const std::string& NcsDescriptor::getKind() const { + return ncsKind_; +} + +uint64_t NcsDescriptor::getbucketId() const { + return bucketId_; +} + +// NcsConnector constructor is already defined in header + +// Already defined above + +} // namespace milvus \ No newline at end of file diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index a602d54..e90fc01 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -10,6 +10,7 @@ # or implied. See the License for the specific language governing permissions and limitations under the License add_subdirectory(test_cachinglayer) +add_subdirectory(test_ncs) find_package(GTest CONFIG REQUIRED) find_package(OpenMP REQUIRED) diff --git a/test/test_ncs/CMakeLists.txt b/test/test_ncs/CMakeLists.txt new file mode 100644 index 0000000..1b47f26 --- /dev/null +++ b/test/test_ncs/CMakeLists.txt @@ -0,0 +1,55 @@ +# NCS tests + +find_package(GTest CONFIG REQUIRED) +find_package(OpenMP REQUIRED) +find_package(BLAS REQUIRED) + +set(IN_MEM_NCS_TEST_FILES + ../init_gtest.cpp + test_inmem_ncs.cpp +) + +add_executable(inmem_ncs_test + ${IN_MEM_NCS_TEST_FILES} +) + +set(REDIS_NCS_TEST_FILES + ../init_gtest.cpp + test_redis_ncs.cpp +) + +add_executable(redis_ncs_test + ${REDIS_NCS_TEST_FILES} +) + +# Ensure SSE4.2 CRC intrinsics are available to match Folly build (x86_64) +if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64" AND NOT MSVC) + target_compile_options(inmem_ncs_test PRIVATE -msse4.2) + target_compile_options(redis_ncs_test PRIVATE -msse4.2) +endif() + +target_link_libraries(inmem_ncs_test + GTest::gtest + GTest::gmock + milvus-common + ${COMMON_LINKER_LIBS} + pthread + atomic + gomp + openblas +) + +target_link_libraries(redis_ncs_test + GTest::gtest + GTest::gmock + milvus-common + ${COMMON_LINKER_LIBS} + pthread + atomic + gomp + openblas +) + + +install(TARGETS inmem_ncs_test redis_ncs_test DESTINATION unittest) + diff --git a/test/test_ncs/test_inmem_ncs.cpp b/test/test_ncs/test_inmem_ncs.cpp new file mode 100644 index 0000000..aa2a435 --- /dev/null +++ b/test/test_ncs/test_inmem_ncs.cpp @@ -0,0 +1,93 @@ +#include +#include "ncs/InMemNcsConnector.h" +#include "ncs/InMemoryNcs.h" // provides InMemoryNcsFactory +#include + +namespace milvus { +namespace { + +TEST(InMemNcsConnectorTest, BasicOperations) { + const uint64_t bucketId = 1; + // Register the trivial InMemoryNcs factory in the singleton and get Ncs instance + NcsSingleton::initNcs(InMemoryNcsFactory::KIND); + Ncs* ncs = NcsSingleton::Instance(); + auto createResult = ncs->createBucket(bucketId); + EXPECT_EQ(createResult, NcsStatus::OK); + + // Create descriptor and connector + auto descriptor = std::make_unique(NcsDescriptor("in_memory", bucketId, json::object())); + auto connector = std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor.get())); + + ASSERT_NE(connector, nullptr); + + // Prepare test data with varying sizes + std::vector keys = {1, 2, 3}; + std::vector> values = { + std::vector(100, 0x11), // 100 bytes + std::vector(200, 0x22), // 200 bytes + std::vector(300, 0x33) // 300 bytes + }; + std::vector valueSpans; + for (auto& value : values) { + valueSpans.emplace_back(value.data(), value.size()); + } + + // Test multiPut + auto putResults = connector->multiPut(keys, valueSpans); + ASSERT_EQ(putResults.size(), keys.size()); + for (const auto& status : putResults) { + EXPECT_EQ(status, NcsStatus::OK); + } + + // Prepare buffers for reading + std::vector> readBuffers; + std::vector readSpans; + for (const auto& value : values) { + readBuffers.emplace_back(value.size()); + readSpans.emplace_back(readBuffers.back().data(), readBuffers.back().size()); + } + + // Test multiGet + auto getResults = connector->multiGet(keys, readSpans); + ASSERT_EQ(getResults.size(), keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + EXPECT_EQ(getResults[i], NcsStatus::OK); + EXPECT_EQ(readBuffers[i], values[i]); + } + + // Test with buffer too small + std::vector smallBuffer(50); + SpanBytes smallSpan(smallBuffer.data(), smallBuffer.size()); + auto getResult = connector->multiGet({keys[2]}, {smallSpan}); + ASSERT_EQ(getResult.size(), 1); + EXPECT_EQ(getResult[0], NcsStatus::ERROR); // Should fail as buffer is too small + + // Test multiDelete for specific keys + auto deleteResult = connector->multiDelete(keys); + ASSERT_EQ(deleteResult.size(), keys.size()); + for (const auto& status : deleteResult) { + EXPECT_EQ(status, NcsStatus::OK); + } + + // Verify data is deleted + auto getResultsAfterDelete = connector->multiGet(keys, readSpans); + for (const auto& status : getResultsAfterDelete) { + EXPECT_EQ(status, NcsStatus::ERROR); + } + + // Test bucket deletion + auto bucketDeleteResult = ncs->deleteBucket(bucketId); + EXPECT_EQ(bucketDeleteResult, NcsStatus::OK); + + // Negative test: try to put into a bucket that was not created + const uint32_t missingBucket = 999; + auto descriptor2 = std::make_unique("in_memory", missingBucket, json::object()); + auto connector2 = std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor2.get())); + // Connector creation should fail (return nullptr) because bucket wasn't created + EXPECT_EQ(connector2, nullptr); +} + +} // namespace +} // namespace milvus \ No newline at end of file diff --git a/test/test_ncs/test_redis_ncs.cpp b/test/test_ncs/test_redis_ncs.cpp new file mode 100644 index 0000000..de6dca1 --- /dev/null +++ b/test/test_ncs/test_redis_ncs.cpp @@ -0,0 +1,140 @@ +#include +#include "ncs/RedisNcs.h" +#include "ncs/RedisNcsConnector.h" +#include +#include + +namespace milvus { +namespace { + +// Note: These tests require a Redis server running on localhost:6379 +// To skip these tests if Redis is not available, they can be marked as DISABLED_ + +TEST(RedisNcsTest, BasicBucketOperations) { + // Initialize RedisNcs with host and port + json config; + config["redis_host"] = "localhost"; + config["redis_port"] = 6379; + + NcsSingleton::initNcs(RedisNcsFactory::KIND, config); + Ncs* ncs = NcsSingleton::Instance(); + + const uint64_t bucketId = 100; + + // Test bucket creation + auto createResult = ncs->createBucket(bucketId); + EXPECT_EQ(createResult, NcsStatus::OK); + + // Test bucket existence check + bool exists = ncs->isBucketExist(bucketId); + EXPECT_TRUE(exists); + + // Test bucket deletion + auto deleteResult = ncs->deleteBucket(bucketId); + EXPECT_EQ(deleteResult, NcsStatus::OK); + + // Verify bucket no longer exists + exists = ncs->isBucketExist(bucketId); + EXPECT_FALSE(exists); +} + +TEST(RedisNcsConnectorTest, MultiGetPutDelete) { + const uint64_t bucketId = 101; + + // Initialize RedisNcs + json config; + config["redis_host"] = "localhost"; + config["redis_port"] = 6379; + + NcsSingleton::initNcs(RedisNcsFactory::KIND, config); + Ncs* ncs = NcsSingleton::Instance(); + + // Create bucket + auto createResult = ncs->createBucket(bucketId); + ASSERT_EQ(createResult, NcsStatus::OK); + + auto descriptor = std::make_unique("redis", bucketId, config); + auto connector = std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor.get())); + + ASSERT_NE(connector, nullptr); + + // Prepare test data + std::vector keys = {1, 2, 3}; + std::vector> values = { + std::vector(100, 0x11), // 100 bytes + std::vector(200, 0x22), // 200 bytes + std::vector(300, 0x33) // 300 bytes + }; + + // Create SpanBytes for put operation + std::vector putBuffs; + for (const auto& value : values) { + putBuffs.emplace_back(const_cast(value.data()), value.size()); + } + + // Test multiPut + auto putResults = connector->multiPut(keys, putBuffs); + EXPECT_EQ(putResults.size(), keys.size()); + for (const auto& result : putResults) { + EXPECT_EQ(result, NcsStatus::OK); + } + + // Test multiGet + std::vector> getBuffers(keys.size()); + std::vector getBuffs; + for (size_t i = 0; i < keys.size(); ++i) { + getBuffers[i].resize(values[i].size()); + getBuffs.emplace_back(getBuffers[i].data(), getBuffers[i].size()); + } + + auto getResults = connector->multiGet(keys, getBuffs); + EXPECT_EQ(getResults.size(), keys.size()); + for (size_t i = 0; i < getResults.size(); ++i) { + EXPECT_EQ(getResults[i], NcsStatus::OK); + EXPECT_EQ(getBuffers[i], values[i]); + } + + // Test multiDelete + auto deleteResults = connector->multiDelete(keys); + EXPECT_EQ(deleteResults.size(), keys.size()); + for (const auto& result : deleteResults) { + EXPECT_EQ(result, NcsStatus::OK); + } + + // Verify deletion - get should return empty/error + auto verifyResults = connector->multiGet(keys, getBuffs); + for (const auto& result : verifyResults) { + EXPECT_NE(result, NcsStatus::OK); + } + + // Cleanup + ncs->deleteBucket(bucketId); +} + +TEST(RedisNcsTest, LargeBucketId) { + // Initialize RedisNcs with host and port + json config; + config["redis_host"] = "localhost"; + config["redis_port"] = 6379; + + NcsSingleton::initNcs(RedisNcsFactory::KIND, config); + Ncs* ncs = NcsSingleton::Instance(); + + const uint64_t bucketId = 463281186922614943ULL; + + // Test bucket creation + auto createResult = ncs->createBucket(bucketId); + EXPECT_EQ(createResult, NcsStatus::OK); + + // Test bucket existence check + bool exists = ncs->isBucketExist(bucketId); + EXPECT_TRUE(exists); + + // Test bucket deletion + auto deleteResult = ncs->deleteBucket(bucketId); + EXPECT_EQ(deleteResult, NcsStatus::OK); +} + +} // namespace +} // namespace milvus From b0a5e9ba1720ed6508e9a6aafd0307f9967d668b Mon Sep 17 00:00:00 2001 From: Ron Marcus Date: Tue, 6 Jan 2026 12:55:13 +0200 Subject: [PATCH 2/3] - Add documentation for ncs.h - Improve Redis NCS connector to use Redis pipelining - Merge unit tests for all NCS kinds - Add concurency test using multiple NcsConnector --- include/ncs/RedisNcsConnector.h | 21 +- include/ncs/ncs.h | 171 +++++++++++ src/ncs/RedisNcsConnector.cpp | 181 +++++++++-- test/test_ncs/CMakeLists.txt | 38 +-- test/test_ncs/test_inmem_ncs.cpp | 93 ------ test/test_ncs/test_ncs_all.cpp | 513 +++++++++++++++++++++++++++++++ test/test_ncs/test_redis_ncs.cpp | 140 --------- 7 files changed, 866 insertions(+), 291 deletions(-) delete mode 100644 test/test_ncs/test_inmem_ncs.cpp create mode 100644 test/test_ncs/test_ncs_all.cpp delete mode 100644 test/test_ncs/test_redis_ncs.cpp diff --git a/include/ncs/RedisNcsConnector.h b/include/ncs/RedisNcsConnector.h index 8dfcd85..1402c2a 100644 --- a/include/ncs/RedisNcsConnector.h +++ b/include/ncs/RedisNcsConnector.h @@ -4,10 +4,22 @@ #include "log/Log.h" #include #include +#include #include namespace milvus { +/** + * @brief Redis-based NCS connector with single connection. + * + * This connector uses a single Redis connection and is NOT thread-safe. + * Each thread should have its own connector instance. + * + * Thread-safety is achieved at a higher level (e.g., NCSReader) by using + * thread_local connectors, one per thread. + * + * Uses Redis pipelining for efficient batch operations. + */ class RedisNcsConnector : public NcsConnector { public: ~RedisNcsConnector() override; @@ -17,7 +29,14 @@ class RedisNcsConnector : public NcsConnector { private: explicit RedisNcsConnector(uint64_t bucketId, const std::string& host, int port); - redisContext* context_; + + /** + * @brief Ensure connection is valid, reconnect if needed. + * @return true if connection is valid, false otherwise. + */ + bool ensureConnected(); + + redisContext* ctx_ = nullptr; std::string host_; int port_; diff --git a/include/ncs/ncs.h b/include/ncs/ncs.h index bef34f0..89b69a8 100644 --- a/include/ncs/ncs.h +++ b/include/ncs/ncs.h @@ -1,5 +1,61 @@ #pragma once +/** + * @file ncs.h + * @brief Near Compute Storage (NCS) abstraction layer. + * + * This header defines the interfaces for NCS implementations (InMemory, Redis). + * All implementations must adhere to the contracts specified in this file. + * + * ## Implementation Requirements + * + * ### Thread Safety - Single-Threaded Connector Model + * + * **NcsConnector instances are NOT required to be thread-safe.** + * + * Concurrency is handled at a higher level by the consumer (e.g., NCSReader in DiskANN). + * Each thread should create and use its own NcsConnector instance via thread_local storage. + * + * #### Why this design? + * - Simplifies NcsConnector implementations (no need for connection pools or mutexes) + * - Eliminates lock contention in high-throughput scenarios + * - Natural isolation - each thread has its own connection to the backend + * - Backend connections (for example: Redis) are typically not thread-safe anyway + * + * #### Concurrency Model (handled by NCSReader): + * ``` + * Thread 1 ──► NcsConnector instance 1 ──►┐ + * Thread 2 ──► NcsConnector instance 2 ──►├──► NCS Backend (for example: Redis) + * Thread 3 ──► NcsConnector instance 3 ──►┘ + * ``` + * + * #### Consumer Responsibilities: + * - NCSReader uses `thread_local` storage to create one connector per thread + * - Each thread's connector is created lazily on first access + * - Connectors are destroyed when the NCSReader is destroyed (for current thread) + * or when threads exit + * + * #### Implementation Note for Backends: + * - **Redis**: Each connector holds a single `redisContext*` (not thread-safe) + * - **InMemory**: Uses shared unordered_map; put/delete operations are NOT thread-safe + * (acceptable if data is populated before concurrent reads begin) + * + * ### Bucket Management + * - Buckets must be created via `Ncs::createBucket()` before any operations. + * - `NcsConnectorCreator::factoryMethod()` MUST return `nullptr` if the bucket does not exist. + * Implementations should query the backend directly to verify bucket existence. + * + * ### Error Handling + * - `multiGet` returns `NcsStatus::ERROR` for keys that don't exist or if buffer is too small. + * - `multiPut` returns `NcsStatus::ERROR` if the write operation fails. + * - `multiDelete` returns `NcsStatus::OK` even if the key doesn't exist (idempotent delete). + * + * ### Performance Guidelines + * - Implementations should use batching/pipelining for `multi*` operations when possible. + * - Redis: Use `redisAppendCommand()` + `redisGetReply()` for pipelining. + * - Avoid sequential operations in loops; prefer bulk commands supported by the backend. + */ + #include #include #include @@ -45,12 +101,44 @@ class NcsBucketStatus{ }; +/** + * @brief Abstract interface for NCS bucket management. + * + * Provides bucket lifecycle operations. Each NCS implementation (InMemory, Redis) + * must implement this interface. + * + * Thread Safety: Implementations must ensure thread-safe bucket operations. + */ class Ncs{ public: + /** + * @brief Create a new bucket. + * @param bucketId Unique identifier for the bucket. + * @return NcsStatus::OK on success, NcsStatus::ERROR on failure. + */ virtual NcsStatus createBucket(uint64_t bucketId) = 0; + + /** + * @brief Delete a bucket and all its contents. + * @param bucketId The bucket to delete. + * @return NcsStatus::OK on success (or if bucket doesn't exist), NcsStatus::ERROR on failure. + */ virtual NcsStatus deleteBucket(uint64_t bucketId) = 0; + + /** + * @brief Get status information about a bucket. + * @param bucketId The bucket to query. + * @return Bucket status information. + */ virtual NcsBucketStatus getBucketNcsStatus(uint64_t bucketId) = 0; + + /** + * @brief Check if a bucket exists. + * @param bucketId The bucket to check. + * @return true if bucket exists, false otherwise. + */ virtual bool isBucketExist(uint64_t bucketId) = 0; + virtual ~Ncs() = default; protected: Ncs() = default; @@ -81,6 +169,20 @@ class NcsSingleton final{ static void initNcs(const std::string& kind, const json& extras = json{}); static Ncs* Instance(); + /** + * @brief Reset the singleton instance. For testing purposes only. + * + * This destroys the current NCS instance and clears the configuration, + * allowing a new NCS type to be initialized via initNcs(). + * + * WARNING: Only use in test code. Do not use in production. + */ + static void reset() { + instance_.reset(); + kind_.clear(); + extras_.clear(); + } + private: inline static std::string kind_; inline static json extras_; @@ -88,11 +190,54 @@ class NcsSingleton final{ }; +/** + * @brief Abstract interface for NCS data operations within a bucket. + * + * Provides key-value operations (get, put, delete) for a specific bucket. + * + * ## Thread Safety + * NcsConnector instances are NOT required to be thread-safe. + * + * ## Performance Requirements + * Implementations SHOULD use batching/pipelining for multi* operations: + * - Redis: Use pipelining instead of sequential GET/SET commands + */ class NcsConnector { public: virtual ~NcsConnector() = default; + + /** + * @brief Read multiple key-value pairs. + * @param keys Vector of keys to read. + * @param buffs Vector of buffers to receive the values. Must be same size as keys. + * Each buffer must be large enough to hold the corresponding value. + * @return Vector of status codes, one per key: + * - NcsStatus::OK if read succeeded + * - NcsStatus::ERROR if key doesn't exist or buffer too small + */ virtual std::vector multiGet(const std::vector& keys, const std::vector& buffs) = 0; + + /** + * @brief Write multiple key-value pairs. + * @param keys Vector of keys to write. + * @param buffs Vector of buffers containing the values. Must be same size as keys. + * @return Vector of status codes, one per key: + * - NcsStatus::OK if write succeeded + * - NcsStatus::ERROR if write failed + * + * Note: If a key already exists, its value is overwritten. + */ virtual std::vector multiPut(const std::vector& keys, const std::vector& buffs) = 0; + + /** + * @brief Delete multiple keys. + * @param keys Vector of keys to delete. + * @return Vector of status codes, one per key: + * - NcsStatus::OK if delete succeeded (including if key didn't exist) + * - NcsStatus::ERROR if delete operation failed + * + * Note: Deleting a non-existent key is not an error (idempotent). + */ virtual std::vector multiDelete(const std::vector& keys) = 0; protected: @@ -100,10 +245,36 @@ class NcsConnector { const uint64_t bucketId_; }; +/** + * @brief Factory interface for creating NcsConnector instances. + * + * Each NCS implementation must provide a concrete NcsConnectorCreator. + * + * ## Implementation Requirements + * - `factoryMethod()` MUST check bucket existence before creating a connector. + * - If the bucket does not exist, `factoryMethod()` MUST return `nullptr`. + * - Bucket existence check MUST query the backend directly (e.g., Redis EXISTS) + * rather than using NcsSingleton, to avoid coupling. + */ class NcsConnectorCreator { public: virtual ~NcsConnectorCreator() = default; + + /** + * @brief Create an NcsConnector for the given descriptor. + * @param descriptor Contains bucket ID, NCS kind, and configuration. + * @return Pointer to new NcsConnector, or nullptr if: + * - The bucket does not exist + * - Required configuration is missing + * - Connection to backend fails + * + * Note: Caller takes ownership of the returned pointer. + */ virtual NcsConnector* factoryMethod(const NcsDescriptor* descriptor) = 0; + + /** + * @brief Get the NCS kind this creator handles (e.g., "redis", "in_memory"). + */ virtual const std::string& getKind() const = 0; }; diff --git a/src/ncs/RedisNcsConnector.cpp b/src/ncs/RedisNcsConnector.cpp index 305e30b..73c44f1 100644 --- a/src/ncs/RedisNcsConnector.cpp +++ b/src/ncs/RedisNcsConnector.cpp @@ -8,28 +8,58 @@ namespace milvus { +// ============================================================================ // RedisNcsConnector implementation +// ============================================================================ + RedisNcsConnector::RedisNcsConnector(uint64_t bucketId, const std::string& host, int port) - : NcsConnector(bucketId), context_(nullptr), host_(host), port_(port) { - context_ = redisConnect(host.c_str(), port); - if (context_ == nullptr || context_->err) { - if (context_) { - LOG_ERROR("[RedisNcsConnector] Redis connection error: {}", context_->errstr); - redisFree(context_); - context_ = nullptr; + : NcsConnector(bucketId), ctx_(nullptr), host_(host), port_(port) { + + ctx_ = redisConnect(host.c_str(), port); + if (ctx_ == nullptr || ctx_->err) { + if (ctx_) { + LOG_ERROR("[RedisNcsConnector] Connection error: {}", ctx_->errstr); + redisFree(ctx_); + ctx_ = nullptr; } else { - LOG_ERROR("[RedisNcsConnector] Redis connection error: can't allocate redis context"); + LOG_ERROR("[RedisNcsConnector] Cannot allocate redis context"); } throw std::runtime_error("Failed to connect to Redis at " + host + ":" + std::to_string(port)); } - LOG_DEBUG("[RedisNcsConnector] Connected to Redis at {}:{} for bucket {}", host, port, bucketId); + + LOG_DEBUG("[RedisNcsConnector] Created connector for bucket {} to {}:{}", + bucketId, host, port); } RedisNcsConnector::~RedisNcsConnector() { - if (context_) { - redisFree(context_); - context_ = nullptr; + if (ctx_) { + redisFree(ctx_); + ctx_ = nullptr; } + LOG_DEBUG("[RedisNcsConnector] Destroyed connector for bucket {}", bucketId_); +} + +bool RedisNcsConnector::ensureConnected() { + if (ctx_ != nullptr && !ctx_->err) { + return true; + } + + if (ctx_) { + redisFree(ctx_); + } + + ctx_ = redisConnect(host_.c_str(), port_); + if (ctx_ == nullptr || ctx_->err) { + if (ctx_) { + LOG_ERROR("[RedisNcsConnector] Reconnection error: {}", ctx_->errstr); + redisFree(ctx_); + ctx_ = nullptr; + } + return false; + } + + LOG_DEBUG("[RedisNcsConnector] Reconnected to {}:{}", host_, port_); + return true; } std::vector RedisNcsConnector::multiGet( @@ -38,8 +68,7 @@ std::vector RedisNcsConnector::multiGet( std::vector results(keys.size(), NcsStatus::ERROR); - if (!context_) { - LOG_ERROR("[RedisNcsConnector] Redis context is null"); + if (keys.empty()) { return results; } @@ -48,13 +77,37 @@ std::vector RedisNcsConnector::multiGet( return results; } + if (!ensureConnected()) { + LOG_ERROR("[RedisNcsConnector] Not connected to Redis"); + return results; + } + + // Pipeline all GET commands for (size_t i = 0; i < keys.size(); ++i) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); - redisReply* reply = (redisReply*)redisCommand(context_, "GET %s", redisKey.c_str()); + if (redisAppendCommand(ctx_, "GET %s", redisKey.c_str()) != REDIS_OK) { + LOG_ERROR("[RedisNcsConnector] Failed to append GET command for key {}", redisKey); + for (size_t j = 0; j < i; ++j) { + redisReply* reply = nullptr; + redisGetReply(ctx_, (void**)&reply); + if (reply) freeReplyObject(reply); + } + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + redisReply* reply = nullptr; + + if (redisGetReply(ctx_, (void**)&reply) != REDIS_OK) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + LOG_ERROR("[RedisNcsConnector] Failed to get reply for key {}: {}", + redisKey, ctx_->errstr ? ctx_->errstr : "unknown error"); + continue; + } if (reply == nullptr) { - LOG_ERROR("[RedisNcsConnector] Failed to GET key {}: {}", redisKey, context_->errstr); continue; } @@ -64,13 +117,16 @@ std::vector RedisNcsConnector::multiGet( std::memcpy(buffs[i].data(), reply->str, dataSize); results[i] = NcsStatus::OK; } else { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Buffer too small for key {}: need {}, have {}", redisKey, dataSize, buffs[i].size()); } } else if (reply->type == REDIS_REPLY_NIL) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_WARN("[RedisNcsConnector] Key {} does not exist", redisKey); } else { - LOG_ERROR("[RedisNcsConnector] Unexpected reply type for key {}", redisKey); + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + LOG_ERROR("[RedisNcsConnector] Unexpected reply type {} for key {}", reply->type, redisKey); } freeReplyObject(reply); @@ -85,8 +141,7 @@ std::vector RedisNcsConnector::multiPut( std::vector results(keys.size(), NcsStatus::ERROR); - if (!context_) { - LOG_ERROR("[RedisNcsConnector] Redis context is null"); + if (keys.empty()) { return results; } @@ -95,20 +150,45 @@ std::vector RedisNcsConnector::multiPut( return results; } + if (!ensureConnected()) { + LOG_ERROR("[RedisNcsConnector] Not connected to Redis"); + return results; + } + for (size_t i = 0; i < keys.size(); ++i) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); - redisReply* reply = (redisReply*)redisCommand(context_, "SET %s %b", - redisKey.c_str(), buffs[i].data(), buffs[i].size()); + if (redisAppendCommand(ctx_, "SET %s %b", + redisKey.c_str(), buffs[i].data(), buffs[i].size()) != REDIS_OK) { + LOG_ERROR("[RedisNcsConnector] Failed to append SET command for key {}", redisKey); + for (size_t j = 0; j < i; ++j) { + redisReply* reply = nullptr; + redisGetReply(ctx_, (void**)&reply); + if (reply) freeReplyObject(reply); + } + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + redisReply* reply = nullptr; + + if (redisGetReply(ctx_, (void**)&reply) != REDIS_OK) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + LOG_ERROR("[RedisNcsConnector] Failed to get reply for key {}: {}", + redisKey, ctx_->errstr ? ctx_->errstr : "unknown error"); + continue; + } if (reply == nullptr) { - LOG_ERROR("[RedisNcsConnector] Failed to SET key {}: {}", redisKey, context_->errstr); continue; } - if (reply->type == REDIS_REPLY_STATUS && std::string(reply->str) == "OK") { + if (reply->type == REDIS_REPLY_STATUS && + reply->str && std::string(reply->str) == "OK") { results[i] = NcsStatus::OK; } else { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to SET key {}", redisKey); } @@ -121,24 +201,47 @@ std::vector RedisNcsConnector::multiPut( std::vector RedisNcsConnector::multiDelete(const std::vector& keys) { std::vector results(keys.size(), NcsStatus::ERROR); - if (!context_) { - LOG_ERROR("[RedisNcsConnector] Redis context is null"); + if (keys.empty()) { + return results; + } + + if (!ensureConnected()) { + LOG_ERROR("[RedisNcsConnector] Not connected to Redis"); return results; } for (size_t i = 0; i < keys.size(); ++i) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); - redisReply* reply = (redisReply*)redisCommand(context_, "DEL %s", redisKey.c_str()); + if (redisAppendCommand(ctx_, "DEL %s", redisKey.c_str()) != REDIS_OK) { + LOG_ERROR("[RedisNcsConnector] Failed to append DEL command for key {}", redisKey); + for (size_t j = 0; j < i; ++j) { + redisReply* reply = nullptr; + redisGetReply(ctx_, (void**)&reply); + if (reply) freeReplyObject(reply); + } + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + redisReply* reply = nullptr; + + if (redisGetReply(ctx_, (void**)&reply) != REDIS_OK) { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + LOG_ERROR("[RedisNcsConnector] Failed to get reply for key {}: {}", + redisKey, ctx_->errstr ? ctx_->errstr : "unknown error"); + continue; + } if (reply == nullptr) { - LOG_ERROR("[RedisNcsConnector] Failed to DEL key {}: {}", redisKey, context_->errstr); continue; } if (reply->type == REDIS_REPLY_INTEGER) { results[i] = NcsStatus::OK; } else { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to DEL key {}", redisKey); } @@ -164,6 +267,30 @@ NcsConnector* RedisNcsConnectorCreator::factoryMethod(const NcsDescriptor* descr std::string host = extras["redis_host"].get(); int port = extras["redis_port"].get(); + // Check if bucket exists by querying Redis directly + std::string bucketKey = "bucket_" + std::to_string(descriptor->getbucketId()) + "_valid"; + redisContext* ctx = redisConnect(host.c_str(), port); + if (ctx == nullptr || ctx->err) { + if (ctx) { + LOG_ERROR("[RedisNcsConnectorCreator] Failed to connect to Redis: {}", ctx->errstr); + redisFree(ctx); + } + return nullptr; + } + + redisReply* reply = (redisReply*)redisCommand(ctx, "EXISTS %s", bucketKey.c_str()); + bool bucketExists = (reply != nullptr && reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); + if (reply) freeReplyObject(reply); + redisFree(ctx); + + if (!bucketExists) { + LOG_ERROR("[RedisNcsConnectorCreator] Bucket {} does not exist", descriptor->getbucketId()); + return nullptr; + } + + LOG_INFO("[RedisNcsConnectorCreator] Creating connector for bucket {} with host={}:{}", + descriptor->getbucketId(), host, port); + return new RedisNcsConnector(descriptor->getbucketId(), host, port); } diff --git a/test/test_ncs/CMakeLists.txt b/test/test_ncs/CMakeLists.txt index 1b47f26..8858147 100644 --- a/test/test_ncs/CMakeLists.txt +++ b/test/test_ncs/CMakeLists.txt @@ -4,31 +4,22 @@ find_package(GTest CONFIG REQUIRED) find_package(OpenMP REQUIRED) find_package(BLAS REQUIRED) -set(IN_MEM_NCS_TEST_FILES +# Unified NCS test for all implementations (InMemory, Redis) +set(NCS_ALL_TEST_FILES ../init_gtest.cpp - test_inmem_ncs.cpp + test_ncs_all.cpp ) -add_executable(inmem_ncs_test - ${IN_MEM_NCS_TEST_FILES} -) - -set(REDIS_NCS_TEST_FILES - ../init_gtest.cpp - test_redis_ncs.cpp -) - -add_executable(redis_ncs_test - ${REDIS_NCS_TEST_FILES} +add_executable(ncs_all_test + ${NCS_ALL_TEST_FILES} ) # Ensure SSE4.2 CRC intrinsics are available to match Folly build (x86_64) if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64" AND NOT MSVC) - target_compile_options(inmem_ncs_test PRIVATE -msse4.2) - target_compile_options(redis_ncs_test PRIVATE -msse4.2) + target_compile_options(ncs_all_test PRIVATE -msse4.2) endif() -target_link_libraries(inmem_ncs_test +target_link_libraries(ncs_all_test GTest::gtest GTest::gmock milvus-common @@ -39,17 +30,4 @@ target_link_libraries(inmem_ncs_test openblas ) -target_link_libraries(redis_ncs_test - GTest::gtest - GTest::gmock - milvus-common - ${COMMON_LINKER_LIBS} - pthread - atomic - gomp - openblas -) - - -install(TARGETS inmem_ncs_test redis_ncs_test DESTINATION unittest) - +install(TARGETS ncs_all_test DESTINATION unittest) diff --git a/test/test_ncs/test_inmem_ncs.cpp b/test/test_ncs/test_inmem_ncs.cpp deleted file mode 100644 index aa2a435..0000000 --- a/test/test_ncs/test_inmem_ncs.cpp +++ /dev/null @@ -1,93 +0,0 @@ -#include -#include "ncs/InMemNcsConnector.h" -#include "ncs/InMemoryNcs.h" // provides InMemoryNcsFactory -#include - -namespace milvus { -namespace { - -TEST(InMemNcsConnectorTest, BasicOperations) { - const uint64_t bucketId = 1; - // Register the trivial InMemoryNcs factory in the singleton and get Ncs instance - NcsSingleton::initNcs(InMemoryNcsFactory::KIND); - Ncs* ncs = NcsSingleton::Instance(); - auto createResult = ncs->createBucket(bucketId); - EXPECT_EQ(createResult, NcsStatus::OK); - - // Create descriptor and connector - auto descriptor = std::make_unique(NcsDescriptor("in_memory", bucketId, json::object())); - auto connector = std::unique_ptr( - NcsConnectorFactory::Instance().createConnector(descriptor.get())); - - ASSERT_NE(connector, nullptr); - - // Prepare test data with varying sizes - std::vector keys = {1, 2, 3}; - std::vector> values = { - std::vector(100, 0x11), // 100 bytes - std::vector(200, 0x22), // 200 bytes - std::vector(300, 0x33) // 300 bytes - }; - std::vector valueSpans; - for (auto& value : values) { - valueSpans.emplace_back(value.data(), value.size()); - } - - // Test multiPut - auto putResults = connector->multiPut(keys, valueSpans); - ASSERT_EQ(putResults.size(), keys.size()); - for (const auto& status : putResults) { - EXPECT_EQ(status, NcsStatus::OK); - } - - // Prepare buffers for reading - std::vector> readBuffers; - std::vector readSpans; - for (const auto& value : values) { - readBuffers.emplace_back(value.size()); - readSpans.emplace_back(readBuffers.back().data(), readBuffers.back().size()); - } - - // Test multiGet - auto getResults = connector->multiGet(keys, readSpans); - ASSERT_EQ(getResults.size(), keys.size()); - for (size_t i = 0; i < keys.size(); ++i) { - EXPECT_EQ(getResults[i], NcsStatus::OK); - EXPECT_EQ(readBuffers[i], values[i]); - } - - // Test with buffer too small - std::vector smallBuffer(50); - SpanBytes smallSpan(smallBuffer.data(), smallBuffer.size()); - auto getResult = connector->multiGet({keys[2]}, {smallSpan}); - ASSERT_EQ(getResult.size(), 1); - EXPECT_EQ(getResult[0], NcsStatus::ERROR); // Should fail as buffer is too small - - // Test multiDelete for specific keys - auto deleteResult = connector->multiDelete(keys); - ASSERT_EQ(deleteResult.size(), keys.size()); - for (const auto& status : deleteResult) { - EXPECT_EQ(status, NcsStatus::OK); - } - - // Verify data is deleted - auto getResultsAfterDelete = connector->multiGet(keys, readSpans); - for (const auto& status : getResultsAfterDelete) { - EXPECT_EQ(status, NcsStatus::ERROR); - } - - // Test bucket deletion - auto bucketDeleteResult = ncs->deleteBucket(bucketId); - EXPECT_EQ(bucketDeleteResult, NcsStatus::OK); - - // Negative test: try to put into a bucket that was not created - const uint32_t missingBucket = 999; - auto descriptor2 = std::make_unique("in_memory", missingBucket, json::object()); - auto connector2 = std::unique_ptr( - NcsConnectorFactory::Instance().createConnector(descriptor2.get())); - // Connector creation should fail (return nullptr) because bucket wasn't created - EXPECT_EQ(connector2, nullptr); -} - -} // namespace -} // namespace milvus \ No newline at end of file diff --git a/test/test_ncs/test_ncs_all.cpp b/test/test_ncs/test_ncs_all.cpp new file mode 100644 index 0000000..cc69a95 --- /dev/null +++ b/test/test_ncs/test_ncs_all.cpp @@ -0,0 +1,513 @@ +/** + * @file test_ncs_all.cpp + * @brief Unified parameterized tests for all NCS (Near Compute Storage) implementations. + * + * All tests run on all NCS types (InMemory, Redis) using Google Test's + * parameterized test framework. Only initialization differs per NCS type. + */ + +#include +#include "ncs/InMemNcsConnector.h" +#include "ncs/InMemoryNcs.h" +#include "ncs/RedisNcs.h" +#include "ncs/RedisNcsConnector.h" +#include +#include +#include +#include +#include + +namespace milvus { + +// Anonymous namespace for internal helpers +namespace { + +// ============================================================================ +// NCS Type Enumeration and Configuration +// ============================================================================ + +enum class NcsType { + InMemory, + Redis +}; + +std::string NcsTypeToString(NcsType type) { + switch (type) { + case NcsType::InMemory: return "InMemory"; + case NcsType::Redis: return "Redis"; + } + return "Unknown"; +} + +/** + * @brief Configuration for each NCS type. + */ +struct NcsTestConfig { + NcsType type; + std::string kind; + json config; + + static NcsTestConfig InMemory() { + return {NcsType::InMemory, "in_memory", json::object()}; + } + + static NcsTestConfig Redis() { + json config; + config["redis_host"] = "localhost"; + config["redis_port"] = 6379; + return {NcsType::Redis, "redis", config}; + } +}; + +} // namespace (anonymous) + +// ============================================================================ +// Parameterized Test Fixture +// ============================================================================ + +class NcsTest : public ::testing::TestWithParam { +protected: + static constexpr size_t kValueSize = 4096; // 4KB per value + + void SetUp() override { + config_ = GetParam(); + + // Initialize NCS based on type + try { + switch (config_.type) { + case NcsType::InMemory: + NcsSingleton::initNcs(InMemoryNcsFactory::KIND); + break; + case NcsType::Redis: + NcsSingleton::initNcs(RedisNcsFactory::KIND, config_.config); + break; + } + } catch (const std::exception& e) { + GTEST_SKIP() << NcsTypeToString(config_.type) << " not available: " << e.what(); + } + + ncs_ = NcsSingleton::Instance(); + if (!ncs_) { + GTEST_SKIP() << NcsTypeToString(config_.type) << " NCS not initialized"; + } + + auto createResult = ncs_->createBucket(bucketId_); + if (createResult != NcsStatus::OK) { + GTEST_SKIP() << "Failed to create " << NcsTypeToString(config_.type) << " bucket"; + } + + auto descriptor = std::make_unique(config_.kind, bucketId_, config_.config); + try { + connector_.reset(NcsConnectorFactory::Instance().createConnector(descriptor.get())); + } catch (const std::exception& e) { + GTEST_SKIP() << "Failed to create " << NcsTypeToString(config_.type) << " connector: " << e.what(); + } + + if (!connector_) { + GTEST_SKIP() << NcsTypeToString(config_.type) << " connector is null"; + } + } + + void TearDown() override { + connector_.reset(); // Release connector before deleting bucket + if (ncs_) { + ncs_->deleteBucket(bucketId_); + } + ncs_ = nullptr; + NcsSingleton::reset(); // Reset singleton for next test + } + + std::vector generateTestData(size_t size, uint8_t pattern) { + return std::vector(size, pattern); + } + + /** + * @brief Create a connector with custom config (for high-concurrency tests). + */ + std::unique_ptr createConnectorWithConfig(uint64_t bucketId, const json& customConfig) { + ncs_->createBucket(bucketId); + auto descriptor = std::make_unique(config_.kind, bucketId, customConfig); + return std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor.get())); + } + + /** + * @brief Create a new connector to the test bucket. + * Each connector has its own connection to the backend. + */ + std::unique_ptr createConnector() { + auto descriptor = std::make_unique(config_.kind, bucketId_, config_.config); + return std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor.get())); + } + + /** + * @brief Run concurrent test using multiple connectors (one per thread). + * + * This simulates the thread_local connector model where each thread + * creates and uses its own connector instance. + */ + void runConcurrentTestWithMultipleConnectors( + size_t numThreads, + size_t opsPerThread, + size_t numKeys + ) { + std::atomic successfulOps{0}; + std::atomic failedOps{0}; + std::vector threads; + + // Pre-populate data using the main connector + std::vector allKeys; + std::vector> allValues; + std::vector allSpans; + + for (size_t i = 0; i < numKeys; ++i) { + allKeys.push_back(static_cast(i)); + allValues.push_back(generateTestData(kValueSize, static_cast(i % 256))); + } + for (auto& v : allValues) { + allSpans.emplace_back(v.data(), v.size()); + } + + auto putResults = connector_->multiPut(allKeys, allSpans); + for (const auto& r : putResults) { + ASSERT_EQ(r, NcsStatus::OK) << "Failed to pre-populate data"; + } + + for (size_t t = 0; t < numThreads; ++t) { + threads.emplace_back([this, &successfulOps, &failedOps, numKeys, opsPerThread, t]() { + // Each thread creates its own connector (thread_local model) + std::unique_ptr threadConnector; + try { + threadConnector = createConnector(); + } catch (const std::exception& e) { + failedOps.fetch_add(opsPerThread); + return; + } + + if (!threadConnector) { + failedOps.fetch_add(opsPerThread); + return; + } + + std::mt19937 rng(static_cast(t)); + std::uniform_int_distribution keyDist(0, numKeys - 1); + std::uniform_int_distribution batchDist(1, 10); + + for (size_t op = 0; op < opsPerThread; ++op) { + size_t batchSize = batchDist(rng); + std::vector keys; + std::vector> buffers; + std::vector spans; + + for (size_t b = 0; b < batchSize; ++b) { + size_t keyIdx = keyDist(rng); + keys.push_back(static_cast(keyIdx)); + buffers.emplace_back(kValueSize); + } + for (auto& buf : buffers) { + spans.emplace_back(buf.data(), buf.size()); + } + + auto results = threadConnector->multiGet(keys, spans); + + bool allOk = true; + for (size_t i = 0; i < results.size(); ++i) { + if (results[i] == NcsStatus::OK) { + uint8_t expectedPattern = static_cast(keys[i] % 256); + if (buffers[i][0] != expectedPattern) { + allOk = false; + } + } else { + allOk = false; + } + } + + if (allOk) { + successfulOps.fetch_add(batchSize); + } else { + failedOps.fetch_add(batchSize); + } + } + }); + } + + for (auto& t : threads) { + t.join(); + } + + EXPECT_EQ(failedOps.load(), 0) << "Some operations failed during concurrent access"; + + connector_->multiDelete(allKeys); + } + + static constexpr uint64_t bucketId_ = 1000; + NcsTestConfig config_; + Ncs* ncs_ = nullptr; + std::unique_ptr connector_; +}; + +// ============================================================================ +// Test Cases (run on all NCS types) +// ============================================================================ + +TEST_P(NcsTest, BasicOperations) { + std::vector keys = {1, 2, 3}; + std::vector> values = { + generateTestData(100, 0x11), + generateTestData(200, 0x22), + generateTestData(300, 0x33) + }; + std::vector putSpans; + for (auto& v : values) { + putSpans.emplace_back(v.data(), v.size()); + } + + // Put + auto putResults = connector_->multiPut(keys, putSpans); + ASSERT_EQ(putResults.size(), keys.size()); + for (const auto& r : putResults) { + EXPECT_EQ(r, NcsStatus::OK); + } + + // Get + std::vector> readBuffers; + std::vector readSpans; + for (const auto& v : values) { + readBuffers.emplace_back(v.size()); + readSpans.emplace_back(readBuffers.back().data(), readBuffers.back().size()); + } + + auto getResults = connector_->multiGet(keys, readSpans); + ASSERT_EQ(getResults.size(), keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + EXPECT_EQ(getResults[i], NcsStatus::OK); + EXPECT_EQ(readBuffers[i], values[i]); + } + + // Delete + auto deleteResults = connector_->multiDelete(keys); + for (const auto& r : deleteResults) { + EXPECT_EQ(r, NcsStatus::OK); + } + + // Verify deleted + auto verifyResults = connector_->multiGet(keys, readSpans); + for (const auto& r : verifyResults) { + EXPECT_EQ(r, NcsStatus::ERROR); + } +} + +TEST_P(NcsTest, BufferTooSmall) { + std::vector keys = {100}; + std::vector value(300, 0xAA); + std::vector putSpans = {SpanBytes(value.data(), value.size())}; + + connector_->multiPut(keys, putSpans); + + // Try to read with buffer too small + std::vector smallBuffer(50); + std::vector smallSpans = {SpanBytes(smallBuffer.data(), smallBuffer.size())}; + + auto result = connector_->multiGet(keys, smallSpans); + EXPECT_EQ(result[0], NcsStatus::ERROR); + + connector_->multiDelete(keys); +} + +TEST_P(NcsTest, Overwrite) { + std::vector keys = {10}; + std::vector value1(50, 0xAA); + std::vector putSpans1 = {SpanBytes(value1.data(), value1.size())}; + + connector_->multiPut(keys, putSpans1); + + // Overwrite with new value + std::vector value2(50, 0xBB); + std::vector putSpans2 = {SpanBytes(value2.data(), value2.size())}; + + connector_->multiPut(keys, putSpans2); + + // Verify new value + std::vector readBuffer(50); + std::vector readSpans = {SpanBytes(readBuffer.data(), readBuffer.size())}; + + auto result = connector_->multiGet(keys, readSpans); + EXPECT_EQ(result[0], NcsStatus::OK); + EXPECT_EQ(readBuffer, value2); + + connector_->multiDelete(keys); +} + +TEST_P(NcsTest, LargeBucketId) { + const uint64_t largeBucketId = 463281186922614943ULL; + + auto createResult = ncs_->createBucket(largeBucketId); + EXPECT_EQ(createResult, NcsStatus::OK); + + bool exists = ncs_->isBucketExist(largeBucketId); + EXPECT_TRUE(exists); + + auto deleteResult = ncs_->deleteBucket(largeBucketId); + EXPECT_EQ(deleteResult, NcsStatus::OK); +} + +TEST_P(NcsTest, MissingBucket) { + // Negative test: try to create a connector for a bucket that doesn't exist + const uint64_t missingBucket = 999999; + + // Ensure bucket doesn't exist + ncs_->deleteBucket(missingBucket); + EXPECT_FALSE(ncs_->isBucketExist(missingBucket)); + + auto descriptor = std::make_unique(config_.kind, missingBucket, config_.config); + auto connector = std::unique_ptr( + NcsConnectorFactory::Instance().createConnector(descriptor.get())); + + // Connector creation should fail (return nullptr) because bucket wasn't created + EXPECT_EQ(connector, nullptr); +} + +TEST_P(NcsTest, EmptyBatch) { + std::vector emptyKeys; + std::vector emptySpans; + + auto getResults = connector_->multiGet(emptyKeys, emptySpans); + EXPECT_TRUE(getResults.empty()); + + auto putResults = connector_->multiPut(emptyKeys, emptySpans); + EXPECT_TRUE(putResults.empty()); + + auto deleteResults = connector_->multiDelete(emptyKeys); + EXPECT_TRUE(deleteResults.empty()); +} + +TEST_P(NcsTest, LargeBatch) { + const size_t batchSize = 100; + + std::vector keys; + std::vector> values; + std::vector putSpans; + + for (size_t i = 0; i < batchSize; ++i) { + keys.push_back(static_cast(5000 + i)); + values.push_back(generateTestData(kValueSize, static_cast(i))); + } + for (auto& v : values) { + putSpans.emplace_back(v.data(), v.size()); + } + + // Put batch + auto putResults = connector_->multiPut(keys, putSpans); + ASSERT_EQ(putResults.size(), batchSize); + for (const auto& r : putResults) { + EXPECT_EQ(r, NcsStatus::OK); + } + + // Get batch + std::vector> readBuffers(batchSize, std::vector(kValueSize)); + std::vector readSpans; + for (auto& buf : readBuffers) { + readSpans.emplace_back(buf.data(), buf.size()); + } + + auto getResults = connector_->multiGet(keys, readSpans); + ASSERT_EQ(getResults.size(), batchSize); + for (size_t i = 0; i < batchSize; ++i) { + EXPECT_EQ(getResults[i], NcsStatus::OK); + EXPECT_EQ(readBuffers[i], values[i]); + } + + // Delete batch + auto deleteResults = connector_->multiDelete(keys); + ASSERT_EQ(deleteResults.size(), batchSize); + for (const auto& r : deleteResults) { + EXPECT_EQ(r, NcsStatus::OK); + } +} + +TEST_P(NcsTest, ConcurrentAccess) { + // Test using multiple connectors (one per thread) - simulates thread_local model + runConcurrentTestWithMultipleConnectors( + /*numThreads=*/16, + /*opsPerThread=*/100, + /*numKeys=*/1000); +} + +TEST_P(NcsTest, HighConcurrency) { + // High concurrency test using multiple connectors (one per thread) + // No need for special config since each thread has its own connector + runConcurrentTestWithMultipleConnectors( + /*numThreads=*/80, + /*opsPerThread=*/50, + /*numKeys=*/2000); +} + +TEST_P(NcsTest, MultipleBatches) { + // Test multiple sequential batch operations + const size_t batchSize = 100; + const size_t numBatches = 10; + + std::vector keys; + std::vector> values; + std::vector putSpans; + + for (size_t i = 0; i < batchSize; ++i) { + keys.push_back(static_cast(10000 + i)); + values.push_back(generateTestData(kValueSize, static_cast(i))); + } + for (auto& v : values) { + putSpans.emplace_back(v.data(), v.size()); + } + + for (size_t batch = 0; batch < numBatches; ++batch) { + auto putResults = connector_->multiPut(keys, putSpans); + ASSERT_EQ(putResults.size(), batchSize); + for (const auto& r : putResults) { + ASSERT_EQ(r, NcsStatus::OK); + } + + std::vector> readBuffers(batchSize, std::vector(kValueSize)); + std::vector readSpans; + for (auto& buf : readBuffers) { + readSpans.emplace_back(buf.data(), buf.size()); + } + + auto getResults = connector_->multiGet(keys, readSpans); + ASSERT_EQ(getResults.size(), batchSize); + for (size_t i = 0; i < batchSize; ++i) { + EXPECT_EQ(getResults[i], NcsStatus::OK); + EXPECT_EQ(readBuffers[i], values[i]); + } + } + + connector_->multiDelete(keys); +} + +// ============================================================================ +// Test Instantiation +// ============================================================================ + +namespace { +// Custom name generator for better test output +std::string NcsTestNameGenerator(const ::testing::TestParamInfo& info) { + return NcsTypeToString(info.param.type); +} + +// Build the list of NCS types to test +std::vector GetNcsTestConfigs() { + std::vector configs; + configs.push_back(NcsTestConfig::InMemory()); + configs.push_back(NcsTestConfig::Redis()); + return configs; +} +} // namespace (anonymous) + +INSTANTIATE_TEST_SUITE_P( + AllNcsTypes, + NcsTest, + ::testing::ValuesIn(GetNcsTestConfigs()), + NcsTestNameGenerator +); + +} // namespace milvus diff --git a/test/test_ncs/test_redis_ncs.cpp b/test/test_ncs/test_redis_ncs.cpp deleted file mode 100644 index de6dca1..0000000 --- a/test/test_ncs/test_redis_ncs.cpp +++ /dev/null @@ -1,140 +0,0 @@ -#include -#include "ncs/RedisNcs.h" -#include "ncs/RedisNcsConnector.h" -#include -#include - -namespace milvus { -namespace { - -// Note: These tests require a Redis server running on localhost:6379 -// To skip these tests if Redis is not available, they can be marked as DISABLED_ - -TEST(RedisNcsTest, BasicBucketOperations) { - // Initialize RedisNcs with host and port - json config; - config["redis_host"] = "localhost"; - config["redis_port"] = 6379; - - NcsSingleton::initNcs(RedisNcsFactory::KIND, config); - Ncs* ncs = NcsSingleton::Instance(); - - const uint64_t bucketId = 100; - - // Test bucket creation - auto createResult = ncs->createBucket(bucketId); - EXPECT_EQ(createResult, NcsStatus::OK); - - // Test bucket existence check - bool exists = ncs->isBucketExist(bucketId); - EXPECT_TRUE(exists); - - // Test bucket deletion - auto deleteResult = ncs->deleteBucket(bucketId); - EXPECT_EQ(deleteResult, NcsStatus::OK); - - // Verify bucket no longer exists - exists = ncs->isBucketExist(bucketId); - EXPECT_FALSE(exists); -} - -TEST(RedisNcsConnectorTest, MultiGetPutDelete) { - const uint64_t bucketId = 101; - - // Initialize RedisNcs - json config; - config["redis_host"] = "localhost"; - config["redis_port"] = 6379; - - NcsSingleton::initNcs(RedisNcsFactory::KIND, config); - Ncs* ncs = NcsSingleton::Instance(); - - // Create bucket - auto createResult = ncs->createBucket(bucketId); - ASSERT_EQ(createResult, NcsStatus::OK); - - auto descriptor = std::make_unique("redis", bucketId, config); - auto connector = std::unique_ptr( - NcsConnectorFactory::Instance().createConnector(descriptor.get())); - - ASSERT_NE(connector, nullptr); - - // Prepare test data - std::vector keys = {1, 2, 3}; - std::vector> values = { - std::vector(100, 0x11), // 100 bytes - std::vector(200, 0x22), // 200 bytes - std::vector(300, 0x33) // 300 bytes - }; - - // Create SpanBytes for put operation - std::vector putBuffs; - for (const auto& value : values) { - putBuffs.emplace_back(const_cast(value.data()), value.size()); - } - - // Test multiPut - auto putResults = connector->multiPut(keys, putBuffs); - EXPECT_EQ(putResults.size(), keys.size()); - for (const auto& result : putResults) { - EXPECT_EQ(result, NcsStatus::OK); - } - - // Test multiGet - std::vector> getBuffers(keys.size()); - std::vector getBuffs; - for (size_t i = 0; i < keys.size(); ++i) { - getBuffers[i].resize(values[i].size()); - getBuffs.emplace_back(getBuffers[i].data(), getBuffers[i].size()); - } - - auto getResults = connector->multiGet(keys, getBuffs); - EXPECT_EQ(getResults.size(), keys.size()); - for (size_t i = 0; i < getResults.size(); ++i) { - EXPECT_EQ(getResults[i], NcsStatus::OK); - EXPECT_EQ(getBuffers[i], values[i]); - } - - // Test multiDelete - auto deleteResults = connector->multiDelete(keys); - EXPECT_EQ(deleteResults.size(), keys.size()); - for (const auto& result : deleteResults) { - EXPECT_EQ(result, NcsStatus::OK); - } - - // Verify deletion - get should return empty/error - auto verifyResults = connector->multiGet(keys, getBuffs); - for (const auto& result : verifyResults) { - EXPECT_NE(result, NcsStatus::OK); - } - - // Cleanup - ncs->deleteBucket(bucketId); -} - -TEST(RedisNcsTest, LargeBucketId) { - // Initialize RedisNcs with host and port - json config; - config["redis_host"] = "localhost"; - config["redis_port"] = 6379; - - NcsSingleton::initNcs(RedisNcsFactory::KIND, config); - Ncs* ncs = NcsSingleton::Instance(); - - const uint64_t bucketId = 463281186922614943ULL; - - // Test bucket creation - auto createResult = ncs->createBucket(bucketId); - EXPECT_EQ(createResult, NcsStatus::OK); - - // Test bucket existence check - bool exists = ncs->isBucketExist(bucketId); - EXPECT_TRUE(exists); - - // Test bucket deletion - auto deleteResult = ncs->deleteBucket(bucketId); - EXPECT_EQ(deleteResult, NcsStatus::OK); -} - -} // namespace -} // namespace milvus From 55e16501535fb1eda7e1bbfe39fb30e0ba45cad4 Mon Sep 17 00:00:00 2001 From: Ron Marcus Date: Tue, 13 Jan 2026 15:42:32 +0200 Subject: [PATCH 3/3] - Replace custom SpanBytes with boost::span - Use std::unique_ptr for Redis resources (redisContext/redisReply) - Add USE_REDIS CMake option to make Redis support optional - Delete copy/move operations for NcsConnector - Initialize all member variables with default values - Reorganize error codes (moved NcsUploadError to 2044) - Remove using clauses from header files --- CMakeLists.txt | 26 ++++++- include/common/EasyAssert.h | 18 ++--- include/common/SpanBytes.h | 19 ----- include/ncs/InMemNcsConnector.h | 6 +- include/ncs/InMemoryKV.h | 6 +- include/ncs/InMemoryNcs.h | 2 +- include/ncs/RedisNcs.h | 11 ++- include/ncs/RedisNcsConnector.h | 16 +++-- include/ncs/RedisTypes.h | 17 +++++ include/ncs/ncs.h | 34 ++++----- src/ncs/InMemNcsConnector.cpp | 6 +- src/ncs/InMemoryKV.cpp | 4 +- src/ncs/InMemoryNcs.cpp | 2 +- src/ncs/RedisNcs.cpp | 57 ++++++++------- src/ncs/RedisNcsConnector.cpp | 120 +++++++++++++------------------- src/ncs/ncs.cpp | 6 +- test/test_ncs/test_ncs_all.cpp | 52 +++++++++----- 17 files changed, 209 insertions(+), 193 deletions(-) delete mode 100644 include/common/SpanBytes.h create mode 100644 include/ncs/RedisTypes.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 5fcb8ed..622b37a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,6 +14,7 @@ project(milvus_common CXX C) option(ENABLE_UNIT_TESTS "Enable unit tests" OFF) option(ENABLE_SYNCPOINT "Enable sync point for testing" OFF) +option(USE_REDIS "Enable Redis NCS support" ON) set( CMAKE_CXX_STANDARD 17 ) set( CMAKE_CXX_STANDARD_REQUIRED on ) @@ -41,14 +42,20 @@ find_package(gflags REQUIRED) find_package(glog REQUIRED) find_package(fmt REQUIRED) find_package(prometheus-cpp REQUIRED) -find_package(hiredis REQUIRED) + +if(USE_REDIS) + find_package(hiredis REQUIRED) +endif() list(APPEND COMMON_LINKER_LIBS glog::glog) list(APPEND COMMON_LINKER_LIBS prometheus-cpp::core prometheus-cpp::push) list(APPEND COMMON_LINKER_LIBS fmt::fmt-header-only) list(APPEND COMMON_LINKER_LIBS Folly::folly) list(APPEND COMMON_LINKER_LIBS gflags::gflags) -list(APPEND COMMON_LINKER_LIBS hiredis::hiredis) + +if(USE_REDIS) + list(APPEND COMMON_LINKER_LIBS hiredis::hiredis) +endif() list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_trace) list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_exporter_ostream_span) @@ -74,12 +81,27 @@ else() list(APPEND COMMON_LINKER_LIBS ${LIBAIO_LIBRARY}) endif() +# Conditionally exclude Redis source files if USE_REDIS is OFF +if(NOT USE_REDIS) + list(REMOVE_ITEM SRC_FILES + ${CMAKE_CURRENT_SOURCE_DIR}/src/ncs/RedisNcs.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/ncs/RedisNcsConnector.cpp + ) + message(STATUS "Redis support disabled, excluding Redis NCS files from build") +else() + message(STATUS "Redis support enabled, including Redis NCS files in build") +endif() + add_library(milvus-common SHARED ${SRC_FILES}) if (ENABLE_SYNCPOINT) add_definitions(-DENABLE_SYNCPOINT) endif() +if (USE_REDIS) + target_compile_definitions(milvus-common PUBLIC USE_REDIS) +endif() + target_link_libraries(milvus-common PUBLIC ${COMMON_LINKER_LIBS} ) diff --git a/include/common/EasyAssert.h b/include/common/EasyAssert.h index 3666cfd..6e72809 100644 --- a/include/common/EasyAssert.h +++ b/include/common/EasyAssert.h @@ -63,16 +63,16 @@ enum ErrorCode { MemAllocateFailed = 2034, MemAllocateSizeNotMatch = 2035, MmapError = 2036, - NcsUploadError = 2037, // timeout or cancel related - FollyOtherException = 2038, - FollyCancel = 2039, - OutOfRange = 2040, - GcpNativeError = 2041, - TextIndexNotFound = 2042, - InvalidParameter = 2043, - InsufficientResource = 2044, - + FollyOtherException = 2037, + FollyCancel = 2038, + OutOfRange = 2039, + GcpNativeError = 2040, + TextIndexNotFound = 2041, + InvalidParameter = 2042, + InsufficientResource = 2043, + NcsUploadError = 2044, + KnowhereError = 2099 }; diff --git a/include/common/SpanBytes.h b/include/common/SpanBytes.h deleted file mode 100644 index 05c1493..0000000 --- a/include/common/SpanBytes.h +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include - -namespace milvus { - -class SpanBytes { -public: - SpanBytes(void* data, size_t size) : data_(data), size_(size) {} - - void* data() const { return data_; } - size_t size() const { return size_; } - -private: - void* data_; - size_t size_; -}; - -} // namespace milvus \ No newline at end of file diff --git a/include/ncs/InMemNcsConnector.h b/include/ncs/InMemNcsConnector.h index fc05d3c..d932973 100644 --- a/include/ncs/InMemNcsConnector.h +++ b/include/ncs/InMemNcsConnector.h @@ -6,15 +6,13 @@ namespace milvus { -using std::make_unique; - class InMemNcsConnector : public NcsConnector { public: friend class InMemoryNcsConnectorCreator; // Interface implementations - std::vector multiGet(const std::vector& keys, const std::vector& buffs) override; - std::vector multiPut(const std::vector& keys, const std::vector& buffs) override; + std::vector multiGet(const std::vector& keys, const std::vector>& buffs) override; + std::vector multiPut(const std::vector& keys, const std::vector>& buffs) override; std::vector multiDelete(const std::vector& keys) override; private: diff --git a/include/ncs/InMemoryKV.h b/include/ncs/InMemoryKV.h index d351a14..ad214bc 100644 --- a/include/ncs/InMemoryKV.h +++ b/include/ncs/InMemoryKV.h @@ -5,7 +5,7 @@ #include #include -#include "common/SpanBytes.h" +#include namespace milvus { @@ -15,11 +15,11 @@ class InMemoryKV { // Put a value into an existing bucket. Returns true on success, false if the // bucket does not exist. - bool put(uint64_t bucketId, uint32_t key, const SpanBytes& value); + bool put(uint64_t bucketId, uint32_t key, const boost::span& value); bool createBucket(uint64_t bucketId); bool deleteBucket(uint64_t bucketId); bool hasBucket(uint64_t bucketId) const; - bool get(uint64_t bucketId, uint32_t key, const SpanBytes& buff) const; + bool get(uint64_t bucketId, uint32_t key, const boost::span& buff) const; bool deleteKey(uint64_t bucketId, uint32_t key); private: InMemoryKV() = default; diff --git a/include/ncs/InMemoryNcs.h b/include/ncs/InMemoryNcs.h index becaa2d..40fd0fc 100644 --- a/include/ncs/InMemoryNcs.h +++ b/include/ncs/InMemoryNcs.h @@ -23,7 +23,7 @@ class InMemoryNcs : public Ncs { class InMemoryNcsFactory : public NcsFactory { public: static const std::string KIND; - std::unique_ptr createNcs(const json& params = json{}) override; + std::unique_ptr createNcs(const nlohmann::json& params = nlohmann::json{}) override; const std::string& getKind() const override; }; diff --git a/include/ncs/RedisNcs.h b/include/ncs/RedisNcs.h index f866797..a5fbbf5 100644 --- a/include/ncs/RedisNcs.h +++ b/include/ncs/RedisNcs.h @@ -1,6 +1,9 @@ #pragma once +#ifdef USE_REDIS + #include "ncs/ncs.h" +#include "ncs/RedisTypes.h" #include "log/Log.h" #include #include @@ -19,9 +22,9 @@ class RedisNcs : public Ncs { private: explicit RedisNcs(const std::string& host, int port); - redisContext* context_; + ncs::RedisContextPtr context_; std::string host_; - int port_; + int port_ = 0; std::mutex mutex_; friend class RedisNcsFactory; @@ -30,8 +33,10 @@ class RedisNcs : public Ncs { class RedisNcsFactory : public NcsFactory { public: static const std::string KIND; - std::unique_ptr createNcs(const json& params = json{}) override; + std::unique_ptr createNcs(const nlohmann::json& params = nlohmann::json{}) override; const std::string& getKind() const override; }; } // namespace milvus + +#endif // USE_REDIS diff --git a/include/ncs/RedisNcsConnector.h b/include/ncs/RedisNcsConnector.h index 1402c2a..84b885f 100644 --- a/include/ncs/RedisNcsConnector.h +++ b/include/ncs/RedisNcsConnector.h @@ -1,6 +1,9 @@ #pragma once +#ifdef USE_REDIS + #include "ncs/ncs.h" +#include "ncs/RedisTypes.h" #include "log/Log.h" #include #include @@ -23,8 +26,9 @@ namespace milvus { class RedisNcsConnector : public NcsConnector { public: ~RedisNcsConnector() override; - std::vector multiGet(const std::vector& keys, const std::vector& buffs) override; - std::vector multiPut(const std::vector& keys, const std::vector& buffs) override; + + std::vector multiGet(const std::vector& keys, const std::vector>& buffs) override; + std::vector multiPut(const std::vector& keys, const std::vector>& buffs) override; std::vector multiDelete(const std::vector& keys) override; private: @@ -35,10 +39,12 @@ class RedisNcsConnector : public NcsConnector { * @return true if connection is valid, false otherwise. */ bool ensureConnected(); + + ncs::RedisReplyPtr getSafeReply(); - redisContext* ctx_ = nullptr; + ncs::RedisContextPtr ctx_{nullptr, redisFree}; std::string host_; - int port_; + int port_ = 0; friend class RedisNcsConnectorCreator; }; @@ -51,3 +57,5 @@ class RedisNcsConnectorCreator : public NcsConnectorCreator { }; } // namespace milvus + +#endif // USE_REDIS diff --git a/include/ncs/RedisTypes.h b/include/ncs/RedisTypes.h new file mode 100644 index 0000000..a1ec9e2 --- /dev/null +++ b/include/ncs/RedisTypes.h @@ -0,0 +1,17 @@ +#pragma once + +#ifdef USE_REDIS + +#include +#include + +namespace milvus { +namespace ncs { + +using RedisReplyPtr = std::unique_ptr; +using RedisContextPtr = std::unique_ptr; + +} // namespace ncs +} // namespace milvus + +#endif // USE_REDIS diff --git a/include/ncs/ncs.h b/include/ncs/ncs.h index 89b69a8..68238c9 100644 --- a/include/ncs/ncs.h +++ b/include/ncs/ncs.h @@ -64,17 +64,11 @@ #include #include "nlohmann/json.hpp" -#include "common/SpanBytes.h" +#include namespace milvus { -using std::string; -using std::vector; -using std::unique_ptr; -using std::unordered_map; -using json = nlohmann::json; - enum class NcsStatus { OK, ERROR @@ -82,18 +76,18 @@ enum class NcsStatus { class NcsDescriptor { public: - NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const json& extras); + NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const nlohmann::json& extras); virtual ~NcsDescriptor() = default; const std::string& getKind() const; uint64_t getbucketId() const; - const json& getExtras() const { return extras_; } + const nlohmann::json& getExtras() const { return extras_; } NLOHMANN_DEFINE_TYPE_INTRUSIVE(NcsDescriptor, ncsKind_, bucketId_, extras_) NcsDescriptor() = default; private: std::string ncsKind_; - uint64_t bucketId_; - json extras_; + uint64_t bucketId_ = 0; + nlohmann::json extras_; }; class NcsBucketStatus{ @@ -146,7 +140,7 @@ class Ncs{ class NcsFactory { public: - virtual std::unique_ptr createNcs(const json& params = json{}) = 0; + virtual std::unique_ptr createNcs(const nlohmann::json& params = nlohmann::json{}) = 0; virtual const std::string& getKind() const = 0; virtual ~NcsFactory() = default; }; @@ -155,7 +149,7 @@ class NcsFactoryRegistry { public: static NcsFactoryRegistry& Instance(); void registerFactory(std::unique_ptr factory); - std::unique_ptr createNcs(const std::string& kind, const json& params = json{}); + std::unique_ptr createNcs(const std::string& kind, const nlohmann::json& params = nlohmann::json{}); bool hasKind(const std::string& kind) const; ~NcsFactoryRegistry() = default; @@ -166,7 +160,7 @@ class NcsFactoryRegistry { class NcsSingleton final{ public: - static void initNcs(const std::string& kind, const json& extras = json{}); + static void initNcs(const std::string& kind, const nlohmann::json& extras = nlohmann::json{}); static Ncs* Instance(); /** @@ -185,7 +179,7 @@ class NcsSingleton final{ private: inline static std::string kind_; - inline static json extras_; + inline static nlohmann::json extras_; inline static std::unique_ptr instance_ = nullptr; }; @@ -206,6 +200,12 @@ class NcsConnector { public: virtual ~NcsConnector() = default; + // Delete copy and move operations - connectors are not copyable/movable + NcsConnector(const NcsConnector&) = delete; + NcsConnector& operator=(const NcsConnector&) = delete; + NcsConnector(NcsConnector&&) = delete; + NcsConnector& operator=(NcsConnector&&) = delete; + /** * @brief Read multiple key-value pairs. * @param keys Vector of keys to read. @@ -215,7 +215,7 @@ class NcsConnector { * - NcsStatus::OK if read succeeded * - NcsStatus::ERROR if key doesn't exist or buffer too small */ - virtual std::vector multiGet(const std::vector& keys, const std::vector& buffs) = 0; + virtual std::vector multiGet(const std::vector& keys, const std::vector>& buffs) = 0; /** * @brief Write multiple key-value pairs. @@ -227,7 +227,7 @@ class NcsConnector { * * Note: If a key already exists, its value is overwritten. */ - virtual std::vector multiPut(const std::vector& keys, const std::vector& buffs) = 0; + virtual std::vector multiPut(const std::vector& keys, const std::vector>& buffs) = 0; /** * @brief Delete multiple keys. diff --git a/src/ncs/InMemNcsConnector.cpp b/src/ncs/InMemNcsConnector.cpp index 2c0a608..3b48e75 100644 --- a/src/ncs/InMemNcsConnector.cpp +++ b/src/ncs/InMemNcsConnector.cpp @@ -5,8 +5,6 @@ namespace milvus { -using std::make_unique; - InMemNcsConnector::InMemNcsConnector(uint64_t bucketId) : NcsConnector(bucketId) { @@ -15,7 +13,7 @@ InMemNcsConnector::InMemNcsConnector(uint64_t bucketId) std::vector InMemNcsConnector::multiGet(const std::vector& keys, - const std::vector& buffs) { + const std::vector>& buffs) { std::vector results; results.reserve(keys.size()); @@ -28,7 +26,7 @@ std::vector InMemNcsConnector::multiGet(const std::vector& } std::vector InMemNcsConnector::multiPut(const std::vector& keys, - const std::vector& buffs) { + const std::vector>& buffs) { std::vector results; results.reserve(keys.size()); diff --git a/src/ncs/InMemoryKV.cpp b/src/ncs/InMemoryKV.cpp index ddd469b..94252cd 100644 --- a/src/ncs/InMemoryKV.cpp +++ b/src/ncs/InMemoryKV.cpp @@ -8,7 +8,7 @@ InMemoryKV* InMemoryKV::Instance() { return &instance; } -bool InMemoryKV::put(uint64_t bucketId, uint32_t key, const SpanBytes& value) { +bool InMemoryKV::put(uint64_t bucketId, uint32_t key, const boost::span& value) { auto bucket_it = data_.find(bucketId); if (bucket_it == data_.end()) { // bucket not created @@ -37,7 +37,7 @@ bool InMemoryKV::hasBucket(uint64_t bucketId) const { return data_.find(bucketId) != data_.end(); } -bool InMemoryKV::get(uint64_t bucketId, uint32_t key, const SpanBytes& buff) const { +bool InMemoryKV::get(uint64_t bucketId, uint32_t key, const boost::span& buff) const { auto bucket_it = data_.find(bucketId); if (bucket_it == data_.end()) return false; diff --git a/src/ncs/InMemoryNcs.cpp b/src/ncs/InMemoryNcs.cpp index 5b68627..505cf39 100644 --- a/src/ncs/InMemoryNcs.cpp +++ b/src/ncs/InMemoryNcs.cpp @@ -25,7 +25,7 @@ bool InMemoryNcs::isBucketExist(uint64_t bucketId) { // InMemoryNcsFactory implementation const std::string InMemoryNcsFactory::KIND = "in_memory"; -std::unique_ptr InMemoryNcsFactory::createNcs(const json& params) { +std::unique_ptr InMemoryNcsFactory::createNcs(const nlohmann::json& params) { return std::unique_ptr(new InMemoryNcs()); } diff --git a/src/ncs/RedisNcs.cpp b/src/ncs/RedisNcs.cpp index 788d2e8..262705c 100644 --- a/src/ncs/RedisNcs.cpp +++ b/src/ncs/RedisNcs.cpp @@ -1,3 +1,5 @@ +#ifdef USE_REDIS + #include "ncs/RedisNcs.h" #include "log/Log.h" #include @@ -9,26 +11,19 @@ namespace milvus { // RedisNcs implementation RedisNcs::RedisNcs(const std::string& host, int port) - : context_(nullptr), host_(host), port_(port) { - context_ = redisConnect(host.c_str(), port); + : context_(nullptr, redisFree), host_(host), port_(port) { + context_.reset(redisConnect(host.c_str(), port)); if (context_ == nullptr || context_->err) { - if (context_) { - LOG_ERROR("[RedisNcs] Redis connection error: {}", context_->errstr); - redisFree(context_); - context_ = nullptr; - } else { - LOG_ERROR("[RedisNcs] Redis connection error: can't allocate redis context"); - } + std::string error_msg = context_ ? context_->errstr : "can't allocate redis context"; + LOG_ERROR("[RedisNcs] Redis connection error: {}", error_msg); + context_.reset(); throw std::runtime_error("Failed to connect to Redis at " + host + ":" + std::to_string(port)); } LOG_INFO("[RedisNcs] Connected to Redis at {}:{}", host, port); } RedisNcs::~RedisNcs() { - if (context_) { - redisFree(context_); - context_ = nullptr; - } + // unique_ptr automatically calls redisFree } NcsStatus RedisNcs::createBucket(uint64_t bucketId) { @@ -39,7 +34,10 @@ NcsStatus RedisNcs::createBucket(uint64_t bucketId) { } std::string key = "bucket_" + std::to_string(bucketId) + "_valid"; - redisReply* reply = (redisReply*)redisCommand(context_, "SET %s true", key.c_str()); + ncs::RedisReplyPtr reply( + (redisReply*)redisCommand(context_.get(), "SET %s true", key.c_str()), + freeReplyObject + ); if (reply == nullptr) { LOG_ERROR("[RedisNcs] Failed to create bucket {}: {}", bucketId, context_->errstr); @@ -47,14 +45,12 @@ NcsStatus RedisNcs::createBucket(uint64_t bucketId) { } if (reply->type == REDIS_REPLY_STATUS && std::string(reply->str) == "OK") { - freeReplyObject(reply); LOG_INFO("[RedisNcs] Created bucket {}", bucketId); return NcsStatus::OK; } LOG_ERROR("[RedisNcs] Failed to create bucket {}. Reply type: {}, str: {}", bucketId, reply->type, (reply->str ? reply->str : "null")); - freeReplyObject(reply); return NcsStatus::ERROR; } @@ -71,12 +67,13 @@ NcsStatus RedisNcs::deleteBucket(uint64_t bucketId) { // Use SCAN to find all keys matching the pattern int cursor = 0; do { - redisReply* reply = (redisReply*)redisCommand(context_, - "SCAN %d MATCH %s COUNT 100", cursor, pattern.c_str()); + ncs::RedisReplyPtr reply( + (redisReply*)redisCommand(context_.get(), "SCAN %d MATCH %s COUNT 100", cursor, pattern.c_str()), + freeReplyObject + ); if (reply == nullptr || reply->type != REDIS_REPLY_ARRAY) { LOG_ERROR("[RedisNcs] Failed to scan keys for bucket {}", bucketId); - if (reply) freeReplyObject(reply); return NcsStatus::ERROR; } @@ -88,8 +85,6 @@ NcsStatus RedisNcs::deleteBucket(uint64_t bucketId) { for (size_t i = 0; i < keysArray->elements; ++i) { keysToDelete.push_back(keysArray->element[i]->str); } - - freeReplyObject(reply); } while (cursor != 0); // Delete all found keys using UNLINK (async delete) @@ -100,7 +95,10 @@ NcsStatus RedisNcs::deleteBucket(uint64_t bucketId) { cmd += " " + key; } - redisReply* reply = (redisReply*)redisCommand(context_, cmd.c_str()); + ncs::RedisReplyPtr reply( + (redisReply*)redisCommand(context_.get(), cmd.c_str()), + freeReplyObject + ); if (reply == nullptr) { LOG_ERROR("[RedisNcs] Failed to delete keys for bucket {}: {}", bucketId, context_->errstr); @@ -108,7 +106,6 @@ NcsStatus RedisNcs::deleteBucket(uint64_t bucketId) { } LOG_INFO("[RedisNcs] Deleted {} keys for bucket {}", reply->integer, bucketId); - freeReplyObject(reply); } return NcsStatus::OK; @@ -126,23 +123,23 @@ bool RedisNcs::isBucketExist(uint64_t bucketId) { } std::string key = "bucket_" + std::to_string(bucketId) + "_valid"; - redisReply* reply = (redisReply*)redisCommand(context_, "EXISTS %s", key.c_str()); + ncs::RedisReplyPtr reply( + (redisReply*)redisCommand(context_.get(), "EXISTS %s", key.c_str()), + freeReplyObject + ); if (reply == nullptr) { LOG_ERROR("[RedisNcs] Failed to check bucket existence: {}", context_->errstr); return false; } - bool exists = (reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); - freeReplyObject(reply); - - return exists; + return (reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); } // RedisNcsFactory implementation const std::string RedisNcsFactory::KIND = "redis"; -std::unique_ptr RedisNcsFactory::createNcs(const json& params) { +std::unique_ptr RedisNcsFactory::createNcs(const nlohmann::json& params) { if (!params.contains("redis_host")) { throw std::runtime_error("RedisNcsFactory: 'redis_host' is required in params"); } @@ -171,3 +168,5 @@ namespace { } } // namespace milvus + +#endif // USE_REDIS diff --git a/src/ncs/RedisNcsConnector.cpp b/src/ncs/RedisNcsConnector.cpp index 73c44f1..6bee5df 100644 --- a/src/ncs/RedisNcsConnector.cpp +++ b/src/ncs/RedisNcsConnector.cpp @@ -1,3 +1,5 @@ +#ifdef USE_REDIS + #include "ncs/RedisNcsConnector.h" #include "log/Log.h" #include @@ -13,17 +15,13 @@ namespace milvus { // ============================================================================ RedisNcsConnector::RedisNcsConnector(uint64_t bucketId, const std::string& host, int port) - : NcsConnector(bucketId), ctx_(nullptr), host_(host), port_(port) { + : NcsConnector(bucketId), ctx_(nullptr, redisFree), host_(host), port_(port) { - ctx_ = redisConnect(host.c_str(), port); + ctx_.reset(redisConnect(host.c_str(), port)); if (ctx_ == nullptr || ctx_->err) { - if (ctx_) { - LOG_ERROR("[RedisNcsConnector] Connection error: {}", ctx_->errstr); - redisFree(ctx_); - ctx_ = nullptr; - } else { - LOG_ERROR("[RedisNcsConnector] Cannot allocate redis context"); - } + std::string error_msg = ctx_ ? ctx_->errstr : "Cannot allocate redis context"; + LOG_ERROR("[RedisNcsConnector] Connection error: {}", error_msg); + ctx_.reset(); throw std::runtime_error("Failed to connect to Redis at " + host + ":" + std::to_string(port)); } @@ -32,10 +30,6 @@ RedisNcsConnector::RedisNcsConnector(uint64_t bucketId, const std::string& host, } RedisNcsConnector::~RedisNcsConnector() { - if (ctx_) { - redisFree(ctx_); - ctx_ = nullptr; - } LOG_DEBUG("[RedisNcsConnector] Destroyed connector for bucket {}", bucketId_); } @@ -44,17 +38,11 @@ bool RedisNcsConnector::ensureConnected() { return true; } - if (ctx_) { - redisFree(ctx_); - } - - ctx_ = redisConnect(host_.c_str(), port_); + ctx_.reset(redisConnect(host_.c_str(), port_)); if (ctx_ == nullptr || ctx_->err) { - if (ctx_) { - LOG_ERROR("[RedisNcsConnector] Reconnection error: {}", ctx_->errstr); - redisFree(ctx_); - ctx_ = nullptr; - } + std::string error_msg = ctx_ ? ctx_->errstr : "Cannot allocate redis context"; + LOG_ERROR("[RedisNcsConnector] Reconnection error: {}", error_msg); + ctx_.reset(); return false; } @@ -62,9 +50,17 @@ bool RedisNcsConnector::ensureConnected() { return true; } +ncs::RedisReplyPtr RedisNcsConnector::getSafeReply() { + redisReply* raw_reply = nullptr; + if (redisGetReply(ctx_.get(), (void**)&raw_reply) != REDIS_OK) { + return ncs::RedisReplyPtr(nullptr, freeReplyObject); + } + return ncs::RedisReplyPtr(raw_reply, freeReplyObject); +} + std::vector RedisNcsConnector::multiGet( const std::vector& keys, - const std::vector& buffs) { + const std::vector>& buffs) { std::vector results(keys.size(), NcsStatus::ERROR); @@ -86,31 +82,24 @@ std::vector RedisNcsConnector::multiGet( for (size_t i = 0; i < keys.size(); ++i) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); - if (redisAppendCommand(ctx_, "GET %s", redisKey.c_str()) != REDIS_OK) { + if (redisAppendCommand(ctx_.get(), "GET %s", redisKey.c_str()) != REDIS_OK) { LOG_ERROR("[RedisNcsConnector] Failed to append GET command for key {}", redisKey); - for (size_t j = 0; j < i; ++j) { - redisReply* reply = nullptr; - redisGetReply(ctx_, (void**)&reply); - if (reply) freeReplyObject(reply); - } + // Connection is in bad state, force reconnect on next call + ctx_.reset(); return results; } } for (size_t i = 0; i < keys.size(); ++i) { - redisReply* reply = nullptr; + auto reply = getSafeReply(); - if (redisGetReply(ctx_, (void**)&reply) != REDIS_OK) { + if (!reply) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to get reply for key {}: {}", redisKey, ctx_->errstr ? ctx_->errstr : "unknown error"); continue; } - if (reply == nullptr) { - continue; - } - if (reply->type == REDIS_REPLY_STRING) { size_t dataSize = reply->len; if (dataSize <= buffs[i].size()) { @@ -128,8 +117,6 @@ std::vector RedisNcsConnector::multiGet( std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Unexpected reply type {} for key {}", reply->type, redisKey); } - - freeReplyObject(reply); } return results; @@ -137,7 +124,7 @@ std::vector RedisNcsConnector::multiGet( std::vector RedisNcsConnector::multiPut( const std::vector& keys, - const std::vector& buffs) { + const std::vector>& buffs) { std::vector results(keys.size(), NcsStatus::ERROR); @@ -158,32 +145,25 @@ std::vector RedisNcsConnector::multiPut( for (size_t i = 0; i < keys.size(); ++i) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); - if (redisAppendCommand(ctx_, "SET %s %b", + if (redisAppendCommand(ctx_.get(), "SET %s %b", redisKey.c_str(), buffs[i].data(), buffs[i].size()) != REDIS_OK) { LOG_ERROR("[RedisNcsConnector] Failed to append SET command for key {}", redisKey); - for (size_t j = 0; j < i; ++j) { - redisReply* reply = nullptr; - redisGetReply(ctx_, (void**)&reply); - if (reply) freeReplyObject(reply); - } + // Connection is in bad state, force reconnect on next call + ctx_.reset(); return results; } } for (size_t i = 0; i < keys.size(); ++i) { - redisReply* reply = nullptr; + auto reply = getSafeReply(); - if (redisGetReply(ctx_, (void**)&reply) != REDIS_OK) { + if (!reply) { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to get reply for key {}: {}", redisKey, ctx_->errstr ? ctx_->errstr : "unknown error"); continue; } - if (reply == nullptr) { - continue; - } - if (reply->type == REDIS_REPLY_STATUS && reply->str && std::string(reply->str) == "OK") { results[i] = NcsStatus::OK; @@ -191,8 +171,6 @@ std::vector RedisNcsConnector::multiPut( std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to SET key {}", redisKey); } - - freeReplyObject(reply); } return results; @@ -213,39 +191,30 @@ std::vector RedisNcsConnector::multiDelete(const std::vectorerrstr ? ctx_->errstr : "unknown error"); continue; } - if (reply == nullptr) { - continue; - } - if (reply->type == REDIS_REPLY_INTEGER) { results[i] = NcsStatus::OK; } else { std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); LOG_ERROR("[RedisNcsConnector] Failed to DEL key {}", redisKey); } - - freeReplyObject(reply); } return results; @@ -255,7 +224,7 @@ std::vector RedisNcsConnector::multiDelete(const std::vectorgetExtras(); + const nlohmann::json& extras = descriptor->getExtras(); if (!extras.contains("redis_host")) { throw std::runtime_error("RedisNcsConnectorCreator: 'redis_host' is required in descriptor extras"); @@ -269,19 +238,22 @@ NcsConnector* RedisNcsConnectorCreator::factoryMethod(const NcsDescriptor* descr // Check if bucket exists by querying Redis directly std::string bucketKey = "bucket_" + std::to_string(descriptor->getbucketId()) + "_valid"; - redisContext* ctx = redisConnect(host.c_str(), port); + ncs::RedisContextPtr ctx( + redisConnect(host.c_str(), port), + redisFree + ); if (ctx == nullptr || ctx->err) { if (ctx) { LOG_ERROR("[RedisNcsConnectorCreator] Failed to connect to Redis: {}", ctx->errstr); - redisFree(ctx); } return nullptr; } - redisReply* reply = (redisReply*)redisCommand(ctx, "EXISTS %s", bucketKey.c_str()); + ncs::RedisReplyPtr reply( + (redisReply*)redisCommand(ctx.get(), "EXISTS %s", bucketKey.c_str()), + freeReplyObject + ); bool bucketExists = (reply != nullptr && reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); - if (reply) freeReplyObject(reply); - redisFree(ctx); if (!bucketExists) { LOG_ERROR("[RedisNcsConnectorCreator] Bucket {} does not exist", descriptor->getbucketId()); @@ -309,3 +281,5 @@ namespace { } } // namespace milvus + +#endif // USE_REDIS diff --git a/src/ncs/ncs.cpp b/src/ncs/ncs.cpp index 1f942e5..657cc2a 100644 --- a/src/ncs/ncs.cpp +++ b/src/ncs/ncs.cpp @@ -16,7 +16,7 @@ void NcsFactoryRegistry::registerFactory(std::unique_ptr factory) { registry_[kind] = std::move(factory); } -std::unique_ptr NcsFactoryRegistry::createNcs(const std::string& kind, const json& params) { +std::unique_ptr NcsFactoryRegistry::createNcs(const std::string& kind, const nlohmann::json& params) { LOG_DEBUG("[NCS] Creating NCS of kind: {} with params: {}", kind, params.dump()); auto it = registry_.find(kind); if (it != registry_.end()) { @@ -38,7 +38,7 @@ bool NcsFactoryRegistry::hasKind(const std::string& kind) const { } // NcsSingleton implementation -void NcsSingleton::initNcs(const std::string& kind, const json& extras) { +void NcsSingleton::initNcs(const std::string& kind, const nlohmann::json& extras) { if (!NcsFactoryRegistry::Instance().hasKind(kind)) { throw std::runtime_error("NCS Factory kind '" + kind + "' is not registered."); } @@ -91,7 +91,7 @@ NcsConnectorFactory& NcsConnectorFactory::Instance() { // Already defaulted in header // NcsDescriptor implementation -NcsDescriptor::NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const json& extras) +NcsDescriptor::NcsDescriptor(const std::string& ncsKind, uint64_t bucketId, const nlohmann::json& extras) : ncsKind_(ncsKind), bucketId_(bucketId), extras_(extras) { } diff --git a/test/test_ncs/test_ncs_all.cpp b/test/test_ncs/test_ncs_all.cpp index cc69a95..125b2d5 100644 --- a/test/test_ncs/test_ncs_all.cpp +++ b/test/test_ncs/test_ncs_all.cpp @@ -9,8 +9,10 @@ #include #include "ncs/InMemNcsConnector.h" #include "ncs/InMemoryNcs.h" +#ifdef USE_REDIS #include "ncs/RedisNcs.h" #include "ncs/RedisNcsConnector.h" +#endif #include #include #include @@ -28,13 +30,17 @@ namespace { enum class NcsType { InMemory, +#ifdef USE_REDIS Redis +#endif }; std::string NcsTypeToString(NcsType type) { switch (type) { case NcsType::InMemory: return "InMemory"; +#ifdef USE_REDIS case NcsType::Redis: return "Redis"; +#endif } return "Unknown"; } @@ -45,18 +51,20 @@ std::string NcsTypeToString(NcsType type) { struct NcsTestConfig { NcsType type; std::string kind; - json config; + nlohmann::json config; static NcsTestConfig InMemory() { - return {NcsType::InMemory, "in_memory", json::object()}; + return {NcsType::InMemory, "in_memory", nlohmann::json::object()}; } +#ifdef USE_REDIS static NcsTestConfig Redis() { - json config; + nlohmann::json config; config["redis_host"] = "localhost"; config["redis_port"] = 6379; return {NcsType::Redis, "redis", config}; } +#endif }; } // namespace (anonymous) @@ -78,9 +86,11 @@ class NcsTest : public ::testing::TestWithParam { case NcsType::InMemory: NcsSingleton::initNcs(InMemoryNcsFactory::KIND); break; +#ifdef USE_REDIS case NcsType::Redis: NcsSingleton::initNcs(RedisNcsFactory::KIND, config_.config); break; +#endif } } catch (const std::exception& e) { GTEST_SKIP() << NcsTypeToString(config_.type) << " not available: " << e.what(); @@ -124,7 +134,7 @@ class NcsTest : public ::testing::TestWithParam { /** * @brief Create a connector with custom config (for high-concurrency tests). */ - std::unique_ptr createConnectorWithConfig(uint64_t bucketId, const json& customConfig) { + std::unique_ptr createConnectorWithConfig(uint64_t bucketId, const nlohmann::json& customConfig) { ncs_->createBucket(bucketId); auto descriptor = std::make_unique(config_.kind, bucketId, customConfig); return std::unique_ptr( @@ -159,7 +169,7 @@ class NcsTest : public ::testing::TestWithParam { // Pre-populate data using the main connector std::vector allKeys; std::vector> allValues; - std::vector allSpans; + std::vector> allSpans; for (size_t i = 0; i < numKeys; ++i) { allKeys.push_back(static_cast(i)); @@ -198,7 +208,7 @@ class NcsTest : public ::testing::TestWithParam { size_t batchSize = batchDist(rng); std::vector keys; std::vector> buffers; - std::vector spans; + std::vector> spans; for (size_t b = 0; b < batchSize; ++b) { size_t keyIdx = keyDist(rng); @@ -258,7 +268,7 @@ TEST_P(NcsTest, BasicOperations) { generateTestData(200, 0x22), generateTestData(300, 0x33) }; - std::vector putSpans; + std::vector> putSpans; for (auto& v : values) { putSpans.emplace_back(v.data(), v.size()); } @@ -272,10 +282,12 @@ TEST_P(NcsTest, BasicOperations) { // Get std::vector> readBuffers; - std::vector readSpans; + std::vector> readSpans; for (const auto& v : values) { readBuffers.emplace_back(v.size()); - readSpans.emplace_back(readBuffers.back().data(), readBuffers.back().size()); + } + for (auto& buf : readBuffers) { + readSpans.emplace_back(buf.data(), buf.size()); } auto getResults = connector_->multiGet(keys, readSpans); @@ -301,13 +313,13 @@ TEST_P(NcsTest, BasicOperations) { TEST_P(NcsTest, BufferTooSmall) { std::vector keys = {100}; std::vector value(300, 0xAA); - std::vector putSpans = {SpanBytes(value.data(), value.size())}; + std::vector> putSpans = {boost::span(value.data(), value.size())}; connector_->multiPut(keys, putSpans); // Try to read with buffer too small std::vector smallBuffer(50); - std::vector smallSpans = {SpanBytes(smallBuffer.data(), smallBuffer.size())}; + std::vector> smallSpans = {boost::span(smallBuffer.data(), smallBuffer.size())}; auto result = connector_->multiGet(keys, smallSpans); EXPECT_EQ(result[0], NcsStatus::ERROR); @@ -318,19 +330,19 @@ TEST_P(NcsTest, BufferTooSmall) { TEST_P(NcsTest, Overwrite) { std::vector keys = {10}; std::vector value1(50, 0xAA); - std::vector putSpans1 = {SpanBytes(value1.data(), value1.size())}; + std::vector> putSpans1 = {boost::span(value1.data(), value1.size())}; connector_->multiPut(keys, putSpans1); // Overwrite with new value std::vector value2(50, 0xBB); - std::vector putSpans2 = {SpanBytes(value2.data(), value2.size())}; + std::vector> putSpans2 = {boost::span(value2.data(), value2.size())}; connector_->multiPut(keys, putSpans2); // Verify new value std::vector readBuffer(50); - std::vector readSpans = {SpanBytes(readBuffer.data(), readBuffer.size())}; + std::vector> readSpans = {boost::span(readBuffer.data(), readBuffer.size())}; auto result = connector_->multiGet(keys, readSpans); EXPECT_EQ(result[0], NcsStatus::OK); @@ -370,7 +382,7 @@ TEST_P(NcsTest, MissingBucket) { TEST_P(NcsTest, EmptyBatch) { std::vector emptyKeys; - std::vector emptySpans; + std::vector> emptySpans; auto getResults = connector_->multiGet(emptyKeys, emptySpans); EXPECT_TRUE(getResults.empty()); @@ -387,7 +399,7 @@ TEST_P(NcsTest, LargeBatch) { std::vector keys; std::vector> values; - std::vector putSpans; + std::vector> putSpans; for (size_t i = 0; i < batchSize; ++i) { keys.push_back(static_cast(5000 + i)); @@ -406,7 +418,7 @@ TEST_P(NcsTest, LargeBatch) { // Get batch std::vector> readBuffers(batchSize, std::vector(kValueSize)); - std::vector readSpans; + std::vector> readSpans; for (auto& buf : readBuffers) { readSpans.emplace_back(buf.data(), buf.size()); } @@ -450,7 +462,7 @@ TEST_P(NcsTest, MultipleBatches) { std::vector keys; std::vector> values; - std::vector putSpans; + std::vector> putSpans; for (size_t i = 0; i < batchSize; ++i) { keys.push_back(static_cast(10000 + i)); @@ -468,7 +480,7 @@ TEST_P(NcsTest, MultipleBatches) { } std::vector> readBuffers(batchSize, std::vector(kValueSize)); - std::vector readSpans; + std::vector> readSpans; for (auto& buf : readBuffers) { readSpans.emplace_back(buf.data(), buf.size()); } @@ -498,7 +510,9 @@ std::string NcsTestNameGenerator(const ::testing::TestParamInfo& std::vector GetNcsTestConfigs() { std::vector configs; configs.push_back(NcsTestConfig::InMemory()); +#ifdef USE_REDIS configs.push_back(NcsTestConfig::Redis()); +#endif return configs; } } // namespace (anonymous)