diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e4018b..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 ) @@ -25,6 +26,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) @@ -41,12 +43,20 @@ find_package(glog REQUIRED) find_package(fmt REQUIRED) find_package(prometheus-cpp 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) +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) list(APPEND COMMON_LINKER_LIBS opentelemetry-cpp::opentelemetry_exporter_otlp_grpc) @@ -54,6 +64,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" @@ -70,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/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..6e72809 100644 --- a/include/common/EasyAssert.h +++ b/include/common/EasyAssert.h @@ -71,7 +71,8 @@ enum ErrorCode { TextIndexNotFound = 2041, InvalidParameter = 2042, InsufficientResource = 2043, - + NcsUploadError = 2044, + KnowhereError = 2099 }; diff --git a/include/ncs/InMemNcsConnector.h b/include/ncs/InMemNcsConnector.h new file mode 100644 index 0000000..d932973 --- /dev/null +++ b/include/ncs/InMemNcsConnector.h @@ -0,0 +1,31 @@ +#pragma once + +#include "ncs/ncs.h" +#include "ncs/InMemoryKV.h" +#include + +namespace milvus { + +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..ad214bc --- /dev/null +++ b/include/ncs/InMemoryKV.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include +#include + +#include + +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 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 boost::span& 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..40fd0fc --- /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 nlohmann::json& params = nlohmann::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..a5fbbf5 --- /dev/null +++ b/include/ncs/RedisNcs.h @@ -0,0 +1,42 @@ +#pragma once + +#ifdef USE_REDIS + +#include "ncs/ncs.h" +#include "ncs/RedisTypes.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); + ncs::RedisContextPtr context_; + std::string host_; + int port_ = 0; + std::mutex mutex_; + + friend class RedisNcsFactory; +}; + +class RedisNcsFactory : public NcsFactory { +public: + static const std::string KIND; + 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 new file mode 100644 index 0000000..84b885f --- /dev/null +++ b/include/ncs/RedisNcsConnector.h @@ -0,0 +1,61 @@ +#pragma once + +#ifdef USE_REDIS + +#include "ncs/ncs.h" +#include "ncs/RedisTypes.h" +#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; + + 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); + + /** + * @brief Ensure connection is valid, reconnect if needed. + * @return true if connection is valid, false otherwise. + */ + bool ensureConnected(); + + ncs::RedisReplyPtr getSafeReply(); + + ncs::RedisContextPtr ctx_{nullptr, redisFree}; + std::string host_; + int port_ = 0; + + 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 + +#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 new file mode 100644 index 0000000..68238c9 --- /dev/null +++ b/include/ncs/ncs.h @@ -0,0 +1,295 @@ +#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 +#include +#include +#include +#include "nlohmann/json.hpp" + +#include + + +namespace milvus { + +enum class NcsStatus { + OK, + ERROR +}; + +class NcsDescriptor { +public: + 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 nlohmann::json& getExtras() const { return extras_; } + + NLOHMANN_DEFINE_TYPE_INTRUSIVE(NcsDescriptor, ncsKind_, bucketId_, extras_) + NcsDescriptor() = default; +private: + std::string ncsKind_; + uint64_t bucketId_ = 0; + nlohmann::json extras_; +}; + +class NcsBucketStatus{ + // TBD (capacity, occupancy, etc) +}; + + +/** + * @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; +}; + +class NcsFactory { +public: + virtual std::unique_ptr createNcs(const nlohmann::json& params = nlohmann::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 nlohmann::json& params = nlohmann::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 nlohmann::json& extras = nlohmann::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 nlohmann::json extras_; + inline static std::unique_ptr instance_ = nullptr; +}; + + +/** + * @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; + + // 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. + * @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: + explicit NcsConnector(uint64_t bucketId) : bucketId_(bucketId) {} + 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; +}; + + + +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..3b48e75 --- /dev/null +++ b/src/ncs/InMemNcsConnector.cpp @@ -0,0 +1,82 @@ +#include "ncs/InMemNcsConnector.h" +#include "log/Log.h" +#include +#include + +namespace milvus { + + +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..94252cd --- /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 boost::span& 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 boost::span& 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..505cf39 --- /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 nlohmann::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..262705c --- /dev/null +++ b/src/ncs/RedisNcs.cpp @@ -0,0 +1,172 @@ +#ifdef USE_REDIS + +#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, redisFree), host_(host), port_(port) { + context_.reset(redisConnect(host.c_str(), port)); + if (context_ == nullptr || context_->err) { + 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() { + // unique_ptr automatically calls redisFree +} + +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"; + 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); + return NcsStatus::ERROR; + } + + if (reply->type == REDIS_REPLY_STATUS && std::string(reply->str) == "OK") { + 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")); + 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 { + 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); + 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); + } + } 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; + } + + 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); + return NcsStatus::ERROR; + } + + LOG_INFO("[RedisNcs] Deleted {} keys for bucket {}", reply->integer, bucketId); + } + + 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"; + 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; + } + + return (reply->type == REDIS_REPLY_INTEGER && reply->integer == 1); +} + +// RedisNcsFactory implementation +const std::string RedisNcsFactory::KIND = "redis"; + +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"); + } + 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 + +#endif // USE_REDIS diff --git a/src/ncs/RedisNcsConnector.cpp b/src/ncs/RedisNcsConnector.cpp new file mode 100644 index 0000000..6bee5df --- /dev/null +++ b/src/ncs/RedisNcsConnector.cpp @@ -0,0 +1,285 @@ +#ifdef USE_REDIS + +#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), ctx_(nullptr, redisFree), host_(host), port_(port) { + + ctx_.reset(redisConnect(host.c_str(), port)); + if (ctx_ == nullptr || ctx_->err) { + 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)); + } + + LOG_DEBUG("[RedisNcsConnector] Created connector for bucket {} to {}:{}", + bucketId, host, port); +} + +RedisNcsConnector::~RedisNcsConnector() { + LOG_DEBUG("[RedisNcsConnector] Destroyed connector for bucket {}", bucketId_); +} + +bool RedisNcsConnector::ensureConnected() { + if (ctx_ != nullptr && !ctx_->err) { + return true; + } + + ctx_.reset(redisConnect(host_.c_str(), port_)); + if (ctx_ == nullptr || ctx_->err) { + std::string error_msg = ctx_ ? ctx_->errstr : "Cannot allocate redis context"; + LOG_ERROR("[RedisNcsConnector] Reconnection error: {}", error_msg); + ctx_.reset(); + return false; + } + + LOG_DEBUG("[RedisNcsConnector] Reconnected to {}:{}", host_, port_); + 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) { + + std::vector results(keys.size(), NcsStatus::ERROR); + + if (keys.empty()) { + return results; + } + + if (keys.size() != buffs.size()) { + LOG_ERROR("[RedisNcsConnector] Keys and buffers size mismatch"); + 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]); + + if (redisAppendCommand(ctx_.get(), "GET %s", redisKey.c_str()) != REDIS_OK) { + LOG_ERROR("[RedisNcsConnector] Failed to append GET command for key {}", redisKey); + // Connection is in bad state, force reconnect on next call + ctx_.reset(); + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + auto reply = getSafeReply(); + + 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->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 { + 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 { + std::string redisKey = "bucket_" + std::to_string(bucketId_) + "_" + std::to_string(keys[i]); + LOG_ERROR("[RedisNcsConnector] Unexpected reply type {} for key {}", reply->type, redisKey); + } + } + + return results; +} + +std::vector RedisNcsConnector::multiPut( + const std::vector& keys, + const std::vector>& buffs) { + + std::vector results(keys.size(), NcsStatus::ERROR); + + if (keys.empty()) { + return results; + } + + if (keys.size() != buffs.size()) { + LOG_ERROR("[RedisNcsConnector] Keys and buffers size mismatch"); + 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]); + + 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); + // Connection is in bad state, force reconnect on next call + ctx_.reset(); + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + auto reply = getSafeReply(); + + 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->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); + } + } + + return results; +} + +std::vector RedisNcsConnector::multiDelete(const std::vector& keys) { + std::vector results(keys.size(), NcsStatus::ERROR); + + 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]); + + if (redisAppendCommand(ctx_.get(), "DEL %s", redisKey.c_str()) != REDIS_OK) { + LOG_ERROR("[RedisNcsConnector] Failed to append DEL command for key {}", redisKey); + // Connection is in bad state, force reconnect on next call + ctx_.reset(); + return results; + } + } + + for (size_t i = 0; i < keys.size(); ++i) { + auto reply = getSafeReply(); + + 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->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); + } + } + + return results; +} + +// RedisNcsConnectorCreator implementation +const std::string RedisNcsConnectorCreator::KIND = "redis"; + +NcsConnector* RedisNcsConnectorCreator::factoryMethod(const NcsDescriptor* descriptor) { + const nlohmann::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(); + + // Check if bucket exists by querying Redis directly + std::string bucketKey = "bucket_" + std::to_string(descriptor->getbucketId()) + "_valid"; + 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); + } + return nullptr; + } + + 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 (!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); +} + +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 + +#endif // USE_REDIS diff --git a/src/ncs/ncs.cpp b/src/ncs/ncs.cpp new file mode 100644 index 0000000..657cc2a --- /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 nlohmann::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 nlohmann::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 nlohmann::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..8858147 --- /dev/null +++ b/test/test_ncs/CMakeLists.txt @@ -0,0 +1,33 @@ +# NCS tests + +find_package(GTest CONFIG REQUIRED) +find_package(OpenMP REQUIRED) +find_package(BLAS REQUIRED) + +# Unified NCS test for all implementations (InMemory, Redis) +set(NCS_ALL_TEST_FILES + ../init_gtest.cpp + test_ncs_all.cpp +) + +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(ncs_all_test PRIVATE -msse4.2) +endif() + +target_link_libraries(ncs_all_test + GTest::gtest + GTest::gmock + milvus-common + ${COMMON_LINKER_LIBS} + pthread + atomic + gomp + openblas +) + +install(TARGETS ncs_all_test DESTINATION unittest) diff --git a/test/test_ncs/test_ncs_all.cpp b/test/test_ncs/test_ncs_all.cpp new file mode 100644 index 0000000..125b2d5 --- /dev/null +++ b/test/test_ncs/test_ncs_all.cpp @@ -0,0 +1,527 @@ +/** + * @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" +#ifdef USE_REDIS +#include "ncs/RedisNcs.h" +#include "ncs/RedisNcsConnector.h" +#endif +#include +#include +#include +#include +#include + +namespace milvus { + +// Anonymous namespace for internal helpers +namespace { + +// ============================================================================ +// NCS Type Enumeration and Configuration +// ============================================================================ + +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"; +} + +/** + * @brief Configuration for each NCS type. + */ +struct NcsTestConfig { + NcsType type; + std::string kind; + nlohmann::json config; + + static NcsTestConfig InMemory() { + return {NcsType::InMemory, "in_memory", nlohmann::json::object()}; + } + +#ifdef USE_REDIS + static NcsTestConfig Redis() { + nlohmann::json config; + config["redis_host"] = "localhost"; + config["redis_port"] = 6379; + return {NcsType::Redis, "redis", config}; + } +#endif +}; + +} // 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; +#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(); + } + + 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 nlohmann::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()); + } + for (auto& buf : readBuffers) { + readSpans.emplace_back(buf.data(), buf.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 = {boost::span(value.data(), value.size())}; + + connector_->multiPut(keys, putSpans); + + // Try to read with buffer too small + std::vector smallBuffer(50); + std::vector> smallSpans = {boost::span(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 = {boost::span(value1.data(), value1.size())}; + + connector_->multiPut(keys, putSpans1); + + // Overwrite with new value + std::vector value2(50, 0xBB); + std::vector> putSpans2 = {boost::span(value2.data(), value2.size())}; + + connector_->multiPut(keys, putSpans2); + + // Verify new value + std::vector readBuffer(50); + std::vector> readSpans = {boost::span(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()); +#ifdef USE_REDIS + configs.push_back(NcsTestConfig::Redis()); +#endif + return configs; +} +} // namespace (anonymous) + +INSTANTIATE_TEST_SUITE_P( + AllNcsTypes, + NcsTest, + ::testing::ValuesIn(GetNcsTestConfigs()), + NcsTestNameGenerator +); + +} // namespace milvus