diff --git a/src/impl/ClientRequestContext.cpp b/src/impl/ClientRequestContext.cpp new file mode 100644 index 00000000..f70f0a22 --- /dev/null +++ b/src/impl/ClientRequestContext.cpp @@ -0,0 +1,70 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include "milvus/ClientRequestContext.h" + +#include +#include +#include +#include + +namespace milvus { +namespace { +thread_local std::string request_id; +} + +void +ClientRequestContext::Set(const std::string& value) { + request_id = value; +} + +const std::string& +ClientRequestContext::Get() { + return request_id; +} + +void +ClientRequestContext::Clear() { + request_id.clear(); +} + +std::string +ClientRequestContext::NewRequestId() { + std::random_device device; + std::mt19937_64 generator(device()); + std::uniform_int_distribution distribution; + uint64_t high = 0; + uint64_t low = 0; + do { + high = distribution(generator); + low = distribution(generator); + } while (high == 0 && low == 0); + std::ostringstream stream; + stream << std::hex << std::setfill('0') << std::setw(16) << high << std::setw(16) << low; + return stream.str(); +} + +bool +ClientRequestContext::IsValid(const std::string& value) { + return value.size() == 32 && value != std::string(32, '0') && + std::all_of(value.begin(), value.end(), [](char character) { + return (character >= '0' && character <= '9') || (character >= 'a' && character <= 'f'); + }); +} + +ScopedClientRequestId::ScopedClientRequestId(const std::string& value) : previous_(ClientRequestContext::Get()) { + ClientRequestContext::Set(value); +} + +ScopedClientRequestId::~ScopedClientRequestId() { + ClientRequestContext::Set(previous_); +} + +} // namespace milvus diff --git a/src/impl/ClientTelemetry.cpp b/src/impl/ClientTelemetry.cpp new file mode 100644 index 00000000..4fc3b3cb --- /dev/null +++ b/src/impl/ClientTelemetry.cpp @@ -0,0 +1,1402 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include "milvus/ClientTelemetry.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +#include +#endif + +#include +#include + +#include + +#include "common.pb.h" +#include "milvus.grpc.pb.h" +#include "milvus.pb.h" +#include "milvus/ClientRequestContext.h" + +namespace milvus { +namespace { + +constexpr size_t kSampleBufferSize = 1000; +constexpr size_t kStoredQuantileSampleCount = 128; +constexpr int64_t kMaxHistoryRangeMs = 60 * 60 * 1000; +// A one-second heartbeat produces at most 3,601 boundary-inclusive snapshots in +// one hour. Keep the complete protocol window plus margin while retaining a hard +// memory bound for shorter server-pushed intervals. Each operation/window stores +// at most 128 sorted, evenly spaced quantile samples including its minimum and +// maximum. Across seven supported operations this caps latency history at +// 3,670,016 doubles (about 28 MiB), plus snapshot/collection metric metadata. +constexpr size_t kSnapshotLimit = 4096; +// Fixed-point unit for accumulating a fractional sampling rate. A rate becomes an integer +// step of this many units, so the smallest rate that still samples is 1e-9 -- far below +// anything an operator would set, which is the point: a configured rate must never round +// down to "off". +constexpr uint64_t kSamplingScale = 1000000000ULL; +constexpr size_t kMaxReplyBytes = 1024 * 1024; +constexpr uint64_t kMaxUnsupportedBackoffMs = 30 * 60 * 1000; + +TelemetryConfig +NormalizedTelemetryConfig(TelemetryConfig config) { + if (config.heartbeat_interval_ms == 0) { + config.heartbeat_interval_ms = 10000; + } + config.sampling_rate = std::max(0.0, std::min(1.0, config.sampling_rate)); + if (config.error_max_count == 0) { + config.error_max_count = 100; + } + return config; +} + +bool +SameTelemetryConfig(const TelemetryConfig& left, const TelemetryConfig& right) { + return left.enabled == right.enabled && left.heartbeat_interval_ms == right.heartbeat_interval_ms && + left.sampling_rate == right.sampling_rate && left.error_max_count == right.error_max_count && + left.client_id == right.client_id; +} + +int64_t +NowMillis() { + return std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()) + .count(); +} + +std::string +RandomUuid() { + std::random_device device; + std::mt19937_64 generator(device()); + std::uniform_int_distribution distribution; + auto high = distribution(generator); + auto low = distribution(generator); + std::ostringstream stream; + stream << std::hex << std::setfill('0') << std::setw(8) << static_cast(high >> 32) << "-" << std::setw(4) + << static_cast(high >> 16) << "-" << std::setw(4) << static_cast(high) << "-" + << std::setw(4) << static_cast(low >> 48) << "-" << std::setw(12) << (low & 0x0000FFFFFFFFFFFFULL); + return stream.str(); +} + +std::string +LocalTimeString() { + auto now = std::chrono::system_clock::now(); + auto value = std::chrono::system_clock::to_time_t(now); + std::tm time{}; +#ifdef _WIN32 + gmtime_s(&time, &value); +#else + gmtime_r(&value, &time); +#endif + std::ostringstream stream; + stream << std::put_time(&time, "%Y-%m-%dT%H:%M:%SZ"); + return stream.str(); +} + +std::string +HostName() { +#ifdef _WIN32 + const char* value = std::getenv("COMPUTERNAME"); + return value == nullptr ? "Unknown" : value; +#else + std::array buffer{}; + if (gethostname(buffer.data(), buffer.size()) != 0) { + return "Unknown"; + } + buffer.back() = '\0'; + return {buffer.data()}; +#endif +} + +std::string +CollectionName(const google::protobuf::Message& request) { + const auto* field = request.GetDescriptor()->FindFieldByName("collection_name"); + if (field == nullptr || field->cpp_type() != google::protobuf::FieldDescriptor::CPPTYPE_STRING) { + return ""; + } + return request.GetReflection()->GetString(request, field); +} + +uint32_t +RotateRight(uint32_t value, uint32_t count) { + return (value >> count) | (value << (32 - count)); +} + +// Small self-contained SHA-256 implementation keeps the SDK independent from a specific TLS provider. +class Sha256 { + public: + Sha256() : state_{0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19} { + } + + void + Update(const std::string& value) { + update(reinterpret_cast(value.data()), value.size()); + } + + std::string + Finish() { + uint64_t bit_length = total_size_ * 8; + buffer_[buffer_size_++] = 0x80; + if (buffer_size_ > 56) { + while (buffer_size_ < 64) { + buffer_[buffer_size_++] = 0; + } + transform(buffer_.data()); + buffer_size_ = 0; + } + while (buffer_size_ < 56) { + buffer_[buffer_size_++] = 0; + } + for (int index = 7; index >= 0; --index) { + buffer_[buffer_size_++] = static_cast(bit_length >> (index * 8)); + } + transform(buffer_.data()); + + std::ostringstream stream; + stream << std::hex << std::setfill('0'); + for (auto value : state_) { + stream << std::setw(8) << value; + } + return stream.str(); + } + + private: + void + update(const uint8_t* data, size_t size) { + total_size_ += size; + for (size_t index = 0; index < size; ++index) { + buffer_[buffer_size_++] = data[index]; + if (buffer_size_ == 64) { + transform(buffer_.data()); + buffer_size_ = 0; + } + } + } + + void + transform(const uint8_t* block) { + static const std::array constants = { + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, + 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, + 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, + 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, + 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, + 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, + 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, + 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2}; + std::array schedule{}; + for (size_t index = 0; index < 16; ++index) { + schedule[index] = + (static_cast(block[index * 4]) << 24) | (static_cast(block[index * 4 + 1]) << 16) | + (static_cast(block[index * 4 + 2]) << 8) | static_cast(block[index * 4 + 3]); + } + for (size_t index = 16; index < 64; ++index) { + uint32_t first = RotateRight(schedule[index - 15], 7) ^ RotateRight(schedule[index - 15], 18) ^ + (schedule[index - 15] >> 3); + uint32_t second = RotateRight(schedule[index - 2], 17) ^ RotateRight(schedule[index - 2], 19) ^ + (schedule[index - 2] >> 10); + schedule[index] = schedule[index - 16] + first + schedule[index - 7] + second; + } + uint32_t a = state_[0], b = state_[1], c = state_[2], d = state_[3]; + uint32_t e = state_[4], f = state_[5], g = state_[6], h = state_[7]; + for (size_t index = 0; index < 64; ++index) { + uint32_t sum1 = RotateRight(e, 6) ^ RotateRight(e, 11) ^ RotateRight(e, 25); + uint32_t choice = (e & f) ^ ((~e) & g); + uint32_t temp1 = h + sum1 + choice + constants[index] + schedule[index]; + uint32_t sum0 = RotateRight(a, 2) ^ RotateRight(a, 13) ^ RotateRight(a, 22); + uint32_t majority = (a & b) ^ (a & c) ^ (b & c); + uint32_t temp2 = sum0 + majority; + h = g; + g = f; + f = e; + e = d + temp1; + d = c; + c = b; + b = a; + a = temp1 + temp2; + } + state_[0] += a; + state_[1] += b; + state_[2] += c; + state_[3] += d; + state_[4] += e; + state_[5] += f; + state_[6] += g; + state_[7] += h; + } + + std::array state_; + std::array buffer_{}; + size_t buffer_size_{0}; + uint64_t total_size_{0}; +}; + +int64_t +DaysFromCivil(int year, unsigned month, unsigned day) { + year -= month <= 2; + const int era = (year >= 0 ? year : year - 399) / 400; + const auto year_of_era = static_cast(year - era * 400); + const unsigned adjusted_month = month > 2 ? month - 3 : month + 9; + const unsigned day_of_year = (153 * adjusted_month + 2) / 5 + day - 1; + const unsigned day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + return static_cast(era) * 146097 + static_cast(day_of_era) - 719468; +} + +int64_t +ParseRfc3339Millis(const std::string& value) { + if (value.size() < 20 || value[4] != '-' || value[7] != '-' || value[10] != 'T' || value[13] != ':' || + value[16] != ':') { + throw std::invalid_argument("invalid RFC3339 timestamp"); + } + for (size_t index = 0; index < 19; ++index) { + if (index == 4 || index == 7 || index == 10 || index == 13 || index == 16) { + continue; + } + if (!std::isdigit(static_cast(value[index]))) { + throw std::invalid_argument("invalid RFC3339 timestamp"); + } + } + const int year = std::stoi(value.substr(0, 4)); + const int month = std::stoi(value.substr(5, 2)); + const int day = std::stoi(value.substr(8, 2)); + const int hour = std::stoi(value.substr(11, 2)); + const int minute = std::stoi(value.substr(14, 2)); + const int second = std::stoi(value.substr(17, 2)); + const bool leap_year = year % 4 == 0 && (year % 100 != 0 || year % 400 == 0); + constexpr std::array kDaysPerMonth = {31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}; + if (month < 1 || month > 12 || day < 1 || + day > kDaysPerMonth[static_cast(month - 1)] + (month == 2 && leap_year ? 1 : 0) || hour > 23 || + minute > 59 || second > 59) { + throw std::invalid_argument("invalid RFC3339 timestamp"); + } + size_t position = 19; + int64_t milliseconds = 0; + if (position < value.size() && value[position] == '.') { + size_t end = position + 1; + while (end < value.size() && std::isdigit(static_cast(value[end]))) { + ++end; + } + if (end == position + 1) { + throw std::invalid_argument("invalid RFC3339 timestamp"); + } + auto fraction = value.substr(position + 1, end - position - 1); + while (fraction.size() < 3) { + fraction.push_back('0'); + } + milliseconds = std::stoll(fraction.substr(0, 3)); + position = end; + } + int offset_seconds = 0; + if (position >= value.size()) { + throw std::invalid_argument("invalid RFC3339 timezone"); + } + if (value[position] == 'Z') { + ++position; + } else { + if (position + 6 != value.size() || (value[position] != '+' && value[position] != '-') || + value[position + 3] != ':') { + throw std::invalid_argument("invalid RFC3339 timezone"); + } + for (auto index : {position + 1, position + 2, position + 4, position + 5}) { + if (!std::isdigit(static_cast(value[index]))) { + throw std::invalid_argument("invalid RFC3339 timezone"); + } + } + int sign = value[position] == '+' ? 1 : -1; + int hours = std::stoi(value.substr(position + 1, 2)); + int minutes = std::stoi(value.substr(position + 4, 2)); + if (hours > 23 || minutes > 59) { + throw std::invalid_argument("invalid RFC3339 timezone"); + } + offset_seconds = sign * (hours * 3600 + minutes * 60); + position += 6; + } + if (position != value.size()) { + throw std::invalid_argument("invalid RFC3339 timestamp"); + } + const auto seconds = DaysFromCivil(year, static_cast(month), static_cast(day)) * 86400 + + hour * 3600 + minute * 60 + second - offset_seconds; + return seconds * 1000 + milliseconds; +} + +struct MetricBucket { + int64_t requests{0}; + int64_t successes{0}; + int64_t errors{0}; + double total_latency_ms{0}; + double max_latency_ms{0}; + std::deque samples; + + void + Record(double latency_ms, bool success) { + ++requests; + success ? ++successes : ++errors; + total_latency_ms += latency_ms; + max_latency_ms = std::max(max_latency_ms, latency_ms); + samples.push_back(latency_ms); + if (samples.size() > kSampleBufferSize) { + samples.pop_front(); + } + } + + TelemetryMetric + Snapshot(std::vector* quantile_samples = nullptr) const { + std::vector sorted(samples.begin(), samples.end()); + std::sort(sorted.begin(), sorted.end()); + if (quantile_samples != nullptr) { + quantile_samples->clear(); + const auto count = std::min(sorted.size(), kStoredQuantileSampleCount); + quantile_samples->reserve(count); + if (count == 1) { + quantile_samples->push_back(sorted.front()); + } else if (count > 1) { + // Index 0 and count-1 map exactly to the minimum and maximum. The + // intermediate integer indices are evenly spaced over the sorted + // reservoir and keep the output compact and already merge-ready. + for (size_t index = 0; index < count; ++index) { + const auto source_index = index * (sorted.size() - 1) / (count - 1); + quantile_samples->push_back(sorted[source_index]); + } + } + } + auto index = sorted.empty() ? 0 : std::min(sorted.size() - 1, static_cast(sorted.size() * 0.99)); + return {requests, + successes, + errors, + requests == 0 ? 0 : total_latency_ms / requests, + sorted.empty() ? 0 : sorted[index], + max_latency_ms}; + } +}; + +struct OperationCollector { + MetricBucket global; + std::unordered_map collections; +}; + +struct StoredSnapshot { + TelemetrySnapshot snapshot; + std::unordered_map> global_samples; +}; + +proto::common::Metrics +ToProtoMetric(const TelemetryMetric& metric) { + proto::common::Metrics result; + result.set_request_count(metric.request_count); + result.set_success_count(metric.success_count); + result.set_error_count(metric.error_count); + result.set_avg_latency_ms(metric.avg_latency_ms); + result.set_p99_latency_ms(metric.p99_latency_ms); + result.set_max_latency_ms(metric.max_latency_ms); + return result; +} + +nlohmann::json +MetricJson(const TelemetryMetric& metric) { + return {{"request_count", metric.request_count}, {"success_count", metric.success_count}, + {"error_count", metric.error_count}, {"avg_latency_ms", metric.avg_latency_ms}, + {"p99_latency_ms", metric.p99_latency_ms}, {"max_latency_ms", metric.max_latency_ms}}; +} + +TelemetryCommandReply +SuccessReply(const std::string& command_id, std::string payload = "") { + return {command_id, true, "", std::move(payload)}; +} + +TelemetryCommandReply +FailedReply(const std::string& command_id, const std::string& error) { + return {command_id, false, error, ""}; +} + +} // namespace + +class ClientTelemetryManager::Impl : public std::enable_shared_from_this { + public: + Impl(const TelemetryConfig& value, const std::string& runtime_client_id) + : config(NormalizedTelemetryConfig(value)), + connection_config(config), + stable_client_id(!value.client_id.empty()), + client_id(stable_client_id ? value.client_id + : (runtime_client_id.empty() ? RandomUuid() : runtime_client_id)) { + RegisterDefaultHandlers(); + } + + ~Impl() { + Stop(); + } + + void + AttachChannel(const std::shared_ptr& channel, const std::string& user, const std::string& db, + const std::string& endpoint, const std::string& version, const std::string& scope) { + // Construct every potentially-throwing part before taking either lock. Once locked, only + // noexcept shared_ptr/string moves and scalar updates remain, so a failed candidate cannot + // partially replace the live telemetry transport or its identity fields. + auto stub_holder = proto::milvus::ClientTelemetryService::NewStub(channel); + auto new_stub = std::shared_ptr(std::move(stub_holder)); + std::string new_username = user; + std::string new_database = db; + std::string new_uri = endpoint; + std::string new_sdk_version = version; + std::string new_connection_scope = scope; + + std::lock_guard command_lock(command_mutex); + std::lock_guard lock(mutex); + stub = std::move(new_stub); + ++channel_generation; + username = std::move(new_username); + database = std::move(new_database); + uri = std::move(new_uri); + sdk_version = std::move(new_sdk_version); + connection_scope = std::move(new_connection_scope); + } + + void + Start() { + std::unique_lock lock(mutex); + if (ready) { + return; + } + const bool called_from_worker = worker_running && worker_id == std::this_thread::get_id(); + if (called_from_worker) { + // Stop()/Start() is used by reconnect handlers. Reuse this worker only + // when no external caller has claimed it for joining. An external stop + // always wins over a self-restart that races with it. + if (join_in_progress || external_stop_requested) { + return; + } + ready = true; + control_plane_activated = control_plane_activated || config.enabled; + if (!control_plane_activated) { + stopped = true; + return; + } + stopped = false; + return; + } + + condition.wait(lock, [this]() { return !join_in_progress; }); + if (ready) { + return; + } + if (worker.joinable()) { + join_in_progress = true; + auto previous_worker = std::move(worker); + lock.unlock(); + previous_worker.join(); + lock.lock(); + worker_id = {}; + worker_running = false; + join_in_progress = false; + condition.notify_all(); + if (ready) { + return; + } + } + external_stop_requested = false; + ready = true; + // Initial enabled=false is an explicit opt-out and creates no control-plane traffic. + // Once activated, keep the control plane sticky across dynamic disable and + // Stop/Start reconnects so the server can later re-enable telemetry. + control_plane_activated = control_plane_activated || config.enabled; + if (!control_plane_activated) { + stopped = true; + return; + } + stopped = false; + worker_running = true; + try { + auto self = shared_from_this(); + worker = std::thread([self = std::move(self)]() { self->HeartbeatLoop(); }); + worker_id = worker.get_id(); + } catch (...) { + // Telemetry startup is best-effort. Restore a clean, retryable stopped state instead + // of leaking an exception through Connect() or leaving worker_running without a thread. + ready = false; + stopped = true; + worker_running = false; + worker_id = {}; + condition.notify_all(); + } + } + + void + Stop() { + std::unique_lock lock(mutex); + stopped = true; + ready = false; + condition.notify_all(); + + const bool called_from_worker = worker_running && worker_id == std::this_thread::get_id(); + if (called_from_worker) { + return; + } + + external_stop_requested = true; + condition.wait(lock, [this]() { return !join_in_progress; }); + if (!worker.joinable()) { + worker_id = {}; + worker_running = false; + return; + } + + // Move the thread object while holding the state mutex. This gives + // exactly one external caller ownership of join(); other Stop()/Start() + // calls wait on join_in_progress instead of touching the same std::thread. + join_in_progress = true; + auto current_worker = std::move(worker); + lock.unlock(); + current_worker.join(); + lock.lock(); + worker_id = {}; + worker_running = false; + join_in_progress = false; + condition.notify_all(); + } + + void + HeartbeatLoop() { + while (true) { + try { + CreateSnapshot(); + SendHeartbeat(); + } catch (...) { + // No telemetry collection, serialization, transport, or extension failure may + // escape a std::thread entry and terminate the process. Keep the control plane + // alive; the next iteration can retry on the same or a reattached transport. + try { + std::lock_guard lock(mutex); + last_heartbeat_error = "unexpected client telemetry heartbeat failure"; + } catch (...) { + // Even recording a best-effort diagnostic can fail under memory pressure. + } + } + std::unique_lock lock(mutex); + if (stopped) { + break; + } + uint64_t delay = config.heartbeat_interval_ms; + if (unsupported_streak > 0) { + uint64_t backed_off = std::min(delay, kMaxUnsupportedBackoffMs); + for (int index = 0; index < unsupported_streak && backed_off < kMaxUnsupportedBackoffMs; ++index) { + backed_off = backed_off > kMaxUnsupportedBackoffMs / 2 ? kMaxUnsupportedBackoffMs : backed_off * 2; + } + delay = std::max(delay, backed_off); + } + condition.wait_for(lock, std::chrono::milliseconds(delay), [this]() { return stopped; }); + if (stopped) { + break; + } + } + std::lock_guard lock(mutex); + worker_running = false; + // If Stop() was called by a command handler, no external thread owns + // join(). Detach only after the loop has finished using this object; + // the worker's shared_ptr keeps Impl alive until this function returns. + if (worker.joinable() && worker.get_id() == std::this_thread::get_id()) { + worker.detach(); + worker_id = {}; + } + condition.notify_all(); + } + + void + CreateSnapshot() { + std::lock_guard lock(mutex); + if (!config.enabled) { + return; + } + StoredSnapshot stored; + auto& snapshot = stored.snapshot; + snapshot.end_time = NowMillis(); + snapshot.timestamp = last_snapshot_end == 0 || last_snapshot_end > snapshot.end_time + ? snapshot.end_time - config.heartbeat_interval_ms + : last_snapshot_end; + last_snapshot_end = snapshot.end_time; + for (auto& entry : collectors) { + if (entry.second.global.requests == 0) { + continue; + } + TelemetryOperationMetrics operation; + operation.operation = entry.first; + std::vector quantile_samples; + operation.global = entry.second.global.Snapshot(&quantile_samples); + stored.global_samples.emplace(entry.first, std::move(quantile_samples)); + for (const auto& collection : entry.second.collections) { + if (all_collections_enabled || enabled_collections.count(collection.first) > 0) { + operation.collection_metrics.emplace(collection.first, collection.second.Snapshot()); + } + } + snapshot.metrics.emplace_back(std::move(operation)); + entry.second = OperationCollector{}; + } + // Heartbeats report exactly the collector interval that just ended. + // Keep this separate from retained history so skipping an empty history + // entry cannot make the next heartbeat resend an older non-empty sample. + latest_snapshot = snapshot; + const auto history_start = snapshot.end_time - kMaxHistoryRangeMs; + while (!snapshots.empty() && snapshots.front().snapshot.end_time < history_start) { + snapshots.pop_front(); + } + if (snapshot.metrics.empty()) { + return; + } + snapshots.push_back(std::move(stored)); + while (snapshots.size() > kSnapshotLimit) { + snapshots.pop_front(); + } + } + + void + SendHeartbeat() { + proto::milvus::ClientHeartbeatRequest request; + std::shared_ptr heartbeat_stub; + uint64_t heartbeat_generation = 0; + size_t reply_count = 0; + { + std::lock_guard lock(mutex); + if (stub == nullptr) { + return; + } + auto* info = request.mutable_client_info(); + info->set_sdk_type("CPP"); + info->set_sdk_version(sdk_version); + info->set_local_time(LocalTimeString()); + info->set_user(username); + info->set_host(HostName()); + (*info->mutable_reserved())["client_id"] = client_id; + (*info->mutable_reserved())["client_id_stable"] = stable_client_id ? "true" : "false"; + if (!database.empty()) { + (*info->mutable_reserved())["db_name"] = database; + } + request.set_report_timestamp(NowMillis()); + // Do not resend the final enabled snapshot after collection is disabled. Replies, + // config hash, cursor and incoming commands remain active as the control plane. + if (config.enabled) { + for (const auto& operation : latest_snapshot.metrics) { + auto* output = request.add_metrics(); + output->set_operation(operation.operation); + *output->mutable_global() = ToProtoMetric(operation.global); + for (const auto& collection : operation.collection_metrics) { + if (all_collections_enabled || enabled_collections.count(collection.first) > 0) { + (*output->mutable_collection_metrics())[collection.first] = + ToProtoMetric(collection.second); + } + } + } + } + for (const auto& reply : pending_replies) { + auto* output = request.add_command_replies(); + output->set_command_id(reply.command_id); + output->set_success(reply.success); + output->set_error_message(reply.error_message); + output->set_payload(reply.payload); + } + reply_count = pending_replies.size(); + request.set_config_hash(config_hash); + request.set_last_command_timestamp(last_command_timestamp); + heartbeat_stub = stub; + heartbeat_generation = channel_generation; + } + + grpc::ClientContext context; + context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10)); + proto::milvus::ClientHeartbeatResponse response; + auto grpc_status = heartbeat_stub->ClientHeartbeat(&context, request, &response); + if (!grpc_status.ok()) { + std::lock_guard lock(mutex); + if (heartbeat_generation != channel_generation) { + return; + } + last_heartbeat_error = grpc_status.error_message(); + if (grpc_status.error_code() == grpc::StatusCode::UNIMPLEMENTED) { + ++unsupported_streak; + } + return; + } + { + std::lock_guard lock(mutex); + if (heartbeat_generation != channel_generation) { + return; + } + // Reaching any server implementation proves the RPC exists, even when the + // response carries a business error. + unsupported_streak = 0; + if (response.status().code() != 0 || response.status().error_code() != proto::common::ErrorCode::Success) { + last_heartbeat_error = response.status().reason(); + return; + } + } + std::vector commands; + commands.reserve(response.commands_size()); + for (const auto& command : response.commands()) { + commands.push_back({command.command_id(), command.command_type(), command.payload(), command.create_time(), + command.persistent(), command.target_scope()}); + } + { + std::lock_guard lock(mutex); + if (heartbeat_generation != channel_generation) { + return; + } + pending_replies.erase(pending_replies.begin(), + pending_replies.begin() + std::min(reply_count, pending_replies.size())); + last_heartbeat_error.clear(); + } + ProcessCommands(commands, heartbeat_generation); + } + + TelemetryCommandReply + HandleCommand(const TelemetryCommand& command) { + CommandHandler handler; + { + std::lock_guard lock(mutex); + auto iterator = handlers.find(command.command_type); + if (iterator == handlers.end()) { + return FailedReply(command.command_id, "unknown command type: " + command.command_type); + } + handler = iterator->second; + } + try { + return handler(command); + } catch (const std::exception& exception) { + return FailedReply(command.command_id, exception.what()); + } catch (...) { + // Command handlers are extensible application code. Telemetry is best-effort and + // must never terminate the process when a handler throws a non-standard exception. + return FailedReply(command.command_id, "command handler threw a non-standard exception"); + } + } + + void + ProcessCommands(const std::vector& commands, uint64_t expected_generation = 0) { + std::lock_guard command_lock(command_mutex); + int64_t previous_timestamp; + { + std::lock_guard lock(mutex); + if (expected_generation != 0 && expected_generation != channel_generation) { + return; + } + previous_timestamp = last_command_timestamp; + } + int64_t max_timestamp = previous_timestamp; + bool has_persistent = false; + for (const auto& command : commands) { + max_timestamp = std::max(max_timestamp, command.create_time); + has_persistent = has_persistent || command.persistent; + bool skip = false; + { + std::lock_guard lock(mutex); + skip = command.create_time < previous_timestamp || executed_commands.count(command.command_id) > 0; + if (skip) { + pending_replies.push_back(SuccessReply(command.command_id)); + } + } + if (skip) { + continue; + } + auto reply = HandleCommand(command); + std::lock_guard lock(mutex); + executed_commands[command.command_id] = command.create_time; + pending_replies.push_back(std::move(reply)); + } + std::lock_guard lock(mutex); + for (auto iterator = executed_commands.begin(); iterator != executed_commands.end();) { + if (iterator->second < max_timestamp) { + iterator = executed_commands.erase(iterator); + } else { + ++iterator; + } + } + if (has_persistent) { + config_hash = ClientTelemetryManager::CalculateConfigHash(commands); + } + last_command_timestamp = std::max(last_command_timestamp, max_timestamp); + } + + void + RegisterDefaultHandlers() { + handlers["push_config"] = [this](const TelemetryCommand& command) { + auto payload = command.payload.empty() ? nlohmann::json::object() : nlohmann::json::parse(command.payload); + if (!payload.is_object()) { + throw std::invalid_argument("push_config payload must be a JSON object"); + } + + std::vector applied; + std::vector ignored; + bool enabled = false; + int64_t interval = 0; + double sampling_rate = 0; + if (payload.count("enabled")) { + if (!payload["enabled"].is_boolean()) { + throw std::invalid_argument("enabled must be a boolean"); + } + enabled = payload["enabled"].get(); + applied.emplace_back("enabled"); + } + if (payload.count("heartbeat_interval_ms")) { + if (!payload["heartbeat_interval_ms"].is_number_integer() && + !payload["heartbeat_interval_ms"].is_number_unsigned()) { + throw std::invalid_argument("heartbeat_interval_ms must be an integer"); + } + interval = payload["heartbeat_interval_ms"].get(); + if (interval <= 0) { + throw std::invalid_argument("heartbeat_interval_ms must be positive"); + } + applied.emplace_back("heartbeat_interval_ms"); + } + if (payload.count("sampling_rate")) { + if (!payload["sampling_rate"].is_number()) { + throw std::invalid_argument("sampling_rate must be a number"); + } + sampling_rate = payload["sampling_rate"].get(); + if (!std::isfinite(sampling_rate)) { + throw std::invalid_argument("sampling_rate must be finite"); + } + sampling_rate = std::max(0.0, std::min(1.0, sampling_rate)); + applied.emplace_back("sampling_rate"); + } + if (payload.count("ttl_seconds")) { + if (!payload["ttl_seconds"].is_number_integer() && !payload["ttl_seconds"].is_number_unsigned()) { + throw std::invalid_argument("ttl_seconds must be an integer"); + } + (void)payload["ttl_seconds"].get(); + } + for (auto iterator = payload.begin(); iterator != payload.end(); ++iterator) { + if (iterator.key() != "enabled" && iterator.key() != "heartbeat_interval_ms" && + iterator.key() != "sampling_rate") { + ignored.push_back(iterator.key()); + } + } + std::sort(ignored.begin(), ignored.end()); + { + std::lock_guard lock(mutex); + if (payload.count("enabled")) { + config.enabled = enabled; + } + if (payload.count("heartbeat_interval_ms")) { + config.heartbeat_interval_ms = static_cast(interval); + } + if (payload.count("sampling_rate")) { + config.sampling_rate = sampling_rate; + } + condition.notify_all(); + } + return SuccessReply(command.command_id, nlohmann::json{{"applied", applied}, {"ignored", ignored}}.dump()); + }; + handlers["collection_metrics"] = [this](const TelemetryCommand& command) { + std::lock_guard lock(mutex); + if (command.payload.empty()) { + std::vector names(enabled_collections.begin(), enabled_collections.end()); + std::sort(names.begin(), names.end()); + nlohmann::json result = {{"enabled_collections", names}, + {"all_collections_enabled", all_collections_enabled}}; + return SuccessReply(command.command_id, result.dump()); + } + auto payload = nlohmann::json::parse(command.payload); + if (!payload.is_object()) { + throw std::invalid_argument("collection_metrics payload must be a JSON object"); + } + if (payload.count("enabled") && !payload["enabled"].is_boolean()) { + throw std::invalid_argument("enabled must be a boolean"); + } + if (payload.count("collections") && !payload["collections"].is_array()) { + throw std::invalid_argument("collections must be an array"); + } + if (payload.count("metrics_types") && !payload["metrics_types"].is_array()) { + throw std::invalid_argument("metrics_types must be an array"); + } + bool enabled = payload.value("enabled", false); + auto collections = payload.value("collections", std::vector{}); + if (payload.count("metrics_types")) { + (void)payload["metrics_types"].get>(); + } + bool wildcard = std::find(collections.begin(), collections.end(), "*") != collections.end(); + if (enabled) { + if (collections.empty()) { + throw std::invalid_argument("collections list cannot be empty when enabled=true"); + } + if (wildcard) { + all_collections_enabled = true; + } else { + enabled_collections.insert(collections.begin(), collections.end()); + } + } else if (wildcard || collections.empty()) { + all_collections_enabled = false; + enabled_collections.clear(); + } else { + for (const auto& collection : collections) { + enabled_collections.erase(collection); + } + } + return SuccessReply(command.command_id); + }; + handlers["show_errors"] = [this](const TelemetryCommand& command) { + auto payload = command.payload.empty() ? nlohmann::json::object() : nlohmann::json::parse(command.payload); + if (!payload.is_object()) { + throw std::invalid_argument("show_errors payload must be a JSON object"); + } + if (payload.count("max_count") && !payload["max_count"].is_number_integer() && + !payload["max_count"].is_number_unsigned()) { + throw std::invalid_argument("max_count must be an integer"); + } + auto requested = payload.value("max_count", int64_t{100}); + auto max_count = static_cast(requested <= 0 ? 100 : requested); + std::vector values; + { + std::lock_guard lock(mutex); + for (auto iterator = errors.rbegin(); iterator != errors.rend() && values.size() < max_count; + ++iterator) { + values.push_back(*iterator); + } + } + if (values.empty()) { + return SuccessReply(command.command_id); + } + nlohmann::json result = nlohmann::json::array(); + for (const auto& error : values) { + nlohmann::json detail = { + {"timestamp", error.timestamp}, {"operation", error.operation}, {"error_msg", error.error_message}}; + if (!error.collection.empty()) { + detail["collection"] = error.collection; + } + if (!error.request_id.empty()) { + detail["request_id"] = error.request_id; + } + result.push_back(std::move(detail)); + } + while (result.dump().size() > kMaxReplyBytes && result.size() > 1) { + result.erase(result.begin() + result.size() / 2, result.end()); + } + auto encoded = result.dump(); + while (encoded.size() > kMaxReplyBytes && result.size() == 1 && + result.at(0).value("error_msg", std::string{}).size() > 1) { + auto message = result.at(0).at("error_msg").get(); + result.at(0)["error_msg"] = + message.substr(0, std::max(1, message.size() / 2)) + "...(truncated)"; + encoded = result.dump(); + } + if (encoded.size() > kMaxReplyBytes) { + throw std::invalid_argument("show_errors response exceeds the 1MB payload limit"); + } + return SuccessReply(command.command_id, encoded); + }; + handlers["get_config"] = [this](const TelemetryCommand& command) { + std::lock_guard lock(mutex); + std::vector collections(enabled_collections.begin(), enabled_collections.end()); + std::sort(collections.begin(), collections.end()); + nlohmann::json user_config = { + {"address", uri}, + {"username", username}, + {"db_name", database}, + {"telemetry_enabled", config.enabled}, + {"telemetry_heartbeat_interval_ms", config.heartbeat_interval_ms}, + {"telemetry_sampling_rate", config.sampling_rate}, + {"enabled_collections", all_collections_enabled ? std::vector{"*"} : collections}, + {"all_collections_enabled", all_collections_enabled}}; + return SuccessReply(command.command_id, nlohmann::json{{"user_config", user_config}}.dump()); + }; + handlers["show_latency_history"] = [this](const TelemetryCommand& command) { + if (command.payload.empty()) { + throw std::invalid_argument("payload is required with start_time and end_time"); + } + auto payload = nlohmann::json::parse(command.payload); + if (!payload.is_object()) { + throw std::invalid_argument("show_latency_history payload must be a JSON object"); + } + if (payload.count("detail") && !payload["detail"].is_boolean()) { + throw std::invalid_argument("detail must be a boolean"); + } + int64_t start = ParseRfc3339Millis(payload.at("start_time").get()); + int64_t end = ParseRfc3339Millis(payload.at("end_time").get()); + if (end < start) { + throw std::invalid_argument("end_time must be after start_time"); + } + if (end - start > 60 * 60 * 1000) { + throw std::invalid_argument("time range cannot exceed 1 hour"); + } + std::vector selected; + { + std::lock_guard lock(mutex); + selected.reserve(snapshots.size()); + for (const auto& stored : snapshots) { + const auto& snapshot = stored.snapshot; + if (snapshot.end_time >= start && snapshot.timestamp <= end) { + selected.push_back(stored); + } + } + } + nlohmann::json response; + if (payload.value("detail", false)) { + response["snapshots"] = nlohmann::json::array(); + for (const auto& stored : selected) { + const auto& snapshot = stored.snapshot; + nlohmann::json metrics = nlohmann::json::object(); + for (const auto& operation : snapshot.metrics) { + metrics[operation.operation] = MetricJson(operation.global); + } + response["snapshots"].push_back( + {{"timestamp", snapshot.timestamp}, {"end_time", snapshot.end_time}, {"metrics", metrics}}); + } + response["total_snapshots"] = selected.size(); + } else { + struct Total { + int64_t requests{0}; + int64_t successes{0}; + int64_t errors{0}; + double average{0}; + double maximum{0}; + struct WeightedSamples { + const std::vector* values; + double weight; + }; + std::vector latency_windows; + }; + std::map totals; + for (const auto& stored : selected) { + const auto& snapshot = stored.snapshot; + for (const auto& operation : snapshot.metrics) { + auto& total = totals[operation.operation]; + total.requests += operation.global.request_count; + total.successes += operation.global.success_count; + total.errors += operation.global.error_count; + total.average += operation.global.avg_latency_ms * operation.global.request_count; + total.maximum = std::max(total.maximum, operation.global.max_latency_ms); + const auto samples = stored.global_samples.find(operation.operation); + if (samples != stored.global_samples.end() && !samples->second.empty()) { + const auto weight = + static_cast(operation.global.request_count) / samples->second.size(); + total.latency_windows.push_back({&samples->second, weight}); + } + } + } + nlohmann::json metrics = nlohmann::json::object(); + for (const auto& entry : totals) { + struct Cursor { + double latency; + size_t window; + size_t sample; + }; + const auto later = [](const Cursor& left, const Cursor& right) { + return left.latency > right.latency; + }; + std::priority_queue, decltype(later)> samples(later); + for (size_t window = 0; window < entry.second.latency_windows.size(); ++window) { + const auto* values = entry.second.latency_windows[window].values; + if (values != nullptr && !values->empty()) { + samples.push({values->front(), window, 0}); + } + } + double p99 = 0; + if (!samples.empty()) { + const auto target = static_cast(entry.second.requests) * 0.99; + double cumulative = 0; + while (!samples.empty()) { + const auto sample = samples.top(); + samples.pop(); + const auto& window = entry.second.latency_windows[sample.window]; + p99 = sample.latency; + cumulative += window.weight; + if (cumulative > target) { + break; + } + const auto next = sample.sample + 1; + if (next < window.values->size()) { + samples.push({(*window.values)[next], sample.window, next}); + } + } + } + metrics[entry.first] = { + {"request_count", entry.second.requests}, + {"success_count", entry.second.successes}, + {"error_count", entry.second.errors}, + {"avg_latency_ms", + entry.second.requests == 0 ? 0 : entry.second.average / entry.second.requests}, + {"p99_latency_ms", p99}, + {"max_latency_ms", entry.second.maximum}}; + } + response = {{"aggregated", {{"start_time", start}, {"end_time", end}, {"metrics", metrics}}}, + {"snapshot_count", selected.size()}}; + } + auto encoded = response.dump(); + if (encoded.size() > kMaxReplyBytes) { + throw std::invalid_argument("response too large, try a smaller time range"); + } + return SuccessReply(command.command_id, encoded); + }; + } + + mutable std::mutex mutex; + std::condition_variable condition; + TelemetryConfig config; + const TelemetryConfig connection_config; + const bool stable_client_id; + const std::string client_id; + std::shared_ptr stub; + uint64_t channel_generation{0}; + std::string username; + std::string database; + std::string uri; + std::string sdk_version; + std::string connection_scope; + std::unordered_map collectors; + std::deque errors; + TelemetrySnapshot latest_snapshot; + std::deque snapshots; + std::vector pending_replies; + std::unordered_map executed_commands; + std::unordered_map handlers; + std::unordered_set enabled_collections; + bool all_collections_enabled{false}; + bool ready{false}; + bool stopped{true}; + bool control_plane_activated{false}; + bool worker_running{false}; + bool join_in_progress{false}; + bool external_stop_requested{false}; + int unsupported_streak{0}; + // Carries the fractional sampling rate between operations, in kSamplingScale units: + // each operation adds the rate and the one that pushes it past a whole unit is the one + // sampled. + uint64_t sampling_accum{0}; + int64_t last_command_timestamp{0}; + int64_t last_snapshot_end{0}; + std::string config_hash; + std::string last_heartbeat_error; + std::thread worker; + std::thread::id worker_id; + std::recursive_mutex command_mutex; +}; + +ClientTelemetryManager::ClientTelemetryManager(const TelemetryConfig& config, const std::string& runtime_client_id) + : impl_(std::make_shared(config, runtime_client_id)) { +} + +ClientTelemetryManager::~ClientTelemetryManager() { + // The worker also owns Impl while it is running. Always request shutdown + // before releasing the manager's reference so destruction from inside a + // command handler cannot free state that HeartbeatLoop is still using. + impl_->Stop(); +} + +void +ClientTelemetryManager::AttachChannel(const std::shared_ptr& channel, const std::string& username, + const std::string& database, const std::string& uri, + const std::string& sdk_version, const std::string& connection_scope) { + impl_->AttachChannel(channel, username, database, uri, sdk_version, connection_scope); +} + +void +ClientTelemetryManager::Start() { + impl_->Start(); +} + +void +ClientTelemetryManager::Stop() { + impl_->Stop(); +} + +bool +ClientTelemetryManager::IsReady() const { + std::lock_guard lock(impl_->mutex); + return impl_->ready; +} + +bool +ClientTelemetryManager::isWorkerThread() const { + std::lock_guard lock(impl_->mutex); + return impl_->worker_running && impl_->worker_id == std::this_thread::get_id(); +} + +bool +ClientTelemetryManager::IsSupported() const { + std::lock_guard lock(impl_->mutex); + return impl_->unsupported_streak == 0; +} + +const std::string& +ClientTelemetryManager::ClientId() const { + return impl_->client_id; +} + +std::string +ClientTelemetryManager::ConfigHash() const { + std::lock_guard lock(impl_->mutex); + return impl_->config_hash; +} + +int64_t +ClientTelemetryManager::LastCommandTimestamp() const { + std::lock_guard lock(impl_->mutex); + return impl_->last_command_timestamp; +} + +TelemetryConfig +ClientTelemetryManager::Config() const { + std::lock_guard lock(impl_->mutex); + return impl_->config; +} + +bool +ClientTelemetryManager::MatchesConnection(const TelemetryConfig& config, const std::string& connection_scope) const { + std::lock_guard lock(impl_->mutex); + return impl_->connection_scope == connection_scope && + SameTelemetryConfig(impl_->connection_config, NormalizedTelemetryConfig(config)); +} + +std::string +ClientTelemetryManager::LastHeartbeatError() const { + std::lock_guard lock(impl_->mutex); + return impl_->last_heartbeat_error; +} + +void +ClientTelemetryManager::RegisterCommandHandler(const std::string& command_type, CommandHandler handler) { + std::lock_guard lock(impl_->mutex); + impl_->handlers[command_type] = std::move(handler); +} + +void +ClientTelemetryManager::RecordOperation(const std::string& operation, const google::protobuf::Message& request, + std::chrono::steady_clock::time_point started, bool success, + const std::string& error_message, const std::string& request_id) { + RecordOperation(operation, operation == "RunAnalyzer" ? std::string{} : CollectionName(request), started, success, + error_message, request_id); +} + +void +ClientTelemetryManager::RecordOperation(const std::string& operation, const std::string& collection, + std::chrono::steady_clock::time_point started, bool success, + const std::string& error_message, const std::string& request_id) { + static const std::unordered_set operations = {"Insert", "Delete", "Upsert", "Search", + "HybridSearch", "Query", "RunAnalyzer"}; + if (operations.count(operation) == 0) { + return; + } + auto latency_us = + std::chrono::duration_cast(std::chrono::steady_clock::now() - started).count(); + auto latency = static_cast(latency_us) / 1000.0; + std::lock_guard lock(impl_->mutex); + if (!impl_->config.enabled) { + return; + } + auto rate = impl_->config.sampling_rate; + bool sampled = rate >= 1.0; + if (rate > 0.0 && rate < 1.0) { + // Sample on the operation that carries the accumulator across a whole unit, so the + // sampled operations are spread evenly: at 0.25 that is every fourth one. The + // ratio has to hold over any stretch of operations, not only over a long one -- + // metrics are reported per heartbeat window, and a window is tens or hundreds of + // operations, so sampling a contiguous run would make each window either complete + // or empty. A rate too small to represent still samples rarely rather than never. + auto step = static_cast(rate * kSamplingScale); + if (step == 0) { + step = 1; + } + auto before = impl_->sampling_accum; + impl_->sampling_accum = before + step; + sampled = impl_->sampling_accum / kSamplingScale != before / kSamplingScale; + } + if (!sampled) { + return; + } + bool collection_enabled = impl_->all_collections_enabled || impl_->enabled_collections.count(collection) > 0; + auto& collector = impl_->collectors[operation]; + collector.global.Record(latency, success); + if (!collection.empty() && collection_enabled) { + collector.collections[collection].Record(latency, success); + } + if (!success) { + impl_->errors.push_back({NowMillis(), operation, error_message, collection, + ClientRequestContext::IsValid(request_id) ? request_id : std::string{}}); + while (impl_->errors.size() > impl_->config.error_max_count) { + impl_->errors.pop_front(); + } + } +} + +std::vector +ClientTelemetryManager::RecentErrors(size_t max_count) const { + std::lock_guard lock(impl_->mutex); + std::vector result; + for (auto iterator = impl_->errors.rbegin(); iterator != impl_->errors.rend() && result.size() < max_count; + ++iterator) { + result.push_back(*iterator); + } + return result; +} + +std::vector +ClientTelemetryManager::MetricsSnapshots() const { + std::lock_guard lock(impl_->mutex); + std::vector result; + result.reserve(impl_->snapshots.size()); + for (const auto& stored : impl_->snapshots) { + result.push_back(stored.snapshot); + } + return result; +} + +std::vector +ClientTelemetryManager::PendingCommandReplies() const { + std::lock_guard lock(impl_->mutex); + return impl_->pending_replies; +} + +void +ClientTelemetryManager::ProcessCommands(const std::vector& commands) { + impl_->ProcessCommands(commands); +} + +std::string +ClientTelemetryManager::CalculateConfigHash(const std::vector& commands) { + std::vector persistent; + for (const auto& command : commands) { + if (command.persistent) { + persistent.push_back(command); + } + } + if (persistent.empty()) { + return ""; + } + std::sort(persistent.begin(), persistent.end(), [](const TelemetryCommand& left, const TelemetryCommand& right) { + return left.command_id < right.command_id; + }); + Sha256 hash; + for (const auto& command : persistent) { + hash.Update(command.command_id); + hash.Update(command.command_type); + hash.Update(command.payload); + } + return hash.Finish().substr(0, 16); +} + +} // namespace milvus diff --git a/src/impl/MilvusClientImpl.cpp b/src/impl/MilvusClientImpl.cpp index c88ee18a..6799ad4a 100644 --- a/src/impl/MilvusClientImpl.cpp +++ b/src/impl/MilvusClientImpl.cpp @@ -20,7 +20,6 @@ #include #include #include -#include #include #include "rg.pb.h" @@ -31,6 +30,7 @@ #include "utils/DmlUtils.h" #include "utils/DqlUtils.h" #include "utils/FieldDataSchema.h" +#include "utils/TelemetryUtils.h" #include "utils/TypeUtils.h" #include "utils/cache/CollectionTsCache.h" #include "utils/cache/SchemaCache.h" @@ -39,7 +39,11 @@ namespace milvus { std::shared_ptr MilvusClient::Create() { - return std::make_shared(); + return {new MilvusClientImpl(), [](MilvusClientImpl* client) noexcept { + auto telemetry = client->GetTelemetry(); + const bool called_from_telemetry_worker = telemetry != nullptr && telemetry->isWorkerThread(); + DeleteClientWithTelemetryWorkerSafety(client, std::move(telemetry), called_from_telemetry_worker); + }}; } MilvusClientImpl::~MilvusClientImpl() { @@ -56,6 +60,11 @@ MilvusClientImpl::Disconnect() { return connection_.Disconnect(); } +ClientTelemetryManagerPtr +MilvusClientImpl::GetTelemetry() const { + return connection_.GetTelemetry(); +} + Status MilvusClientImpl::SetRpcDeadlineMs(uint64_t timeout_ms) { return connection_.SetRpcDeadlineMs(timeout_ms); @@ -978,7 +987,8 @@ MilvusClientImpl::DropIndexProperties(const std::string& collection_name, const Status MilvusClientImpl::Insert(const std::string& collection_name, const std::string& partition_name, const std::vector& fields, DmlResults& results) { - return insert(collection_name, partition_name, fields, results, true); + return InvokeWithTelemetry(connection_, "Insert", collection_name, + [&]() { return insert(collection_name, partition_name, fields, results, true); }); } Status @@ -1057,7 +1067,8 @@ MilvusClientImpl::insert(const std::string& collection_name, const std::string& Status MilvusClientImpl::Insert(const std::string& collection_name, const std::string& partition_name, const EntityRows& rows, DmlResults& results) { - return insert(collection_name, partition_name, rows, results, true); + return InvokeWithTelemetry(connection_, "Insert", collection_name, + [&]() { return insert(collection_name, partition_name, rows, results, true); }); } Status @@ -1130,7 +1141,8 @@ MilvusClientImpl::insert(const std::string& collection_name, const std::string& Status MilvusClientImpl::Upsert(const std::string& collection_name, const std::string& partition_name, const std::vector& fields, DmlResults& results) { - return upsert(collection_name, partition_name, fields, results, true); + return InvokeWithTelemetry(connection_, "Upsert", collection_name, + [&]() { return upsert(collection_name, partition_name, fields, results, true); }); } Status @@ -1226,7 +1238,8 @@ MilvusClientImpl::upsert(const std::string& collection_name, const std::string& Status MilvusClientImpl::Upsert(const std::string& collection_name, const std::string& partition_name, const EntityRows& rows, DmlResults& results) { - return upsert(collection_name, partition_name, rows, results, true); + return InvokeWithTelemetry(connection_, "Upsert", collection_name, + [&]() { return upsert(collection_name, partition_name, rows, results, true); }); } Status @@ -1321,8 +1334,10 @@ MilvusClientImpl::Delete(const std::string& collection_name, const std::string& return Status::OK(); }; - return connection_.Invoke( - pre, &MilvusConnection::Delete, post); + return InvokeWithTelemetry(connection_, "Delete", collection_name, [&]() { + return connection_.Invoke( + pre, &MilvusConnection::Delete, post); + }); } Status @@ -1350,8 +1365,10 @@ MilvusClientImpl::Search(const SearchArguments& arguments, SearchResults& result return ConvertSearchResults(response, pk_name, results); }; - return connection_.Invoke( - validate, pre, &MilvusConnection::Search, nullptr, post); + return InvokeWithTelemetry(connection_, "Search", arguments.CollectionName(), [&]() { + return connection_.Invoke( + validate, pre, &MilvusConnection::Search, nullptr, post); + }); } Status @@ -1446,8 +1463,10 @@ MilvusClientImpl::HybridSearch(const HybridSearchArguments& arguments, SearchRes return ConvertSearchResults(response, pk_name, results); }; - return connection_.Invoke( - validate, pre, &MilvusConnection::HybridSearch, nullptr, post); + return InvokeWithTelemetry(connection_, "HybridSearch", arguments.CollectionName(), [&]() { + return connection_.Invoke( + validate, pre, &MilvusConnection::HybridSearch, nullptr, post); + }); } Status @@ -1461,8 +1480,10 @@ MilvusClientImpl::Query(const QueryArguments& arguments, QueryResults& results) auto post = [&results](const proto::milvus::QueryResults& response) { return ConvertQueryResults(response, results); }; - return connection_.Invoke(pre, &MilvusConnection::Query, - post); + return InvokeWithTelemetry(connection_, "Query", arguments.CollectionName(), [&]() { + return connection_.Invoke( + pre, &MilvusConnection::Query, post); + }); } Status @@ -1529,8 +1550,10 @@ MilvusClientImpl::RunAnalyzer(const RunAnalyzerArguments& arguments, AnalyzerRes return Status::OK(); }; - return connection_.Invoke( - pre, &MilvusConnection::RunAnalyzer, post); + return InvokeWithTelemetry(connection_, "RunAnalyzer", "", [&]() { + return connection_.Invoke( + pre, &MilvusConnection::RunAnalyzer, post); + }); } Status diff --git a/src/impl/MilvusClientImpl.h b/src/impl/MilvusClientImpl.h index f9d74af6..1e38e658 100644 --- a/src/impl/MilvusClientImpl.h +++ b/src/impl/MilvusClientImpl.h @@ -37,6 +37,9 @@ class MilvusClientImpl : public MilvusClient { Status Disconnect() final; + ClientTelemetryManagerPtr + GetTelemetry() const final; + Status SetRpcDeadlineMs(uint64_t timeout_ms) final; diff --git a/src/impl/MilvusClientV2Impl.cpp b/src/impl/MilvusClientV2Impl.cpp index d7e86f54..63632f83 100644 --- a/src/impl/MilvusClientV2Impl.cpp +++ b/src/impl/MilvusClientV2Impl.cpp @@ -20,7 +20,6 @@ #include #include #include -#include #include #include @@ -34,6 +33,7 @@ #include "utils/DqlUtils.h" #include "utils/FieldDataSchema.h" #include "utils/MiscUtils.h" +#include "utils/TelemetryUtils.h" #include "utils/TypeUtils.h" #include "utils/cache/CollectionTsCache.h" #include "utils/cache/SchemaCache.h" @@ -42,7 +42,11 @@ namespace milvus { std::shared_ptr MilvusClientV2::Create() { - return std::make_shared(); + return {new MilvusClientV2Impl(), [](MilvusClientV2Impl* client) noexcept { + auto telemetry = client->GetTelemetry(); + const bool called_from_telemetry_worker = telemetry != nullptr && telemetry->isWorkerThread(); + DeleteClientWithTelemetryWorkerSafety(client, std::move(telemetry), called_from_telemetry_worker); + }}; } MilvusClientV2Impl::~MilvusClientV2Impl() { @@ -74,6 +78,11 @@ MilvusClientV2Impl::Disconnect() { return connection_.Disconnect(); } +ClientTelemetryManagerPtr +MilvusClientV2Impl::GetTelemetry() const { + return connection_.GetTelemetry(); +} + Status MilvusClientV2Impl::SetRpcDeadlineMs(uint64_t timeout_ms) { return connection_.SetRpcDeadlineMs(timeout_ms); @@ -1712,7 +1721,8 @@ MilvusClientV2Impl::DropIndexProperties(const DropIndexPropertiesRequest& reques Status MilvusClientV2Impl::Insert(const InsertRequest& request, InsertResponse& response) { - return insert(request, response, true); + return InvokeWithTelemetry(connection_, "Insert", request.CollectionName(), + [&]() { return insert(request, response, true); }); } Status @@ -1830,7 +1840,8 @@ MilvusClientV2Impl::insert(const InsertRequest& request, InsertResponse& respons Status MilvusClientV2Impl::Upsert(const UpsertRequest& request, UpsertResponse& response) { - return upsert(request, response, true); + return InvokeWithTelemetry(connection_, "Upsert", request.CollectionName(), + [&]() { return upsert(request, response, true); }); } Status @@ -2022,8 +2033,10 @@ MilvusClientV2Impl::Delete(const DeleteRequest& request, DeleteResponse& respons return Status::OK(); }; - return connection_.Invoke( - pre, &MilvusConnection::Delete, post); + return InvokeWithTelemetry(connection_, "Delete", request.CollectionName(), [&]() { + return connection_.Invoke( + pre, &MilvusConnection::Delete, post); + }); } Status @@ -2083,8 +2096,10 @@ MilvusClientV2Impl::search(const SearchRequest& request, SearchResponse& respons return status; }; - return connection_.Invoke( - validate, pre, &MilvusConnection::Search, nullptr, post); + return InvokeWithTelemetry(connection_, "Search", request.CollectionName(), [&]() { + return connection_.Invoke( + validate, pre, &MilvusConnection::Search, nullptr, post); + }); } Status @@ -2206,8 +2221,10 @@ MilvusClientV2Impl::hybridSearch(const HybridSearchRequest& request, HybridSearc return status; }; - return connection_.Invoke( - pre, &MilvusConnection::HybridSearch, post); + return InvokeWithTelemetry(connection_, "HybridSearch", request.CollectionName(), [&]() { + return connection_.Invoke( + pre, &MilvusConnection::HybridSearch, post); + }); } Status @@ -2219,7 +2236,7 @@ MilvusClientV2Impl::Query(const QueryRequest& request, QueryResponse& response) Status MilvusClientV2Impl::query(const std::string& endpoint, const std::string& database_name, const QueryRequest& request, - QueryResponse& response, const std::string& cluster_id) { + QueryResponse& response, const std::string& cluster_id, bool record_telemetry) { auto pre = [this, &endpoint, &database_name, &request, &cluster_id](proto::milvus::QueryRequest& rpc_request) { const auto id_count = request.IDs().GetRowCount(); if (!request.Filter().empty() && id_count != 0) { @@ -2262,8 +2279,11 @@ MilvusClientV2Impl::query(const std::string& endpoint, const std::string& databa return status; }; - return connection_.Invoke(pre, &MilvusConnection::Query, - post); + auto invoke = [&]() { + return connection_.Invoke( + pre, &MilvusConnection::Query, post); + }; + return record_telemetry ? InvokeWithTelemetry(connection_, "Query", request.CollectionName(), invoke) : invoke(); } Status @@ -2273,6 +2293,13 @@ MilvusClientV2Impl::Get(const GetRequest& request, GetResponse& response) { Status MilvusClientV2Impl::get(const GetRequest& request, GetResponse& response, const std::string& cluster_id) { + return InvokeWithTelemetry(connection_, "Query", request.CollectionName(), + [&]() { return getWithoutTelemetry(request, response, cluster_id); }); +} + +Status +MilvusClientV2Impl::getWithoutTelemetry(const GetRequest& request, GetResponse& response, + const std::string& cluster_id) { const auto endpoint = connection_.CurrentEndpoint(); const auto database_name = connection_.CurrentDbName(request.DatabaseName()); CollectionDescPtr collection_desc; @@ -2308,7 +2335,7 @@ MilvusClientV2Impl::get(const GetRequest& request, GetResponse& response, const .AddFilterTemplate(ids_key, filter_template) .WithOutputFields(std::move(output_fields)); - return query(endpoint, database_name, actual_request, response, cluster_id); + return query(endpoint, database_name, actual_request, response, cluster_id, false); } Status @@ -2385,8 +2412,10 @@ MilvusClientV2Impl::RunAnalyzer(const RunAnalyzerRequest& request, RunAnalyzerRe return Status::OK(); }; - return connection_.Invoke( - pre, &MilvusConnection::RunAnalyzer, post); + return InvokeWithTelemetry(connection_, "RunAnalyzer", "", [&]() { + return connection_.Invoke( + pre, &MilvusConnection::RunAnalyzer, post); + }); } Status diff --git a/src/impl/MilvusClientV2Impl.h b/src/impl/MilvusClientV2Impl.h index b98e3499..2f25af05 100644 --- a/src/impl/MilvusClientV2Impl.h +++ b/src/impl/MilvusClientV2Impl.h @@ -37,6 +37,9 @@ class MilvusClientV2Impl : public MilvusClientV2, public std::enable_shared_from Status Disconnect() final; + ClientTelemetryManagerPtr + GetTelemetry() const final; + Status SetRpcDeadlineMs(uint64_t timeout_ms) final; @@ -438,11 +441,14 @@ class MilvusClientV2Impl : public MilvusClientV2, public std::enable_shared_from Status query(const std::string& endpoint, const std::string& database_name, const QueryRequest& request, - QueryResponse& response, const std::string& cluster_id); + QueryResponse& response, const std::string& cluster_id, bool record_telemetry = true); Status get(const GetRequest& request, GetResponse& response, const std::string& cluster_id); + Status + getWithoutTelemetry(const GetRequest& request, GetResponse& response, const std::string& cluster_id); + Status queryIterator(QueryIteratorRequest& request, QueryIteratorPtr& iterator, const std::string& cluster_id); diff --git a/src/impl/MilvusConnection.cpp b/src/impl/MilvusConnection.cpp index 8729bcef..5a3579ed 100644 --- a/src/impl/MilvusConnection.cpp +++ b/src/impl/MilvusConnection.cpp @@ -103,13 +103,15 @@ MilvusConnection::StatusCodeFromGrpcStatus(const ::grpc::Status& grpc_status) { } Status -MilvusConnection::Connect(const ConnectParam& param) { +MilvusConnection::Connect(const ConnectParam& param, const std::string& runtime_telemetry_client_id, + ClientTelemetryManagerPtr reusable_telemetry, const std::string& telemetry_logical_endpoint) { std::shared_ptr channel; + std::string telemetry_endpoint; try { // ParseURI() might throw exceptions when the uri/port is invalid std::shared_ptr credentials{nullptr}; auto uri = ParseURI(param.Uri()); - auto address = uri.host + ":" + std::to_string(uri.port); + telemetry_endpoint = uri.host + ":" + std::to_string(uri.port); ::grpc::ChannelArguments args; args.SetMaxSendMessageSize(-1); // max send message size: 2GB @@ -136,7 +138,7 @@ MilvusConnection::Connect(const ConnectParam& param) { metadata["dbname"] = db_name; } - channel = CreateChannelWithHeaderInterceptor(address, credentials, args, metadata); + channel = CreateChannelWithHeaderInterceptor(telemetry_endpoint, credentials, args, metadata); } catch (const std::exception& ex) { std::string reason = "Exception caught when creating grpc channel: "; reason += ex.what(); @@ -150,8 +152,28 @@ MilvusConnection::Connect(const ConnectParam& param) { return {StatusCode::NOT_CONNECTED, reason}; } - auto stub_holder = proto::milvus::MilvusService::NewStub(channel); - auto stub = std::shared_ptr(std::move(stub_holder)); + std::shared_ptr stub; + ClientTelemetryManagerPtr telemetry; + std::string reported_endpoint; + std::string connection_scope; + try { + auto stub_holder = proto::milvus::MilvusService::NewStub(channel); + stub = std::shared_ptr(std::move(stub_holder)); + auto reusable_client_id = + runtime_telemetry_client_id.empty() ? telemetry_client_id_ : runtime_telemetry_client_id; + reported_endpoint = telemetry_logical_endpoint.empty() ? telemetry_endpoint : telemetry_logical_endpoint; + connection_scope = reported_endpoint + "#" + std::to_string(std::hash{}(param.Authorizations())); + const bool can_reuse_telemetry = + reusable_telemetry != nullptr && reusable_telemetry->MatchesConnection(param.Telemetry(), connection_scope); + telemetry = can_reuse_telemetry + ? std::move(reusable_telemetry) + : std::make_shared(param.Telemetry(), reusable_client_id); + } catch (const std::exception& exception) { + return {StatusCode::UNKNOWN_ERROR, + std::string("Failed to prepare client telemetry transport: ") + exception.what()}; + } catch (...) { + return {StatusCode::UNKNOWN_ERROR, "Failed to prepare client telemetry transport"}; + } // grpc channel has been create, now we call the proto::milvus::MilvusClient::Connect() interface // to send some basic information of client to the server, including the sdk type, version, etc. @@ -189,12 +211,25 @@ MilvusConnection::Connect(const ConnectParam& param) { return status; } + try { + telemetry->AttachChannel(channel, param.Username(), param.DbName(), reported_endpoint, GetBuildVersion(), + connection_scope); + } catch (const std::exception& exception) { + return {StatusCode::UNKNOWN_ERROR, + std::string("Failed to attach client telemetry transport: ") + exception.what()}; + } catch (...) { + return {StatusCode::UNKNOWN_ERROR, "Failed to attach client telemetry transport"}; + } { std::lock_guard lock(stub_mtx_); param_ = param; channel_ = std::move(channel); stub_ = std::move(stub); + telemetry_ = telemetry; + telemetry_client_id_ = telemetry->ClientId(); + telemetry_logical_endpoint_ = telemetry_logical_endpoint; } + telemetry->Start(); return Status::OK(); } @@ -203,19 +238,48 @@ MilvusConnection::GetConnectParam() { return param_; } +ClientTelemetryManagerPtr +MilvusConnection::GetTelemetry() const { + std::lock_guard lock(stub_mtx_); + return telemetry_; +} + Status -MilvusConnection::Disconnect() { +MilvusConnection::Disconnect(bool stop_telemetry) { + ClientTelemetryManagerPtr telemetry; + { + std::lock_guard lock(stub_mtx_); + telemetry = telemetry_; + } + if (stop_telemetry && telemetry != nullptr) { + telemetry->Stop(); + } std::lock_guard lock(stub_mtx_); stub_.reset(); channel_.reset(); + telemetry_.reset(); return Status::OK(); } Status MilvusConnection::UseDatabase(const std::string& db_name) { - Disconnect(); - param_.SetDbName(db_name); - return Connect(param_); + ConnectParam candidate_param; + ClientTelemetryManagerPtr telemetry; + std::string telemetry_client_id; + std::string telemetry_logical_endpoint; + { + std::lock_guard lock(stub_mtx_); + candidate_param = param_; + telemetry = telemetry_; + telemetry_client_id = telemetry_client_id_; + telemetry_logical_endpoint = telemetry_logical_endpoint_; + } + candidate_param.SetDbName(db_name); + + // Connect builds and validates a private channel/stub, and only replaces this connection's + // published transport after the handshake and telemetry handoff both succeed. On failure the + // existing database, channel, stub, and telemetry manager remain fully usable. + return Connect(candidate_param, telemetry_client_id, std::move(telemetry), telemetry_logical_endpoint); } Status diff --git a/src/impl/MilvusConnection.h b/src/impl/MilvusConnection.h index 31d708ed..59c87d01 100644 --- a/src/impl/MilvusConnection.h +++ b/src/impl/MilvusConnection.h @@ -27,10 +27,13 @@ #include #include #include +#include #include "common.pb.h" #include "milvus.grpc.pb.h" #include "milvus.pb.h" +#include "milvus/ClientRequestContext.h" +#include "milvus/ClientTelemetry.h" #include "milvus/Status.h" #include "milvus/types/ConnectParam.h" #include "schema.pb.h" @@ -45,11 +48,16 @@ class MilvusConnection { struct GrpcContextOptions { /** timeout in milliseconds */ uint64_t timeout{0}; + /** Optional per-call request ID. Falls back to ClientRequestContext. */ + std::string request_id; // constructors GrpcContextOptions() = default; explicit GrpcContextOptions(uint64_t timeout_) : timeout{timeout_} { } + GrpcContextOptions(uint64_t timeout_, std::string request_id_) + : timeout{timeout_}, request_id{std::move(request_id_)} { + } }; MilvusConnection() = default; @@ -57,13 +65,17 @@ class MilvusConnection { virtual ~MilvusConnection(); Status - Connect(const ConnectParam& param); + Connect(const ConnectParam& param, const std::string& runtime_telemetry_client_id = "", + ClientTelemetryManagerPtr reusable_telemetry = nullptr, const std::string& telemetry_logical_endpoint = ""); ConnectParam& GetConnectParam(); + ClientTelemetryManagerPtr + GetTelemetry() const; + Status - Disconnect(); + Disconnect(bool stop_telemetry = true); Status UseDatabase(const std::string& db_name); @@ -503,10 +515,13 @@ class MilvusConnection { const GrpcContextOptions& options); private: - std::mutex stub_mtx_; + mutable std::mutex stub_mtx_; std::shared_ptr stub_; std::shared_ptr channel_; ConnectParam param_; + ClientTelemetryManagerPtr telemetry_; + std::string telemetry_client_id_; + std::string telemetry_logical_endpoint_; static Status StatusByProtoResponse(const proto::common::Status& status); @@ -529,6 +544,7 @@ class MilvusConnection { grpcCall(const char* name, grpc::Status (proto::milvus::MilvusService::Stub::*func)(grpc::ClientContext*, const Request&, Response*), const Request& request, Response& response, const GrpcContextOptions& options) { + (void)name; std::shared_ptr stub; { std::lock_guard lock(stub_mtx_); @@ -539,6 +555,11 @@ class MilvusConnection { } ::grpc::ClientContext context; + const std::string& contextual_request_id = ClientRequestContext::Get(); + const std::string request_id = options.request_id.empty() ? contextual_request_id : options.request_id; + if (ClientRequestContext::IsValid(request_id)) { + context.AddMetadata("client_request_id", request_id); + } if (options.timeout > 0) { auto deadline = std::chrono::system_clock::now() + std::chrono::milliseconds{options.timeout}; context.set_deadline(deadline); @@ -562,7 +583,8 @@ class MilvusConnection { // Some milvus error codes can be retried: // response.status().error_code() == io.milvus.grpc.ErrorCode.RateLimit // or response.status()code() == 8 can be retried - return StatusByProtoResponse(response); + auto status = StatusByProtoResponse(response); + return status; } }; diff --git a/src/impl/types/ConnectParam.cpp b/src/impl/types/ConnectParam.cpp index 8377a150..f384f0c0 100644 --- a/src/impl/types/ConnectParam.cpp +++ b/src/impl/types/ConnectParam.cpp @@ -70,6 +70,7 @@ ConnectParam::operator=(const ConnectParam& other) { username_ = other.username_; token_ = other.token_; db_name_ = other.db_name_; + telemetry_config_ = other.telemetry_config_; } return *this; } @@ -314,4 +315,20 @@ ConnectParam::WithDbName(const std::string& db_name) { return *this; } +const TelemetryConfig& +ConnectParam::Telemetry() const { + return telemetry_config_; +} + +void +ConnectParam::SetTelemetryConfig(const TelemetryConfig& config) { + telemetry_config_ = config; +} + +ConnectParam& +ConnectParam::WithTelemetryConfig(const TelemetryConfig& config) { + SetTelemetryConfig(config); + return *this; +} + } // namespace milvus diff --git a/src/impl/utils/ConnectionHandler.cpp b/src/impl/utils/ConnectionHandler.cpp index 5cd4fbe4..ff26aab7 100644 --- a/src/impl/utils/ConnectionHandler.cpp +++ b/src/impl/utils/ConnectionHandler.cpp @@ -32,42 +32,31 @@ ConnectionHandler::~ConnectionHandler() { Status ConnectionHandler::Connect(const ConnectParam& connect_param) { - // Snapshot the previous global-cluster state so a failed handshake can restore it: the old - // primary connection_ is kept until the new handshake succeeds, and tearing down the refresher - // before the attempt would otherwise drop global tracking from a still-usable connection. - bool was_global = false; - std::string prior_endpoint; - ConnectParam prior_connect_param; - std::unique_ptr prior_refresher; - { - std::lock_guard lock(mtx_); - was_global = global_mode_; - prior_endpoint = global_endpoint_; - prior_connect_param = global_connect_param_; - prior_refresher = std::move(global_refresher_); - // mark global mode off first so an in-flight callback cannot reconnect during teardown - global_mode_ = false; - } - // join the refresher thread outside the lock (it may be inside reconnectToPrimary waiting on mtx_) - if (prior_refresher != nullptr) { - prior_refresher->Stop(); + // Keep the candidate private until its handshake succeeds, but do not hold mtx_ while + // AttachChannel waits for an in-flight command handler. The lifecycle lock also fences + // concurrent explicit connects, disconnects, database switches, and global failovers. + // It must remain fail-fast: a command handler holds telemetry's command_mutex and may + // re-enter a lifecycle API while an external lifecycle call is waiting in AttachChannel. + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return {StatusCode::CLIENT_BUSY, "Connection lifecycle change is already in progress"}; } - // Restore the prior global state on the failure path (only when a previous connection is still - // live), so cache-key scope and automatic failover are preserved for the old primary. - auto restore = [&]() { + // Snapshot only reusable resources. The published connection and global state remain untouched + // until the candidate handshake succeeds, so every failure is a no-op from callers' perspective. + ClientTelemetryManagerPtr reusable_telemetry; + MilvusConnectionPtr old_connection; + { std::lock_guard lock(mtx_); - if (connection_ == nullptr) { - return; - } - global_mode_ = was_global; - global_endpoint_ = prior_endpoint; - global_connect_param_ = prior_connect_param; - global_refresher_ = std::move(prior_refresher); - if (global_refresher_ != nullptr) { - global_refresher_->Start(); + reusable_telemetry = telemetry_; + old_connection = connection_; + if (old_connection != nullptr) { + auto current_telemetry = old_connection->GetTelemetry(); + if (current_telemetry != nullptr) { + reusable_telemetry = std::move(current_telemetry); + } } - }; + } bool is_global = GlobalClusterUtils::IsGlobalEndpoint(connect_param.Uri()); ConnectParam primary_param = connect_param; @@ -78,72 +67,102 @@ ConnectionHandler::Connect(const ConnectParam& connect_param) { // the connection under mtx_ auto status = GlobalClusterUtils::FetchTopology(connect_param.Uri(), connect_param.Token(), initial_topology); if (!status.IsOk()) { - restore(); return status; } const ClusterInfo* primary = initial_topology.Primary(); if (primary == nullptr) { - restore(); return {StatusCode::SERVER_FAILED, "No primary (writable) cluster found in global topology"}; } primary_param.SetUri( GlobalClusterUtils::BuildPrimaryUri(connect_param.Uri(), connect_param.TlsEnabled(), primary->Endpoint())); } - // Serialize the connection handshake with lifecycle and configuration mutations. The candidate - // connection remains private until it succeeds, but setters must not update the current - // connection and then be overwritten by the successful swap below. - MilvusConnectionPtr new_connection; - Status status; + std::unique_ptr new_refresher; + if (is_global) { + new_refresher = std::make_unique( + connect_param.Uri(), connect_param.Token(), initial_topology.Version(), std::chrono::seconds(300), + [this](const GlobalTopology& topology, const std::function& should_stop) { + return reconnectToPrimary(topology, should_stop); + }); + // Starting a private refresher cannot observe or mutate the live lifecycle before commit: + // its first refresh waits for the configured interval. If thread creation throws, the old + // published state is still intact and the exception is converted to a Status below. Starting + // it before the candidate attaches shared telemetry also keeps that handoff as the final + // infallible step before publication. + try { + new_refresher->Start(); + } catch (const std::exception& exception) { + return {StatusCode::UNKNOWN_ERROR, + std::string("Failed to start global topology refresher: ") + exception.what()}; + } + } + + auto new_connection = std::make_shared(); + auto status = new_connection->Connect(primary_param, telemetry_client_id_, reusable_telemetry, + is_global ? connect_param.Uri() : ""); + if (!status.IsOk()) { + if (new_refresher != nullptr) { + new_refresher->Stop(); + } + return status; + } + + auto telemetry = new_connection->GetTelemetry(); + std::unique_ptr old_refresher; { std::lock_guard lock(mtx_); + old_connection = connection_; + old_refresher = std::move(global_refresher_); + global_mode_ = is_global; + global_endpoint_ = is_global ? connect_param.Uri() : std::string{}; + if (telemetry != nullptr) { + telemetry_ = telemetry; + telemetry_client_id_ = telemetry->ClientId(); + } + connection_ = new_connection; - global_mode_ = false; - global_endpoint_.clear(); - - new_connection = std::make_shared(); - status = new_connection->Connect(primary_param); - if (!status.IsOk()) { - // fall through to restore the prior global state after releasing the lock - } else { - if (connection_ != nullptr) { - connection_->Disconnect(); - } - connection_ = std::move(new_connection); - - // commit the global-cluster state only after the primary connect succeeded, so a failed - // Connect() leaves the handler in non-global mode (consistent with connection_ == null) - if (is_global) { - global_mode_ = true; - global_endpoint_ = connect_param.Uri(); - global_connect_param_ = connect_param; - global_refresher_ = std::make_unique( - global_endpoint_, global_connect_param_.Token(), initial_topology.Version(), - std::chrono::seconds(300), - [this](const GlobalTopology& topology, const std::function& should_stop) { - return reconnectToPrimary(topology, should_stop); - }); - global_refresher_->Start(); - } + // Commit global-cluster state only after the primary connect succeeds. + if (is_global) { + global_connect_param_ = connect_param; + global_refresher_ = std::move(new_refresher); } } - if (!status.IsOk()) { - restore(); - return status; + + // The old refresher can be inside its callback. It cannot wait for lifecycle_mtx_ because + // reconnectToPrimary also uses try_to_lock, so Stop()/join is safe while this lifecycle owns it. + if (old_refresher != nullptr) { + old_refresher->Stop(); + } + if (old_connection != nullptr) { + const bool shares_telemetry = telemetry != nullptr && old_connection->GetTelemetry() == telemetry; + old_connection->Disconnect(!shares_telemetry); } return Status::OK(); } Status ConnectionHandler::Disconnect() { + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return {StatusCode::CLIENT_BUSY, "Connection lifecycle change is already in progress"}; + } + // stop the refresher without holding the lock; callbacks see global_mode_==false and no-op stopGlobalRefresher(); - std::lock_guard lock(mtx_); - if (connection_ != nullptr) { - return connection_->Disconnect(); + MilvusConnectionPtr connection; + { + std::lock_guard lock(mtx_); + connection = std::move(connection_); + if (connection == nullptr) { + return Status::OK(); + } + auto telemetry = connection->GetTelemetry(); + if (telemetry != nullptr) { + telemetry_ = std::move(telemetry); + } } - return Status::OK(); + return connection->Disconnect(); } void @@ -171,6 +190,13 @@ ConnectionHandler::TriggerGlobalRefresh() { bool ConnectionHandler::reconnectToPrimary(const GlobalTopology& topology, const std::function& should_stop) { + // A refresher callback must never wait behind Connect()/Disconnect() while those methods + // are stopping and joining the refresher thread. + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return false; + } + const ClusterInfo* primary = topology.Primary(); if (primary == nullptr) { // no writable cluster in this topology; report failure so the refresher retries the @@ -179,6 +205,9 @@ ConnectionHandler::reconnectToPrimary(const GlobalTopology& topology, const std: } ConnectParam primary_param; + ClientTelemetryManagerPtr reusable_telemetry; + std::string telemetry_client_id; + std::string telemetry_logical_endpoint; { std::lock_guard lock(mtx_); if (!global_mode_) { @@ -199,7 +228,13 @@ ConnectionHandler::reconnectToPrimary(const GlobalTopology& topology, const std: if (connection_ != nullptr) { primary_param.SetDbName(connection_->GetConnectParam().DbName()); primary_param.SetRpcDeadlineMs(connection_->GetConnectParam().RpcDeadlineMs()); + reusable_telemetry = connection_->GetTelemetry(); } + if (reusable_telemetry == nullptr) { + reusable_telemetry = telemetry_; + } + telemetry_client_id = telemetry_client_id_; + telemetry_logical_endpoint = global_endpoint_; } // abort promptly when the refresher is stopping rather than starting a fresh gRPC connect @@ -211,36 +246,54 @@ ConnectionHandler::reconnectToPrimary(const GlobalTopology& topology, const std: // WaitForConnected() and the Connect RPC for up to ~2x ConnectTimeout, and holding mtx_ that // long would stall every other SDK operation that snapshots the connection. auto new_connection = std::make_shared(); - auto status = new_connection->Connect(primary_param); + auto status = + new_connection->Connect(primary_param, telemetry_client_id, reusable_telemetry, telemetry_logical_endpoint); if (!status.IsOk()) { // keep the existing connection; report failure so the refresher retries the same version return false; } + auto new_telemetry = new_connection->GetTelemetry(); + MilvusConnectionPtr old_connection; + bool discard_candidate = false; + bool stop_candidate_telemetry = true; + bool reconnect_result = true; { std::lock_guard lock(mtx_); if (!global_mode_) { // disconnected while reconnecting; discard the unused candidate connection - new_connection->Disconnect(); - return true; - } - // re-read live configuration in case SetRpcDeadlineMs()/UseDatabase() ran while the - // candidate was being built outside the lock, so it is not silently dropped on swap - if (connection_ != nullptr) { + discard_candidate = true; + stop_candidate_telemetry = new_telemetry == nullptr || telemetry_ != new_telemetry; + } else if (connection_ != nullptr) { + // Re-read live configuration before swapping the candidate. const ConnectParam& live = connection_->GetConnectParam(); new_connection->GetConnectParam().SetRpcDeadlineMs(live.RpcDeadlineMs()); if (new_connection->GetConnectParam().DbName() != live.DbName()) { // the database changed while reconnecting; drop the stale candidate and retry - new_connection->Disconnect(); - return false; + discard_candidate = true; + stop_candidate_telemetry = new_telemetry == nullptr || connection_->GetTelemetry() != new_telemetry; + reconnect_result = false; } } - auto old_connection = connection_; - connection_ = std::move(new_connection); - if (old_connection != nullptr) { - old_connection->Disconnect(); + + if (!discard_candidate) { + old_connection = connection_; + connection_ = new_connection; + if (new_telemetry != nullptr) { + telemetry_ = new_telemetry; + telemetry_client_id_ = new_telemetry->ClientId(); + } } } + + if (discard_candidate) { + new_connection->Disconnect(stop_candidate_telemetry); + return reconnect_result; + } + if (old_connection != nullptr) { + const bool shares_telemetry = new_telemetry != nullptr && old_connection->GetTelemetry() == new_telemetry; + old_connection->Disconnect(!shares_telemetry); + } return true; } @@ -250,8 +303,18 @@ ConnectionHandler::GetConnection() const { return connection_; } +ClientTelemetryManagerPtr +ConnectionHandler::GetTelemetry() const { + std::lock_guard lock(mtx_); + return connection_ == nullptr ? telemetry_ : connection_->GetTelemetry(); +} + Status ConnectionHandler::SetRpcDeadlineMs(uint64_t timeout_ms) { + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return {StatusCode::CLIENT_BUSY, "Connection lifecycle change is already in progress"}; + } std::lock_guard lock(mtx_); if (connection_ == nullptr) { return {StatusCode::NOT_CONNECTED, "Connection is not created!"}; @@ -271,6 +334,10 @@ ConnectionHandler::GetRpcDeadlineMs() const { Status ConnectionHandler::SetRetryParam(const RetryParam& retry_param) { + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return {StatusCode::CLIENT_BUSY, "Connection lifecycle change is already in progress"}; + } std::lock_guard lock(mtx_); if (connection_ == nullptr) { return {StatusCode::NOT_CONNECTED, "Connection is not created!"}; @@ -287,12 +354,12 @@ ConnectionHandler::GetRetryParam() const { Status ConnectionHandler::UseDatabase(const std::string& db_name) { - std::lock_guard lock(mtx_); - if (connection_ != nullptr) { - return connection_->UseDatabase(db_name); + std::unique_lock lifecycle_lock(lifecycle_mtx_, std::try_to_lock); + if (!lifecycle_lock.owns_lock()) { + return {StatusCode::CLIENT_BUSY, "Connection lifecycle change is already in progress"}; } - - return Status::OK(); + auto connection = GetConnection(); + return connection == nullptr ? Status::OK() : connection->UseDatabase(db_name); } std::string diff --git a/src/impl/utils/ConnectionHandler.h b/src/impl/utils/ConnectionHandler.h index 4c5b7e88..93c974c8 100644 --- a/src/impl/utils/ConnectionHandler.h +++ b/src/impl/utils/ConnectionHandler.h @@ -49,6 +49,9 @@ class ConnectionHandler { MilvusConnectionPtr GetConnection() const; + ClientTelemetryManagerPtr + GetTelemetry() const; + Status SetRpcDeadlineMs(uint64_t timeout_ms); @@ -277,9 +280,20 @@ class ConnectionHandler { } private: + // Lifecycle lock order and re-entrancy contract: + // 1. Every state-mutating lifecycle entry acquires lifecycle_mtx_ with try_to_lock. + // 2. While it owns lifecycle_mtx_, it may briefly acquire mtx_ to snapshot or commit state. + // 3. mtx_ is never held across network I/O, TopologyRefresher::Stop(), telemetry Stop(), + // AttachChannel(), or ProcessCommands(). + // A telemetry command handler already owns command_mutex and may re-enter a lifecycle API. + // Fail-fast lifecycle acquisition prevents a command_mutex -> lifecycle_mtx_ wait from + // forming a cycle with an external lifecycle_mtx_ -> AttachChannel(command_mutex) handoff. + mutable std::mutex lifecycle_mtx_; mutable std::mutex mtx_; MilvusConnectionPtr connection_; + ClientTelemetryManagerPtr telemetry_; RetryParam retry_param_; + std::string telemetry_client_id_; // global-cluster state bool global_mode_{false}; diff --git a/src/impl/utils/TelemetryUtils.h b/src/impl/utils/TelemetryUtils.h new file mode 100644 index 00000000..e52c4509 --- /dev/null +++ b/src/impl/utils/TelemetryUtils.h @@ -0,0 +1,62 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#pragma once + +#include +#include +#include +#include + +#include "ConnectionHandler.h" +#include "milvus/ClientRequestContext.h" + +namespace milvus { + +template +void +DeleteClientWithTelemetryWorkerSafety(Client* client, ClientTelemetryManagerPtr telemetry, + bool called_from_telemetry_worker) noexcept { + if (!called_from_telemetry_worker) { + delete client; + return; + } + + // The global topology refresher may be waiting for the current command handler before + // it can hand off the telemetry channel. Stop/join the worker on another thread, then + // delete the client after the handler has returned and released the command lock. + try { + std::thread([client, telemetry = std::move(telemetry)]() { + telemetry->Stop(); + delete client; + }).detach(); + } catch (...) { + // A deleter must not throw. If the one-shot cleanup thread cannot be created, + // intentionally retain the client rather than terminate or re-enter the known + // refresher/command join cycle. + } +} + +template +Status +InvokeWithTelemetry(ConnectionHandler& connection, const std::string& operation, const std::string& collection, + Callable&& callable) { + auto started = std::chrono::steady_clock::now(); + auto telemetry = connection.GetTelemetry(); + auto status = std::forward(callable)(); + if (telemetry != nullptr) { + const auto& request_id = ClientRequestContext::Get(); + telemetry->RecordOperation(operation, collection, started, status.IsOk(), + status.IsOk() ? std::string{} : status.Message(), request_id); + } + return status; +} + +} // namespace milvus diff --git a/src/include/milvus/ClientRequestContext.h b/src/include/milvus/ClientRequestContext.h new file mode 100644 index 00000000..25909e5d --- /dev/null +++ b/src/include/milvus/ClientRequestContext.h @@ -0,0 +1,54 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#pragma once + +#include + +#include "milvus/Export.h" + +namespace milvus { + +/** Per-thread request ID propagated as the client_request_id gRPC metadata. */ +class MILVUS_SDK_API ClientRequestContext { + public: + static void + Set(const std::string& request_id); + + static const std::string& + Get(); + + static void + Clear(); + + /** Returns a lowercase 32-character OpenTelemetry-compatible trace ID. */ + static std::string + NewRequestId(); + + /** Returns true for a lowercase, non-zero, 32-character OpenTelemetry trace ID. */ + static bool + IsValid(const std::string& request_id); +}; + +/** Restores the previous thread-local request ID when it leaves scope. */ +class MILVUS_SDK_API ScopedClientRequestId { + public: + explicit ScopedClientRequestId(const std::string& request_id); + ~ScopedClientRequestId(); + + ScopedClientRequestId(const ScopedClientRequestId&) = delete; + ScopedClientRequestId& + operator=(const ScopedClientRequestId&) = delete; + + private: + std::string previous_; +}; + +} // namespace milvus diff --git a/src/include/milvus/ClientTelemetry.h b/src/include/milvus/ClientTelemetry.h new file mode 100644 index 00000000..211cd528 --- /dev/null +++ b/src/include/milvus/ClientTelemetry.h @@ -0,0 +1,171 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "milvus/Export.h" +#include "milvus/types/TelemetryConfig.h" + +namespace grpc { +class Channel; +} + +namespace google { +namespace protobuf { +class Message; +} +} // namespace google + +namespace milvus { + +struct MILVUS_SDK_API TelemetryMetric { + int64_t request_count{0}; + int64_t success_count{0}; + int64_t error_count{0}; + double avg_latency_ms{0}; + double p99_latency_ms{0}; + double max_latency_ms{0}; +}; + +struct MILVUS_SDK_API TelemetryOperationMetrics { + std::string operation; + TelemetryMetric global; + std::unordered_map collection_metrics; +}; + +struct MILVUS_SDK_API TelemetrySnapshot { + int64_t timestamp{0}; + int64_t end_time{0}; + std::vector metrics; +}; + +struct MILVUS_SDK_API TelemetryError { + int64_t timestamp{0}; + std::string operation; + std::string error_message; + std::string collection; + std::string request_id; +}; + +struct MILVUS_SDK_API TelemetryCommand { + std::string command_id; + std::string command_type; + std::string payload; + int64_t create_time{0}; + bool persistent{false}; + std::string target_scope; +}; + +struct MILVUS_SDK_API TelemetryCommandReply { + std::string command_id; + bool success{false}; + std::string error_message; + std::string payload; +}; + +/** Client-side metrics, heartbeat, command, and diagnostic manager. */ +class MILVUS_SDK_API ClientTelemetryManager { + public: + using CommandHandler = std::function; + + explicit ClientTelemetryManager(const TelemetryConfig& config = TelemetryConfig{}, + const std::string& runtime_client_id = ""); + ~ClientTelemetryManager(); + + ClientTelemetryManager(const ClientTelemetryManager&) = delete; + ClientTelemetryManager& + operator=(const ClientTelemetryManager&) = delete; + + void + AttachChannel(const std::shared_ptr& channel, const std::string& username, + const std::string& database, const std::string& uri, const std::string& sdk_version, + const std::string& connection_scope = ""); + + void + Start(); + + void + Stop(); + + bool + IsReady() const; + + bool + IsSupported() const; + + const std::string& + ClientId() const; + + std::string + ConfigHash() const; + + int64_t + LastCommandTimestamp() const; + + TelemetryConfig + Config() const; + + /** Whether a reconnect can reuse this manager without changing user-supplied telemetry settings. */ + bool + MatchesConnection(const TelemetryConfig& config, const std::string& connection_scope) const; + + std::string + LastHeartbeatError() const; + + void + RegisterCommandHandler(const std::string& command_type, CommandHandler handler); + + void + RecordOperation(const std::string& operation, const google::protobuf::Message& request, + std::chrono::steady_clock::time_point started, bool success, const std::string& error_message, + const std::string& request_id = ""); + + void + RecordOperation(const std::string& operation, const std::string& collection, + std::chrono::steady_clock::time_point started, bool success, const std::string& error_message, + const std::string& request_id = ""); + + std::vector + RecentErrors(size_t max_count = 100) const; + + std::vector + MetricsSnapshots() const; + + std::vector + PendingCommandReplies() const; + + void + ProcessCommands(const std::vector& commands); + + static std::string + CalculateConfigHash(const std::vector& commands); + + private: + friend class MilvusClient; + friend class MilvusClientV2; + + bool + isWorkerThread() const; + + class Impl; + std::shared_ptr impl_; +}; + +using ClientTelemetryManagerPtr = std::shared_ptr; + +} // namespace milvus diff --git a/src/include/milvus/MilvusClient.h b/src/include/milvus/MilvusClient.h index 66e19993..a03a278a 100644 --- a/src/include/milvus/MilvusClient.h +++ b/src/include/milvus/MilvusClient.h @@ -18,6 +18,8 @@ #include +#include "ClientRequestContext.h" +#include "ClientTelemetry.h" #include "Status.h" #include "milvus/Export.h" #include "types/AliasDesc.h" @@ -89,6 +91,7 @@ class MILVUS_SDK_API MilvusClient { * @brief Connect to Milvus server. * * @param [in] connect_param server address and port + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -97,6 +100,7 @@ class MILVUS_SDK_API MilvusClient { /** * @brief Break connections between client and server. * + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -105,6 +109,7 @@ class MILVUS_SDK_API MilvusClient { /** * @brief Change timeout value in milliseconds for each RPC call. * + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried */ virtual Status SetRpcDeadlineMs(uint64_t timeout_ms) = 0; @@ -113,6 +118,7 @@ class MILVUS_SDK_API MilvusClient { * @brief Reset retry rules for each RPC call. * * @param [in] retry_param retry rules + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried */ virtual Status SetRetryParam(const RetryParam& retry_param) = 0; @@ -455,6 +461,7 @@ class MILVUS_SDK_API MilvusClient { * @brief Switch connection to another database. * * @param [in] db_name name of the database + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -1145,6 +1152,12 @@ class MILVUS_SDK_API MilvusClient { */ virtual Status RemovePrivilegesFromGroup(const std::string& group_name, const std::vector& privileges) = 0; + + /** Returns the telemetry manager for diagnostics and custom command handlers. */ + virtual ClientTelemetryManagerPtr + GetTelemetry() const { + return nullptr; + } }; using MilvusClientPtr = std::shared_ptr; diff --git a/src/include/milvus/MilvusClientV2.h b/src/include/milvus/MilvusClientV2.h index e242fbed..afc0e4aa 100644 --- a/src/include/milvus/MilvusClientV2.h +++ b/src/include/milvus/MilvusClientV2.h @@ -18,6 +18,8 @@ #include +#include "ClientRequestContext.h" +#include "ClientTelemetry.h" #include "MilvusClientV2Session.h" #include "Status.h" #include "milvus/Export.h" @@ -218,6 +220,7 @@ class MILVUS_SDK_API MilvusClientV2 { * @brief Connect to Milvus server. * * @param [in] connect_param server address and port + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -226,6 +229,7 @@ class MILVUS_SDK_API MilvusClientV2 { /** * @brief Close connections between client and server. * + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -234,6 +238,7 @@ class MILVUS_SDK_API MilvusClientV2 { /** * @brief Change timeout value in milliseconds for each RPC call. * + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried */ virtual Status SetRpcDeadlineMs(uint64_t timeout_ms) = 0; @@ -242,6 +247,7 @@ class MILVUS_SDK_API MilvusClientV2 { * @brief Reset retry rules for each RPC call. * * @param [in] retry_param retry rules + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried */ virtual Status SetRetryParam(const RetryParam& retry_param) = 0; @@ -656,6 +662,7 @@ class MILVUS_SDK_API MilvusClientV2 { * @brief Switch connection to another database. * * @param [in] db_name name of the database + * @retval StatusCode::CLIENT_BUSY another connection lifecycle change is in progress; the operation may be retried * @return Status operation successfully or not */ virtual Status @@ -1435,6 +1442,12 @@ class MILVUS_SDK_API MilvusClientV2 { */ virtual Status Session(const std::string& cluster_id, MilvusClientV2SessionPtr& session) = 0; + + /** Returns the telemetry manager for diagnostics and custom command handlers. */ + virtual ClientTelemetryManagerPtr + GetTelemetry() const { + return nullptr; + } }; using MilvusClientV2Ptr = std::shared_ptr; diff --git a/src/include/milvus/Status.h b/src/include/milvus/Status.h index 11e5b2ba..2ea0cad2 100644 --- a/src/include/milvus/Status.h +++ b/src/include/milvus/Status.h @@ -35,6 +35,7 @@ enum class StatusCode { UNKNOWN_ERROR = 1, NOT_SUPPORTED, NOT_CONNECTED, + CLIENT_BUSY, // function error section INVALID_ARGUMENT = 1000, diff --git a/src/include/milvus/types/ConnectParam.h b/src/include/milvus/types/ConnectParam.h index 567b1b0b..35819960 100644 --- a/src/include/milvus/types/ConnectParam.h +++ b/src/include/milvus/types/ConnectParam.h @@ -20,6 +20,7 @@ #include #include "milvus/Export.h" +#include "milvus/types/TelemetryConfig.h" namespace milvus { @@ -338,6 +339,16 @@ class MILVUS_SDK_API ConnectParam { ConnectParam& WithDbName(const std::string& db_name); + /** Client telemetry and command configuration. */ + const TelemetryConfig& + Telemetry() const; + + void + SetTelemetryConfig(const TelemetryConfig& config); + + ConnectParam& + WithTelemetryConfig(const TelemetryConfig& config); + private: std::string uri_ = "http://localhost:19530"; @@ -357,6 +368,7 @@ class MILVUS_SDK_API ConnectParam { std::string username_; std::string token_; std::string db_name_; + TelemetryConfig telemetry_config_; }; } // namespace milvus diff --git a/src/include/milvus/types/TelemetryConfig.h b/src/include/milvus/types/TelemetryConfig.h new file mode 100644 index 00000000..ce0082d4 --- /dev/null +++ b/src/include/milvus/types/TelemetryConfig.h @@ -0,0 +1,36 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#pragma once + +#include +#include +#include + +#include "milvus/Export.h" + +namespace milvus { + +/** Client metrics, heartbeat, and server-pushed command configuration. */ +struct MILVUS_SDK_API TelemetryConfig { + bool enabled{true}; + // Milliseconds between heartbeats, and therefore the metrics window: each heartbeat + // carries the operations since the last one. The coordinator answers a telemetry query + // from the window before the newest, so what a caller reads is between one and two + // intervals old. + uint64_t heartbeat_interval_ms{10000}; + double sampling_rate{1.0}; + size_t error_max_count{100}; + + /** Optional stable identity. A random UUID is used when empty. */ + std::string client_id; +}; + +} // namespace milvus diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index f5269edf..929bd179 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -76,6 +76,11 @@ add_executable(testing-it ${it_files}) target_compile_options(testing-it PRIVATE $<$:/bigobj>) link_milvus_test(testing-it) +if (CMAKE_SYSTEM_NAME MATCHES "(Linux|Darwin)") + add_executable(testing-telemetry-e2e e2e/ClientTelemetryE2E.cpp) + target_link_libraries(testing-telemetry-e2e PRIVATE milvus_sdk) +endif() + # st only available under linux/macos if (CMAKE_SYSTEM_NAME MATCHES "(Linux|Darwin)") set(ST_DIR "${CMAKE_CURRENT_SOURCE_DIR}/st") diff --git a/test/e2e/ClientTelemetryE2E.cpp b/test/e2e/ClientTelemetryE2E.cpp new file mode 100644 index 00000000..e004a94b --- /dev/null +++ b/test/e2e/ClientTelemetryE2E.cpp @@ -0,0 +1,375 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "milvus/ClientRequestContext.h" +#include "milvus/ClientTelemetry.h" +#include "milvus/MilvusClient.h" +#include "milvus/MilvusClientV2.h" +#include "milvus/thirdparty/nlohmann/json.hpp" + +namespace { + +using Json = nlohmann::json; + +const std::string kTelemetryHost = std::getenv("MILVUS_TELEMETRY_HOST") == nullptr + ? "127.0.0.1" + : std::getenv("MILVUS_TELEMETRY_HOST"); +const uint16_t kTelemetryPort = static_cast( + std::getenv("MILVUS_TELEMETRY_PORT") == nullptr ? 9091 : std::stoi(std::getenv("MILVUS_TELEMETRY_PORT"))); +const std::string kTelemetryBase = "/api/v1/_telemetry"; + +void +Require(bool condition, const std::string& message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +std::string +ToLower(std::string value) { + for (auto& character : value) { + if (character >= 'A' && character <= 'Z') { + character = static_cast(character - 'A' + 'a'); + } + } + return value; +} + +std::string +DecodeChunked(const std::string& body) { + std::string decoded; + size_t offset = 0; + while (offset < body.size()) { + auto line_end = body.find("\r\n", offset); + Require(line_end != std::string::npos, "Malformed chunked HTTP response"); + auto size_text = body.substr(offset, line_end - offset); + auto extension = size_text.find(';'); + if (extension != std::string::npos) { + size_text.resize(extension); + } + size_t chunk_size = std::stoul(size_text, nullptr, 16); + offset = line_end + 2; + if (chunk_size == 0) { + break; + } + Require(offset + chunk_size <= body.size(), "Truncated chunked HTTP response"); + decoded.append(body, offset, chunk_size); + offset += chunk_size + 2; + } + return decoded; +} + +std::string +HttpRequest(const std::string& method, const std::string& path, const std::string& body = "") { + addrinfo hints{}; + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + addrinfo* addresses = nullptr; + auto port = std::to_string(kTelemetryPort); + int lookup = getaddrinfo(kTelemetryHost.c_str(), port.c_str(), &hints, &addresses); + Require(lookup == 0, std::string("getaddrinfo failed: ") + gai_strerror(lookup)); + + int socket_fd = -1; + for (auto* address = addresses; address != nullptr; address = address->ai_next) { + socket_fd = socket(address->ai_family, address->ai_socktype, address->ai_protocol); + if (socket_fd >= 0 && connect(socket_fd, address->ai_addr, address->ai_addrlen) == 0) { + break; + } + if (socket_fd >= 0) { + close(socket_fd); + } + socket_fd = -1; + } + freeaddrinfo(addresses); + Require(socket_fd >= 0, "Unable to connect to Milvus telemetry HTTP endpoint"); + + std::ostringstream request; + request << method << " " << path << " HTTP/1.1\r\n" + << "Host: " << kTelemetryHost << ':' << kTelemetryPort << "\r\n" + << "Accept: application/json\r\n" + << "Connection: close\r\n"; + if (!body.empty()) { + request << "Content-Type: application/json\r\n" << "Content-Length: " << body.size() << "\r\n"; + } + request << "\r\n" << body; + auto wire = request.str(); + size_t sent = 0; + while (sent < wire.size()) { + auto count = send(socket_fd, wire.data() + sent, wire.size() - sent, 0); + if (count <= 0) { + close(socket_fd); + throw std::runtime_error("Failed to send telemetry HTTP request"); + } + sent += static_cast(count); + } + + std::string response; + char buffer[8192]; + while (true) { + auto count = recv(socket_fd, buffer, sizeof(buffer), 0); + if (count < 0) { + close(socket_fd); + throw std::runtime_error("Failed to read telemetry HTTP response"); + } + if (count == 0) { + break; + } + response.append(buffer, static_cast(count)); + } + close(socket_fd); + + auto headers_end = response.find("\r\n\r\n"); + Require(headers_end != std::string::npos, "Malformed telemetry HTTP response"); + auto headers = response.substr(0, headers_end); + auto status_end = headers.find("\r\n"); + auto status_line = headers.substr(0, status_end); + Require(status_line.find(" 200 ") != std::string::npos, "Telemetry HTTP request failed: " + status_line); + auto response_body = response.substr(headers_end + 4); + if (ToLower(headers).find("transfer-encoding: chunked") != std::string::npos) { + response_body = DecodeChunked(response_body); + } + return response_body; +} + +Json +ClientState(const std::string& client_id) { + auto response = Json::parse(HttpRequest( + "GET", kTelemetryBase + "/clients?client_id=" + client_id + "&include_metrics=true")); + auto clients = response.value("clients", Json::array()); + return clients.empty() ? Json() : clients.at(0); +} + +Json +WaitFor(const std::string& label, const std::string& client_id, const std::function& predicate) { + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(15); + Json last; + while (std::chrono::steady_clock::now() < deadline) { + last = ClientState(client_id); + if (!last.is_null() && predicate(last)) { + std::cout << "PASS " << label << std::endl; + return last; + } + std::this_thread::sleep_for(std::chrono::milliseconds(250)); + } + throw std::runtime_error("Timed out waiting for " + label + "; last=" + last.dump()); +} + +std::string +PushCommand(const std::string& client_id, const std::string& command_type, const Json& payload, + bool persistent = false) { + Json request = {{"command_type", command_type}, + {"target_client_id", client_id}, + {"payload", payload}, + {"ttl_seconds", 30}, + {"persistent", persistent}}; + auto response = Json::parse(HttpRequest("POST", kTelemetryBase + "/commands", request.dump())); + return response.at("command_id").get(); +} + +Json +FindReply(const Json& state, const std::string& command_id) { + auto replies = state.value("command_replies", Json::array()); + for (const auto& reply : replies) { + if (reply.value("command_id", "") == command_id) { + return reply; + } + } + return Json(); +} + +Json +WaitForReply(const std::string& client_id, const std::string& command_id) { + auto state = WaitFor("command reply " + command_id, client_id, + [&command_id](const Json& candidate) { return !FindReply(candidate, command_id).is_null(); }); + return FindReply(state, command_id); +} + +bool +HasMetric(const Json& state, const std::string& operation, const std::string& counter, int64_t minimum, + const std::string& collection = "") { + auto metrics = state.value("metrics", Json::array()); + for (const auto& metric : metrics) { + if (metric.value("operation", "") != operation) { + continue; + } + auto global = metric.value("global", Json::object()); + if (global.value(counter, static_cast(0)) < minimum) { + continue; + } + if (collection.empty()) { + return true; + } + auto collections = metric.value("collection_metrics", Json::object()); + return collections.find(collection) != collections.end(); + } + return false; +} + +int +Run() { + auto milvus_host = std::getenv("MILVUS_HOST") == nullptr ? "127.0.0.1" : std::getenv("MILVUS_HOST"); + auto milvus_port = static_cast( + std::getenv("MILVUS_PORT") == nullptr ? 19530 : std::stoi(std::getenv("MILVUS_PORT"))); + + { + auto legacy_client_id = "e2e-cpp-legacy-" + milvus::ClientRequestContext::NewRequestId(); + milvus::TelemetryConfig legacy_telemetry; + legacy_telemetry.client_id = legacy_client_id; + legacy_telemetry.heartbeat_interval_ms = 500; + milvus::ConnectParam legacy_param{milvus_host, milvus_port}; + legacy_param.WithTelemetryConfig(legacy_telemetry); + auto legacy_client = milvus::MilvusClient::Create(); + auto status = legacy_client->Connect(legacy_param); + Require(status.IsOk(), "Legacy telemetry connect failed: " + status.Message()); + Require(legacy_client->GetTelemetry()->ClientId() == legacy_client_id, + "Unexpected legacy telemetry client ID"); + WaitFor("legacy client registration", legacy_client_id, + [](const Json& state) { return state.value("status", "") == "active"; }); + status = legacy_client->Disconnect(); + Require(status.IsOk(), "Legacy telemetry disconnect failed: " + status.Message()); + } + + { + milvus::ConnectParam default_param{milvus_host, milvus_port}; + auto default_client = milvus::MilvusClientV2::Create(); + auto status = default_client->Connect(default_param); + Require(status.IsOk(), "Default telemetry connect failed: " + status.Message()); + auto default_manager = default_client->GetTelemetry(); + Require(default_manager != nullptr, "Default telemetry manager is missing"); + Require(default_manager->Config().client_id.empty(), "Generated client ID was marked as stable"); + auto default_client_id = default_manager->ClientId(); + WaitFor("default client registration", default_client_id, + [](const Json& state) { return state.value("status", "") == "active"; }); + + status = default_client->UseDatabase("default"); + Require(status.IsOk(), "Default telemetry reconnect failed: " + status.Message()); + Require(default_client->GetTelemetry() == default_manager, "Reconnect replaced the telemetry manager"); + Require(default_manager->ClientId() == default_client_id, "Reconnect changed the runtime client ID"); + Require(default_manager->Config().client_id.empty(), "Reconnect marked the runtime client ID as stable"); + WaitFor("default client reconnect", default_client_id, + [](const Json& state) { return state.value("status", "") == "active"; }); + status = default_client->Disconnect(); + Require(status.IsOk(), "Default telemetry disconnect failed: " + status.Message()); + } + + auto trace_suffix = milvus::ClientRequestContext::NewRequestId(); + auto client_id = "e2e-cpp-" + trace_suffix; + milvus::TelemetryConfig telemetry_config; + telemetry_config.client_id = client_id; + telemetry_config.heartbeat_interval_ms = 500; + telemetry_config.sampling_rate = 1.0; + + milvus::ConnectParam connect_param{milvus_host, milvus_port}; + connect_param.WithTelemetryConfig(telemetry_config); + + auto client = milvus::MilvusClientV2::Create(); + auto status = client->Connect(connect_param); + Require(status.IsOk(), "Milvus connect failed: " + status.Message()); + try { + auto manager = client->GetTelemetry(); + Require(manager != nullptr, "Telemetry manager is missing"); + Require(manager->ClientId() == client_id, "Unexpected telemetry client ID"); + WaitFor("client registration", client_id, + [](const Json& state) { return state.value("status", "") == "active"; }); + + milvus::RunAnalyzerRequest analyzer_request; + analyzer_request.AddText("hello milvus telemetry") + .WithAnalyzerParams({{"type", "standard"}}) + .WithDetail(true); + milvus::RunAnalyzerResponse analyzer_response; + status = client->RunAnalyzer(analyzer_request, analyzer_response); + Require(status.IsOk(), "RunAnalyzer failed: " + status.Message()); + const auto& tokens = analyzer_response.Results().at(0).Tokens(); + Require(tokens.size() == 3 && tokens[0].token_ == "hello" && tokens[1].token_ == "milvus" && + tokens[2].token_ == "telemetry", + "Unexpected RunAnalyzer tokens"); + WaitFor("RunAnalyzer metric", client_id, + [](const Json& state) { return HasMetric(state, "RunAnalyzer", "success_count", 1); }); + + auto collection_command = + PushCommand(client_id, "collection_metrics", {{"collections", {"*"}}, {"enabled", true}}); + Require(WaitForReply(client_id, collection_command).value("success", false), + "collection_metrics command failed"); + + auto request_id = milvus::ClientRequestContext::NewRequestId(); + { + milvus::ScopedClientRequestId scoped(request_id); + milvus::QueryRequest query_request; + query_request.WithCollectionName("telemetry_e2e_missing").WithFilter("id > 0"); + milvus::QueryResponse query_response; + status = client->Query(query_request, query_response); + } + Require(!status.IsOk(), "Query against missing collection unexpectedly succeeded"); + WaitFor("failed Query collection metric", client_id, [](const Json& state) { + return HasMetric(state, "Query", "error_count", 1, "telemetry_e2e_missing"); + }); + + auto errors_reply = WaitForReply(client_id, PushCommand(client_id, "show_errors", {{"max_count", 10}})); + Require(errors_reply.value("success", false), "show_errors command failed"); + auto errors = Json::parse(errors_reply.at("payload").get()); + bool trace_found = false; + for (const auto& error : errors) { + trace_found = trace_found || (error.value("operation", "") == "Query" && + error.value("request_id", "") == request_id); + } + Require(trace_found, "show_errors did not include the request ID"); + std::cout << "PASS request-id in show_errors" << std::endl; + + Json config_payload = {{"sampling_rate", 0.75}, {"heartbeat_interval_ms", 600}}; + auto config_command = PushCommand(client_id, "push_config", config_payload, true); + Require(WaitForReply(client_id, config_command).value("success", false), "push_config command failed"); + milvus::TelemetryCommand expected_config{config_command, "push_config", config_payload.dump(), 0, true, ""}; + Require(manager->ConfigHash() == milvus::ClientTelemetryManager::CalculateConfigHash({expected_config}), + "Persistent config hash mismatch"); + Require(manager->LastCommandTimestamp() > 0, "Command timestamp was not updated"); + + auto get_config_reply = WaitForReply(client_id, PushCommand(client_id, "get_config", Json::object())); + Require(get_config_reply.value("success", false), "get_config command failed"); + auto user_config = Json::parse(get_config_reply.at("payload").get()).at("user_config"); + Require(user_config.value("telemetry_sampling_rate", 0.0) == 0.75, "Sampling rate was not applied"); + Require(user_config.value("telemetry_heartbeat_interval_ms", 0) == 600, + "Heartbeat interval was not applied"); + Require(user_config.value("all_collections_enabled", false), "Collection metrics wildcard was not applied"); + std::cout << "CPP_E2E_OK " << client_id << std::endl; + } catch (...) { + client->Disconnect(); + throw; + } + status = client->Disconnect(); + Require(status.IsOk(), "Milvus disconnect failed: " + status.Message()); + return 0; +} + +} // namespace + +int +main() { + try { + return Run(); + } catch (const std::exception& exception) { + std::cerr << "CPP_E2E_FAILED: " << exception.what() << std::endl; + return 1; + } +} diff --git a/test/it/v1/TestConnection.cpp b/test/it/v1/TestConnection.cpp index 9a500305..cb622cff 100644 --- a/test/it/v1/TestConnection.cpp +++ b/test/it/v1/TestConnection.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -65,6 +66,11 @@ TEST_F(UnconnectMilvusMockedTest, ResolveDatabaseNamePreservesEmptyForRpc) { ASSERT_TRUE(empty_db_handler.Connect(milvus::ConnectParam{"127.0.0.1", server_.ListenPort()}).IsOk()); EXPECT_EQ(empty_db_handler.CurrentDbName(""), ""); EXPECT_EQ(empty_db_handler.CurrentDbName("request_db"), "request_db"); + auto empty_db_telemetry = empty_db_handler.GetTelemetry(); + empty_db_telemetry->ProcessCommands({{"empty-db", "get_config", "", 1, false, ""}}); + auto empty_db_replies = empty_db_telemetry->PendingCommandReplies(); + ASSERT_FALSE(empty_db_replies.empty()); + EXPECT_EQ(nlohmann::json::parse(empty_db_replies.back().payload)["user_config"]["db_name"], ""); milvus::ConnectParam explicit_db_param{"127.0.0.1", server_.ListenPort()}; explicit_db_param.SetDbName("connected_db"); @@ -72,6 +78,11 @@ TEST_F(UnconnectMilvusMockedTest, ResolveDatabaseNamePreservesEmptyForRpc) { ASSERT_TRUE(explicit_db_handler.Connect(explicit_db_param).IsOk()); EXPECT_EQ(explicit_db_handler.CurrentDbName(""), "connected_db"); EXPECT_EQ(explicit_db_handler.CurrentDbName("request_db"), "request_db"); + auto explicit_db_telemetry = explicit_db_handler.GetTelemetry(); + explicit_db_telemetry->ProcessCommands({{"explicit-db", "get_config", "", 1, false, ""}}); + auto explicit_db_replies = explicit_db_telemetry->PendingCommandReplies(); + ASSERT_FALSE(explicit_db_replies.empty()); + EXPECT_EQ(nlohmann::json::parse(explicit_db_replies.back().payload)["user_config"]["db_name"], "connected_db"); } TEST_F(UnconnectMilvusMockedTest, ConnectServerRejected) { @@ -145,7 +156,7 @@ TEST_F(UnconnectMilvusMockedTest, FailedReconnectPreservesExistingConnection) { EXPECT_TRUE(has_collection); } -TEST_F(UnconnectMilvusMockedTest, ConnectionMutationWaitsForConnectAndAppliesToNewConnection) { +TEST_F(UnconnectMilvusMockedTest, ConnectionMutationFailsFastDuringConnectAndCanBeRetried) { milvus::ConnectionHandler handler; milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; @@ -195,7 +206,7 @@ TEST_F(UnconnectMilvusMockedTest, ConnectionMutationWaitsForConnectAndAppliesToN { std::unique_lock lock(setter_mutex); setter_cv.wait(lock, [&] { return setter_started; }); - EXPECT_FALSE(setter_cv.wait_for(lock, std::chrono::milliseconds(50), [&] { return setter_finished; })); + EXPECT_TRUE(setter_cv.wait_for(lock, std::chrono::seconds(1), [&] { return setter_finished; })); } { @@ -207,7 +218,8 @@ TEST_F(UnconnectMilvusMockedTest, ConnectionMutationWaitsForConnectAndAppliesToN connect_thread.join(); setter_thread.join(); EXPECT_TRUE(connect_status.IsOk()); - EXPECT_TRUE(setter_status.IsOk()); + EXPECT_EQ(setter_status.Code(), StatusCode::CLIENT_BUSY); + EXPECT_TRUE(handler.SetRpcDeadlineMs(1234).IsOk()); EXPECT_EQ(handler.GetRpcDeadlineMs(), 1234); } diff --git a/test/it/v1/TestUseDatabase.cpp b/test/it/v1/TestUseDatabase.cpp index 15f9eccc..95784773 100644 --- a/test/it/v1/TestUseDatabase.cpp +++ b/test/it/v1/TestUseDatabase.cpp @@ -41,9 +41,12 @@ TEST_F(UnconnectMilvusMockedTest, UseDatabase) { status = client_->CurrentUsedDatabase(db_name); EXPECT_TRUE(status.IsOk()); EXPECT_EQ(db_name, "AAA"); + auto telemetry = client_->GetTelemetry(); + ASSERT_NE(telemetry, nullptr); status = client_->UseDatabase("BBB"); EXPECT_TRUE(status.IsOk()); + EXPECT_EQ(client_->GetTelemetry(), telemetry); status = client_->CurrentUsedDatabase(db_name); EXPECT_TRUE(status.IsOk()); diff --git a/test/ut/TestClientRequestContextV1Aggregate.cpp b/test/ut/TestClientRequestContextV1Aggregate.cpp new file mode 100644 index 00000000..7789d612 --- /dev/null +++ b/test/ut/TestClientRequestContextV1Aggregate.cpp @@ -0,0 +1,24 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "milvus/MilvusClient.h" + +TEST(ClientRequestContextAggregateTest, V1HeaderExportsRequestContext) { + milvus::ScopedClientRequestId scoped("v1-aggregate"); + EXPECT_EQ(milvus::ClientRequestContext::Get(), "v1-aggregate"); +} diff --git a/test/ut/TestClientRequestContextV2Aggregate.cpp b/test/ut/TestClientRequestContextV2Aggregate.cpp new file mode 100644 index 00000000..0fbe409c --- /dev/null +++ b/test/ut/TestClientRequestContextV2Aggregate.cpp @@ -0,0 +1,24 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "milvus/MilvusClientV2.h" + +TEST(ClientRequestContextAggregateTest, V2HeaderExportsRequestContext) { + milvus::ScopedClientRequestId scoped("v2-aggregate"); + EXPECT_EQ(milvus::ClientRequestContext::Get(), "v2-aggregate"); +} diff --git a/test/ut/TestClientTelemetry.cpp b/test/ut/TestClientTelemetry.cpp new file mode 100644 index 00000000..d0727b4f --- /dev/null +++ b/test/ut/TestClientTelemetry.cpp @@ -0,0 +1,1148 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "MilvusConnection.h" +#include "milvus.grpc.pb.h" +#include "milvus.pb.h" +#include "milvus/ClientRequestContext.h" +#include "milvus/ClientTelemetry.h" +#include "milvus/MilvusClient.h" +#include "milvus/MilvusClientV2.h" + +namespace { + +std::string +Rfc3339(std::chrono::system_clock::time_point value) { + auto raw = std::chrono::system_clock::to_time_t(value); + std::tm utc{}; +#ifdef _WIN32 + gmtime_s(&utc, &raw); +#else + gmtime_r(&raw, &utc); +#endif + std::ostringstream stream; + stream << std::put_time(&utc, "%Y-%m-%dT%H:%M:%SZ"); + return stream.str(); +} + +class ReconnectTelemetryService final : public milvus::proto::milvus::ClientTelemetryService::Service { + public: + grpc::Status + ClientHeartbeat(grpc::ServerContext*, const milvus::proto::milvus::ClientHeartbeatRequest*, + milvus::proto::milvus::ClientHeartbeatResponse* response) override { + ++heartbeats; + auto* command = response->add_commands(); + command->set_command_id("reconnect"); + command->set_command_type("reconnect"); + command->set_create_time(1); + return grpc::Status::OK; + } + + std::atomic heartbeats{0}; +}; + +class ControlPlaneTelemetryService final : public milvus::proto::milvus::ClientTelemetryService::Service { + public: + grpc::Status + ClientHeartbeat(grpc::ServerContext*, const milvus::proto::milvus::ClientHeartbeatRequest* request, + milvus::proto::milvus::ClientHeartbeatResponse* response) override { + std::unique_lock lock(mutex_); + requests_.push_back(*request); + const auto heartbeat = requests_.size(); + condition_.notify_all(); + if (heartbeat == 1) { + auto* command = response->add_commands(); + command->set_command_id("disable"); + command->set_command_type("push_config"); + command->set_payload(R"({"enabled":false})"); + command->set_create_time(1); + command->set_persistent(true); + } else if (heartbeat == 2) { + condition_.wait(lock, [this]() { return allow_enable_; }); + auto* command = response->add_commands(); + command->set_command_id("enable"); + command->set_command_type("push_config"); + command->set_payload(R"({"enabled":true})"); + command->set_create_time(2); + command->set_persistent(true); + } + return grpc::Status::OK; + } + + bool + WaitForHeartbeats(size_t count) { + std::unique_lock lock(mutex_); + return condition_.wait_for(lock, std::chrono::seconds(2), + [this, count]() { return requests_.size() >= count; }); + } + + void + AllowEnable() { + { + std::lock_guard lock(mutex_); + allow_enable_ = true; + } + condition_.notify_all(); + } + + std::vector + Requests() const { + std::lock_guard lock(mutex_); + return requests_; + } + + size_t + RequestCount() const { + std::lock_guard lock(mutex_); + return requests_.size(); + } + + private: + mutable std::mutex mutex_; + std::condition_variable condition_; + bool allow_enable_{false}; + std::vector requests_; +}; + +class LifecycleTelemetryService final : public milvus::proto::milvus::MilvusService::Service, + public milvus::proto::milvus::ClientTelemetryService::Service { + public: + explicit LifecycleTelemetryService(std::string command_type) : command_type_(std::move(command_type)) { + } + + grpc::Status + Connect(grpc::ServerContext*, const milvus::proto::milvus::ConnectRequest*, + milvus::proto::milvus::ConnectResponse* response) override { + { + std::lock_guard lock(connect_mutex_); + ++connects; + if (fail_next_connect_) { + fail_next_connect_ = false; + response->mutable_status()->set_code(1); + response->mutable_status()->set_reason("rejected connect"); + } + } + connect_condition_.notify_all(); + return grpc::Status::OK; + } + + grpc::Status + HasCollection(grpc::ServerContext*, const milvus::proto::milvus::HasCollectionRequest*, + milvus::proto::milvus::BoolResponse*) override { + ++has_collections; + return grpc::Status::OK; + } + + grpc::Status + ClientHeartbeat(grpc::ServerContext*, const milvus::proto::milvus::ClientHeartbeatRequest* request, + milvus::proto::milvus::ClientHeartbeatResponse* response) override { + ++heartbeats; + const auto database = request->client_info().reserved().find("db_name"); + if (database != request->client_info().reserved().end() && database->second == "secondary") { + saw_secondary_database = true; + } + for (const auto& reply : request->command_replies()) { + if (reply.command_id() == command_type_ && !reply.success()) { + saw_failed_command_reply = true; + } + } + if (commands_enabled.load() && !command_sent.exchange(true)) { + auto* command = response->add_commands(); + command->set_command_id(command_type_); + command->set_command_type(command_type_); + command->set_create_time(1); + } + return grpc::Status::OK; + } + + std::atomic connects{0}; + std::atomic heartbeats{0}; + std::atomic has_collections{0}; + std::atomic saw_secondary_database{false}; + std::atomic saw_failed_command_reply{false}; + + void + EnableCommands() { + commands_enabled = true; + } + + void + FailNextConnect() { + std::lock_guard lock(connect_mutex_); + fail_next_connect_ = true; + } + + bool + WaitForConnects(int count) { + std::unique_lock lock(connect_mutex_); + return connect_condition_.wait_for(lock, std::chrono::seconds(2), + [this, count]() { return connects.load() >= count; }); + } + + private: + std::string command_type_; + std::atomic commands_enabled{false}; + std::atomic command_sent{false}; + std::mutex connect_mutex_; + std::condition_variable connect_condition_; + bool fail_next_connect_{false}; +}; + +class RequestMetadataService final : public milvus::proto::milvus::MilvusService::Service { + public: + grpc::Status + Connect(grpc::ServerContext*, const milvus::proto::milvus::ConnectRequest*, + milvus::proto::milvus::ConnectResponse*) override { + return grpc::Status::OK; + } + + grpc::Status + HasCollection(grpc::ServerContext* context, const milvus::proto::milvus::HasCollectionRequest*, + milvus::proto::milvus::BoolResponse*) override { + const auto& metadata = context->client_metadata(); + const auto request_id = metadata.find("client_request_id"); + std::lock_guard lock(mutex_); + request_ids_.emplace_back(request_id != metadata.end(), + request_id == metadata.end() + ? std::string{} + : std::string(request_id->second.data(), request_id->second.size())); + return grpc::Status::OK; + } + + std::vector> + RequestIds() const { + std::lock_guard lock(mutex_); + return request_ids_; + } + + private: + mutable std::mutex mutex_; + std::vector> request_ids_; +}; + +class DestructionSignal final { + public: + explicit DestructionSignal(std::shared_ptr> signal) : signal_(std::move(signal)) { + } + + ~DestructionSignal() { + signal_->set_value(); + } + + private: + std::shared_ptr> signal_; +}; + +std::unique_ptr +StartLifecycleServer(LifecycleTelemetryService& service, int& port) { + grpc::ServerBuilder builder; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(static_cast(&service)); + builder.RegisterService(static_cast(&service)); + return builder.BuildAndStart(); +} + +milvus::ConnectParam +TelemetryConnectParam(int port) { + milvus::ConnectParam param("http://127.0.0.1:" + std::to_string(port)); + milvus::TelemetryConfig config; + config.heartbeat_interval_ms = 1; + param.SetTelemetryConfig(config); + return param; +} + +} // namespace + +TEST(ClientTelemetryTest, MatchesCrossSdkConfigHashVector) { + std::vector commands = { + {"cfg-b", "push_config", "{\"sampling_rate\":0.5}", 0, true, ""}, + {"cfg-a", "push_config", "{\"heartbeat_interval_ms\":5000}", 0, true, ""}, + }; + EXPECT_EQ(milvus::ClientTelemetryManager::CalculateConfigHash(commands), "a271ff0bb1941777"); +} + +TEST(ClientTelemetryTest, RuntimeClientIdDoesNotBecomeStableConfiguration) { + milvus::TelemetryConfig config; + milvus::ClientTelemetryManager manager(config, "runtime-client-id"); + + EXPECT_EQ(manager.ClientId(), "runtime-client-id"); + EXPECT_TRUE(manager.Config().client_id.empty()); +} + +TEST(ClientTelemetryTest, AppliesCommandsAndDeduplicatesIds) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + int calls = 0; + manager.RegisterCommandHandler("custom", [&calls](const milvus::TelemetryCommand& command) { + ++calls; + return milvus::TelemetryCommandReply{command.command_id, true, "", ""}; + }); + + manager.ProcessCommands({ + {"config", "push_config", "{\"heartbeat_interval_ms\":5000,\"sampling_rate\":0.25}", 1, true, ""}, + {"custom", "custom", "", 2, false, ""}, + }); + manager.ProcessCommands({{"custom", "custom", "", 2, false, ""}}); + manager.ProcessCommands({{"custom", "custom", "", 2, false, ""}}); + + EXPECT_EQ(manager.Config().heartbeat_interval_ms, 5000U); + EXPECT_DOUBLE_EQ(manager.Config().sampling_rate, 0.25); + EXPECT_EQ(manager.LastCommandTimestamp(), 2); + EXPECT_FALSE(manager.ConfigHash().empty()); + EXPECT_EQ(calls, 1); +} + +TEST(ClientTelemetryTest, RetainsMoreThanOneHundredTwentySnapshotsWithinOneHour) { + milvus::TelemetryConfig config; + milvus::ClientTelemetryManager manager(config); + milvus::proto::milvus::SearchRequest request; + + constexpr size_t expected_snapshots = 121; + for (size_t index = 0; index < expected_snapshots; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now(), true, ""); + manager.Start(); + for (int retry = 0; retry < 1000 && manager.MetricsSnapshots().size() <= index; ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + ASSERT_GT(manager.MetricsSnapshots().size(), index); + } + + EXPECT_EQ(manager.MetricsSnapshots().size(), expected_snapshots); +} + +TEST(ClientTelemetryTest, RetainsOneSecondHeartbeatWindowAndSkipsEmptyIntervals) { + milvus::TelemetryConfig config; + config.heartbeat_interval_ms = 1; + milvus::ClientTelemetryManager manager(config); + milvus::proto::milvus::SearchRequest request; + + manager.Start(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + manager.Stop(); + EXPECT_TRUE(manager.MetricsSnapshots().empty()); + + constexpr size_t generated_snapshots = 4104; + constexpr size_t retained_snapshots = 4096; + for (size_t index = 0; index < generated_snapshots; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now(), true, ""); + manager.Start(); + manager.Stop(); + } + + EXPECT_EQ(manager.MetricsSnapshots().size(), retained_snapshots); +} + +TEST(ClientTelemetryTest, AggregatesP99FromRetainedLatencySamples) { + milvus::TelemetryConfig config; + milvus::ClientTelemetryManager manager(config); + milvus::proto::milvus::SearchRequest request; + + for (int index = 0; index < 100; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now() - std::chrono::milliseconds(1), + true, ""); + } + manager.Start(); + for (int retry = 0; retry < 1000 && manager.MetricsSnapshots().empty(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + + for (int index = 0; index < 100; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now() - std::chrono::milliseconds(100), + true, ""); + } + manager.Start(); + for (int retry = 0; retry < 1000 && manager.MetricsSnapshots().size() < 2; ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + + const auto now = std::chrono::system_clock::now(); + const auto payload = nlohmann::json{ + {"start_time", Rfc3339(now - std::chrono::minutes(1))}, + {"end_time", Rfc3339(now + std::chrono::minutes(1))}, + {"detail", false}}.dump(); + manager.ProcessCommands({{"history", "show_latency_history", payload, 1, false, ""}}); + + const auto replies = manager.PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto response = nlohmann::json::parse(replies.back().payload); + EXPECT_GT(response["aggregated"]["metrics"]["Search"]["p99_latency_ms"].get(), 90.0); +} + +TEST(ClientTelemetryTest, CompressedQuantileHistoryPreservesEndpointsAndSlowTail) { + milvus::TelemetryConfig config; + milvus::ClientTelemetryManager manager(config); + milvus::proto::milvus::SearchRequest request; + + // The exact per-window p99 is still in the fast group (indices 0..990), while + // the 128-point history compression must include the slow endpoint/tail that + // starts at evenly-spaced source index 991. + for (int index = 0; index < 991; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now() - std::chrono::milliseconds(1), + true, ""); + } + for (int index = 0; index < 9; ++index) { + manager.RecordOperation("Search", request, std::chrono::steady_clock::now() - std::chrono::milliseconds(250), + true, ""); + } + manager.Start(); + for (int retry = 0; retry < 1000 && manager.MetricsSnapshots().empty(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + + const auto snapshots = manager.MetricsSnapshots(); + ASSERT_EQ(snapshots.size(), 1U); + ASSERT_EQ(snapshots.front().metrics.size(), 1U); + EXPECT_LT(snapshots.front().metrics.front().global.p99_latency_ms, 50.0); + EXPECT_GT(snapshots.front().metrics.front().global.max_latency_ms, 200.0); + + const auto now = std::chrono::system_clock::now(); + const auto payload = nlohmann::json{ + {"start_time", Rfc3339(now - std::chrono::minutes(1))}, + {"end_time", Rfc3339(now + std::chrono::minutes(1))}, + {"detail", false}}.dump(); + manager.ProcessCommands({{"compressed-history", "show_latency_history", payload, 1, false, ""}}); + + const auto replies = manager.PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto response = nlohmann::json::parse(replies.back().payload); + const auto metric = response["aggregated"]["metrics"]["Search"]; + EXPECT_GT(metric["p99_latency_ms"].get(), 200.0); + EXPECT_GT(metric["max_latency_ms"].get(), 200.0); +} + +TEST(ClientTelemetryTest, ReusesHeartbeatWorkerWhenReconnectRunsInCommandHandler) { + ReconnectTelemetryService service; + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + + std::atomic reconnects{0}; + { + milvus::TelemetryConfig config; + config.heartbeat_interval_ms = 1; + milvus::ClientTelemetryManager manager(config); + manager.AttachChannel( + grpc::CreateChannel("127.0.0.1:" + std::to_string(port), grpc::InsecureChannelCredentials()), "", "", "", + "", ""); + manager.RegisterCommandHandler("reconnect", [&manager, &reconnects](const milvus::TelemetryCommand& command) { + ++reconnects; + manager.Stop(); + manager.Start(); + return milvus::TelemetryCommandReply{command.command_id, true, "", ""}; + }); + + manager.Start(); + for (int retry = 0; retry < 2000 && service.heartbeats.load() < 2; ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + } + server->Shutdown(); + + EXPECT_GE(service.heartbeats.load(), 2); + EXPECT_EQ(reconnects.load(), 1); +} + +TEST(ClientTelemetryTest, ExternalStopWinsRaceWithHeartbeatSelfRestart) { + ReconnectTelemetryService service; + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + + milvus::TelemetryConfig config; + config.heartbeat_interval_ms = 1; + milvus::ClientTelemetryManager manager(config); + manager.AttachChannel(grpc::CreateChannel("127.0.0.1:" + std::to_string(port), grpc::InsecureChannelCredentials()), + "", "", "", "", ""); + + std::atomic handler_entered{false}; + std::atomic allow_self_restart{false}; + std::atomic self_restart_ready{true}; + manager.RegisterCommandHandler("reconnect", [&](const milvus::TelemetryCommand& command) { + handler_entered = true; + while (!allow_self_restart.load()) { + std::this_thread::yield(); + } + manager.Stop(); + manager.Start(); + self_restart_ready = manager.IsReady(); + return milvus::TelemetryCommandReply{command.command_id, true, "", ""}; + }); + + manager.Start(); + for (int retry = 0; retry < 2000 && !handler_entered.load(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_TRUE(handler_entered.load()); + ASSERT_TRUE(manager.IsReady()); + + std::promise external_stop_finished; + auto external_stop_future = external_stop_finished.get_future(); + std::thread external_stopper([&]() { + manager.Stop(); + external_stop_finished.set_value(); + }); + + // ready=false is written before Stop() releases the manager mutex to join, + // so observing it makes the intended race ordering deterministic. + for (int retry = 0; retry < 2000 && manager.IsReady(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_FALSE(manager.IsReady()); + EXPECT_EQ(external_stop_future.wait_for(std::chrono::milliseconds(10)), std::future_status::timeout); + + allow_self_restart = true; + EXPECT_EQ(external_stop_future.wait_for(std::chrono::seconds(2)), std::future_status::ready); + external_stopper.join(); + server->Shutdown(); + + EXPECT_FALSE(self_restart_ready.load()); + EXPECT_FALSE(manager.IsReady()); +} + +TEST(ClientTelemetryTest, DisabledTelemetryKeepsControlPlaneHeartbeatAlive) { + ControlPlaneTelemetryService service; + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + + milvus::TelemetryConfig config; + config.heartbeat_interval_ms = 100; + milvus::ClientTelemetryManager manager(config); + manager.AttachChannel(grpc::CreateChannel("127.0.0.1:" + std::to_string(port), grpc::InsecureChannelCredentials()), + "", "", "", "", ""); + milvus::proto::milvus::SearchRequest request; + manager.RecordOperation("Search", request, std::chrono::steady_clock::now(), true, ""); + + manager.Start(); + for (int retry = 0; retry < 2000 && manager.Config().enabled; ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_FALSE(manager.Config().enabled); + manager.Stop(); + manager.Start(); + ASSERT_TRUE(service.WaitForHeartbeats(2)); + ASSERT_FALSE(manager.Config().enabled); + manager.RecordOperation("Search", request, std::chrono::steady_clock::now(), true, ""); + service.AllowEnable(); + ASSERT_TRUE(service.WaitForHeartbeats(3)); + manager.Stop(); + + const auto requests = service.Requests(); + ASSERT_GE(requests.size(), 3U); + EXPECT_EQ(requests[0].metrics_size(), 1); + EXPECT_EQ(requests[1].metrics_size(), 0); + ASSERT_EQ(requests[1].command_replies_size(), 1); + EXPECT_EQ(requests[1].command_replies(0).command_id(), "disable"); + EXPECT_TRUE(requests[1].command_replies(0).success()); + const milvus::TelemetryCommand disable{"disable", "push_config", R"({"enabled":false})", 1, true, ""}; + EXPECT_EQ(requests[1].config_hash(), milvus::ClientTelemetryManager::CalculateConfigHash({disable})); + + EXPECT_EQ(requests[2].metrics_size(), 0); + ASSERT_EQ(requests[2].command_replies_size(), 1); + EXPECT_EQ(requests[2].command_replies(0).command_id(), "enable"); + EXPECT_TRUE(requests[2].command_replies(0).success()); + const milvus::TelemetryCommand enable{"enable", "push_config", R"({"enabled":true})", 2, true, ""}; + EXPECT_EQ(requests[2].config_hash(), milvus::ClientTelemetryManager::CalculateConfigHash({enable})); + EXPECT_TRUE(manager.Config().enabled); + + server->Shutdown(); +} + +TEST(ClientTelemetryTest, InitialDisabledConfigDoesNotActivateControlPlane) { + ControlPlaneTelemetryService service; + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + + milvus::TelemetryConfig config; + config.enabled = false; + config.heartbeat_interval_ms = 1; + milvus::ClientTelemetryManager manager(config); + manager.AttachChannel(grpc::CreateChannel("127.0.0.1:" + std::to_string(port), grpc::InsecureChannelCredentials()), + "", "", "", "", ""); + + manager.Start(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + manager.Stop(); + manager.Start(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + manager.Stop(); + + EXPECT_EQ(service.RequestCount(), 0U); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, UseDatabaseFromHeartbeatCommandReusesWorker) { + LifecycleTelemetryService service("use_database"); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(TelemetryConnectParam(port)).IsOk()); + std::atomic command_finished{false}; + std::atomic use_database_succeeded{false}; + std::weak_ptr weak_client = client; + client->GetTelemetry()->RegisterCommandHandler("use_database", [&](const milvus::TelemetryCommand& command) { + auto current_client = weak_client.lock(); + use_database_succeeded = current_client != nullptr && current_client->UseDatabase("secondary").IsOk(); + command_finished = true; + return milvus::TelemetryCommandReply{command.command_id, use_database_succeeded.load(), "", ""}; + }); + service.EnableCommands(); + + for (int retry = 0; retry < 3000 && (!command_finished.load() || service.heartbeats.load() < 2 || + !service.saw_secondary_database.load()); + ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + EXPECT_TRUE(command_finished.load()); + EXPECT_TRUE(use_database_succeeded.load()); + EXPECT_GE(service.connects.load(), 2); + EXPECT_GE(service.heartbeats.load(), 2); + EXPECT_TRUE(service.saw_secondary_database.load()); + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, AllLifecycleEntriesFailFastWhenConcurrentConnectWaitsForCommandLock) { + LifecycleTelemetryService service("connect"); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + const auto connect_param = TelemetryConnectParam(port); + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(connect_param).IsOk()); + + std::atomic handler_started{false}; + auto handler_finished = std::make_shared>>(); + auto handler_future = handler_finished->get_future(); + client->GetTelemetry()->RegisterCommandHandler( + "connect", [&, handler_finished](const milvus::TelemetryCommand& command) { + handler_started = true; + if (!service.WaitForConnects(2)) { + handler_finished->set_value({milvus::StatusCode::UNKNOWN_ERROR}); + return milvus::TelemetryCommandReply{command.command_id, false, "concurrent connect did not run", ""}; + } + + // The external Connect has completed its server handshake and owns lifecycle_mtx_, + // but its AttachChannel is blocked on this handler's command_mutex. Every re-entrant + // lifecycle entry must fail fast rather than wait and complete the lock cycle. + std::vector codes; + codes.push_back(client->Connect(connect_param).Code()); + codes.push_back(client->UseDatabase("other").Code()); + codes.push_back(client->Disconnect().Code()); + codes.push_back(client->SetRpcDeadlineMs(1234).Code()); + codes.push_back(client->SetRetryParam(milvus::RetryParam{}).Code()); + bool all_busy = true; + for (auto code : codes) { + all_busy = all_busy && code == milvus::StatusCode::CLIENT_BUSY; + } + handler_finished->set_value(codes); + return milvus::TelemetryCommandReply{command.command_id, all_busy, "", ""}; + }); + + service.EnableCommands(); + for (int retry = 0; retry < 2000 && !handler_started.load(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_TRUE(handler_started.load()); + + milvus::Status concurrent_status; + std::thread concurrent_connect([&]() { concurrent_status = client->Connect(connect_param); }); + const bool handler_failed_fast = handler_future.wait_for(std::chrono::seconds(1)) == std::future_status::ready; + concurrent_connect.join(); + + EXPECT_TRUE(handler_failed_fast); + ASSERT_EQ(handler_future.wait_for(std::chrono::seconds(2)), std::future_status::ready); + const auto codes = handler_future.get(); + ASSERT_EQ(codes.size(), 5U); + for (auto code : codes) { + EXPECT_EQ(code, milvus::StatusCode::CLIENT_BUSY); + } + EXPECT_TRUE(concurrent_status.IsOk()) << concurrent_status.Message(); + + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, FailedUseDatabasePreservesPublishedConnectionAndTelemetry) { + LifecycleTelemetryService service(""); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + auto connect_param = TelemetryConnectParam(port); + connect_param.SetDbName("primary"); + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(connect_param).IsOk()); + auto manager = client->GetTelemetry(); + ASSERT_NE(manager, nullptr); + const auto client_id = manager->ClientId(); + + service.FailNextConnect(); + const auto status = client->UseDatabase("secondary"); + EXPECT_EQ(status.Code(), milvus::StatusCode::SERVER_FAILED); + ASSERT_EQ(client->GetTelemetry(), manager); + EXPECT_EQ(manager->ClientId(), client_id); + + bool has_collection = false; + EXPECT_TRUE(client->HasCollection("still-connected", has_collection).IsOk()); + EXPECT_EQ(service.has_collections.load(), 1); + + manager->ProcessCommands({{"config-after-failure", "get_config", "", 10, false, ""}}); + const auto replies = manager->PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto user_config = nlohmann::json::parse(replies.back().payload).at("user_config"); + EXPECT_EQ(user_config.at("db_name"), "primary"); + + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, FailedConnectPreservesPublishedConnectionAndTelemetry) { + LifecycleTelemetryService first_service(""); + LifecycleTelemetryService rejected_service(""); + int first_port = 0; + int rejected_port = 0; + auto first_server = StartLifecycleServer(first_service, first_port); + auto rejected_server = StartLifecycleServer(rejected_service, rejected_port); + ASSERT_NE(first_server, nullptr); + ASSERT_NE(rejected_server, nullptr); + + auto first_param = TelemetryConnectParam(first_port); + first_param.SetDbName("primary"); + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(first_param).IsOk()); + auto manager = client->GetTelemetry(); + ASSERT_NE(manager, nullptr); + const auto client_id = manager->ClientId(); + + rejected_service.FailNextConnect(); + const auto status = client->Connect(TelemetryConnectParam(rejected_port)); + EXPECT_EQ(status.Code(), milvus::StatusCode::SERVER_FAILED); + ASSERT_EQ(client->GetTelemetry(), manager); + EXPECT_EQ(manager->ClientId(), client_id); + + bool has_collection = false; + EXPECT_TRUE(client->HasCollection("still-on-first", has_collection).IsOk()); + EXPECT_EQ(first_service.has_collections.load(), 1); + EXPECT_EQ(rejected_service.has_collections.load(), 0); + + manager->ProcessCommands({{"config-after-rejected-connect", "get_config", "", 10, false, ""}}); + const auto replies = manager->PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto user_config = nlohmann::json::parse(replies.back().payload).at("user_config"); + EXPECT_EQ(user_config.at("address"), "127.0.0.1:" + std::to_string(first_port)); + EXPECT_EQ(user_config.at("db_name"), "primary"); + + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + first_server->Shutdown(); + rejected_server->Shutdown(); +} + +TEST(ClientTelemetryTest, NonStandardCommandExceptionReturnsFailureAndHeartbeatContinues) { + LifecycleTelemetryService service("throw_non_standard"); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(TelemetryConnectParam(port)).IsOk()); + client->GetTelemetry()->RegisterCommandHandler( + "throw_non_standard", [](const milvus::TelemetryCommand&) -> milvus::TelemetryCommandReply { throw 42; }); + service.EnableCommands(); + + for (int retry = 0; retry < 3000 && (service.heartbeats.load() < 3 || !service.saw_failed_command_reply.load()); + ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_GE(service.heartbeats.load(), 3); + EXPECT_TRUE(service.saw_failed_command_reply.load()); + + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, LastManagerReferenceCanBeDestroyedByHeartbeatCommand) { + LifecycleTelemetryService service("disconnect_and_destroy"); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + auto client = milvus::MilvusClient::Create(); + ASSERT_TRUE(client->Connect(TelemetryConnectParam(port)).IsOk()); + auto destroyed = std::make_shared>(); + auto destroyed_future = destroyed->get_future(); + auto marker = std::make_shared(destroyed); + client->GetTelemetry()->RegisterCommandHandler( + "disconnect_and_destroy", [&client, marker](const milvus::TelemetryCommand& command) { + auto status = client->Disconnect(); + client.reset(); + return milvus::TelemetryCommandReply{command.command_id, status.IsOk(), "", ""}; + }); + marker.reset(); + service.EnableCommands(); + + ASSERT_EQ(destroyed_future.wait_for(std::chrono::seconds(3)), std::future_status::ready); + EXPECT_EQ(client, nullptr); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, V2LastManagerReferenceCanBeDestroyedByHeartbeatCommand) { + LifecycleTelemetryService service("disconnect_and_destroy_v2"); + int port = 0; + auto server = StartLifecycleServer(service, port); + ASSERT_NE(server, nullptr); + + auto client = milvus::MilvusClientV2::Create(); + ASSERT_TRUE(client->Connect(TelemetryConnectParam(port)).IsOk()); + auto destroyed = std::make_shared>(); + auto destroyed_future = destroyed->get_future(); + auto marker = std::make_shared(destroyed); + client->GetTelemetry()->RegisterCommandHandler( + "disconnect_and_destroy_v2", [&client, marker](const milvus::TelemetryCommand& command) { + auto status = client->Disconnect(); + client.reset(); + return milvus::TelemetryCommandReply{command.command_id, status.IsOk(), "", ""}; + }); + marker.reset(); + service.EnableCommands(); + + ASSERT_EQ(destroyed_future.wait_for(std::chrono::seconds(3)), std::future_status::ready); + EXPECT_EQ(client, nullptr); + server->Shutdown(); +} + +TEST(ClientRequestContextTest, PropagatesOnlyValidTraceIdMetadataOnWire) { + RequestMetadataService service; + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + + auto client = milvus::MilvusClient::Create(); + auto connect_param = TelemetryConnectParam(port); + milvus::TelemetryConfig telemetry_config; + telemetry_config.enabled = false; + connect_param.SetTelemetryConfig(telemetry_config); + ASSERT_TRUE(client->Connect(connect_param).IsOk()); + + bool has_collection = false; + constexpr const char* valid_request_id = "4bf92f3577b34da6a3ce929d0e0e4736"; + { + milvus::ScopedClientRequestId request_id(valid_request_id); + EXPECT_TRUE(client->HasCollection("valid", has_collection).IsOk()); + } + { + milvus::ScopedClientRequestId request_id(""); + EXPECT_TRUE(client->HasCollection("empty", has_collection).IsOk()); + } + { + milvus::ScopedClientRequestId request_id("ABCDEF0123456789ABCDEF0123456789"); + EXPECT_TRUE(client->HasCollection("invalid", has_collection).IsOk()); + } + + const auto request_ids = service.RequestIds(); + ASSERT_EQ(request_ids.size(), 3U); + EXPECT_TRUE(request_ids[0].first); + EXPECT_EQ(request_ids[0].second, valid_request_id); + EXPECT_FALSE(request_ids[1].first); + EXPECT_FALSE(request_ids[2].first); + + EXPECT_TRUE(client->Disconnect().IsOk()); + client.reset(); + server->Shutdown(); +} + +TEST(ClientTelemetryTest, ReconnectReuseMatchesOriginalUserConfig) { + milvus::TelemetryConfig config; + config.enabled = false; + config.sampling_rate = 0.5; + milvus::ClientTelemetryManager manager(config); + + manager.ProcessCommands({{"remote", "push_config", R"({"sampling_rate":0.25})", 1, true, ""}}); + EXPECT_DOUBLE_EQ(manager.Config().sampling_rate, 0.25); + EXPECT_TRUE(manager.MatchesConnection(config, "")); + + auto changed = config; + changed.enabled = true; + EXPECT_FALSE(manager.MatchesConnection(changed, "")); +} + +TEST(ClientTelemetryTest, GlobalPhysicalHandoffAndUseDatabaseKeepLogicalIdentityAndState) { + LifecycleTelemetryService first_service(""); + LifecycleTelemetryService second_service(""); + int first_port = 0; + int second_port = 0; + auto first_server = StartLifecycleServer(first_service, first_port); + auto second_server = StartLifecycleServer(second_service, second_port); + ASSERT_NE(first_server, nullptr); + ASSERT_NE(second_server, nullptr); + + constexpr const char* logical_endpoint = "https://tenant.global-cluster.example.com"; + milvus::ConnectParam first_param("http://127.0.0.1:" + std::to_string(first_port)); + first_param.SetDbName("primary_db"); + milvus::TelemetryConfig telemetry_config; + telemetry_config.enabled = false; + first_param.SetTelemetryConfig(telemetry_config); + + auto first_connection = std::make_shared(); + ASSERT_TRUE(first_connection->Connect(first_param, "", nullptr, logical_endpoint).IsOk()); + auto manager = first_connection->GetTelemetry(); + ASSERT_NE(manager, nullptr); + const auto client_id = manager->ClientId(); + manager->ProcessCommands({ + {"config", "push_config", R"({"sampling_rate":0.25})", 1, true, ""}, + {"collections", "collection_metrics", R"({"enabled":true,"collections":["before_failover"]})", 2, false, ""}, + }); + const auto config_hash = manager->ConfigHash(); + const auto replies_before_failover = manager->PendingCommandReplies(); + ASSERT_EQ(replies_before_failover.size(), 2U); + + milvus::ConnectParam second_param = first_param; + second_param.SetUri("http://127.0.0.1:" + std::to_string(second_port)); + auto second_connection = std::make_shared(); + ASSERT_TRUE(second_connection->Connect(second_param, client_id, manager, logical_endpoint).IsOk()); + ASSERT_EQ(second_connection->GetTelemetry(), manager); + EXPECT_EQ(manager->ClientId(), client_id); + EXPECT_EQ(manager->ConfigHash(), config_hash); + EXPECT_EQ(manager->LastCommandTimestamp(), 2); + EXPECT_DOUBLE_EQ(manager->Config().sampling_rate, 0.25); + EXPECT_EQ(manager->PendingCommandReplies().size(), replies_before_failover.size()); + EXPECT_TRUE(first_connection->Disconnect(false).IsOk()); + + ASSERT_TRUE(second_connection->UseDatabase("secondary_db").IsOk()); + ASSERT_EQ(second_connection->GetTelemetry(), manager); + EXPECT_EQ(manager->ClientId(), client_id); + EXPECT_EQ(manager->ConfigHash(), config_hash); + EXPECT_EQ(manager->LastCommandTimestamp(), 2); + + manager->ProcessCommands({{"config-after-use-db", "get_config", "", 3, false, ""}}); + const auto replies = manager->PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto user_config = nlohmann::json::parse(replies.back().payload).at("user_config"); + EXPECT_EQ(user_config.at("address"), logical_endpoint); + EXPECT_EQ(user_config.at("db_name"), "secondary_db"); + + EXPECT_TRUE(second_connection->Disconnect().IsOk()); + first_server->Shutdown(); + second_server->Shutdown(); +} + +TEST(ClientTelemetryTest, PushConfigIsAtomicAndReportsAppliedAndIgnoredKeys) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + + manager.ProcessCommands( + {{"invalid", "push_config", R"({"enabled":true,"heartbeat_interval_ms":0})", 1, false, ""}}); + EXPECT_FALSE(manager.Config().enabled); + ASSERT_EQ(manager.PendingCommandReplies().size(), 1U); + EXPECT_FALSE(manager.PendingCommandReplies().back().success); + + manager.ProcessCommands( + {{"valid", "push_config", + R"({"unknown_b":1,"sampling_rate":2,"enabled":true,"ttl_seconds":3,"heartbeat_interval_ms":2500,"unknown_a":2})", + 2, false, ""}}); + auto updated = manager.Config(); + EXPECT_TRUE(updated.enabled); + EXPECT_EQ(updated.heartbeat_interval_ms, 2500U); + EXPECT_DOUBLE_EQ(updated.sampling_rate, 1.0); + + auto replies = manager.PendingCommandReplies(); + ASSERT_EQ(replies.size(), 2U); + ASSERT_TRUE(replies.back().success); + auto payload = nlohmann::json::parse(replies.back().payload); + EXPECT_EQ(payload["applied"], nlohmann::json({"enabled", "heartbeat_interval_ms", "sampling_rate"})); + EXPECT_EQ(payload["ignored"], nlohmann::json({"ttl_seconds", "unknown_a", "unknown_b"})); +} + +TEST(ClientTelemetryTest, RejectsWrongCommandPayloadTypes) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + + manager.ProcessCommands( + {{"push", "push_config", R"({"enabled":"false"})", 1, false, ""}, + {"collection", "collection_metrics", R"({"enabled":false,"collections":"books"})", 2, false, ""}, + {"ttl", "push_config", R"({"enabled":true,"ttl_seconds":"bad"})", 3, false, ""}}); + + auto replies = manager.PendingCommandReplies(); + ASSERT_EQ(replies.size(), 3U); + EXPECT_FALSE(replies[0].success); + EXPECT_FALSE(replies[1].success); + EXPECT_FALSE(replies[2].success); + EXPECT_FALSE(manager.Config().enabled); +} + +TEST(ClientTelemetryTest, RejectsInvalidRfc3339CalendarTimes) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + + manager.ProcessCommands( + {{"invalid-day", "show_latency_history", + R"({"start_time":"2026-02-30T00:00:00Z","end_time":"2026-03-01T00:00:00Z","detail":false})", 1, false, ""}, + {"leap-second", "show_latency_history", + R"({"start_time":"2026-02-28T23:59:60Z","end_time":"2026-03-01T00:00:00Z","detail":false})", 2, false, ""}, + {"missing-seconds", "show_latency_history", + R"({"start_time":"2026-02-28T23:59Z","end_time":"2026-03-01T00:00:00Z","detail":false})", 3, false, ""}}); + + const auto replies = manager.PendingCommandReplies(); + ASSERT_EQ(replies.size(), 3U); + EXPECT_FALSE(replies[0].success); + EXPECT_FALSE(replies[1].success); + EXPECT_FALSE(replies[2].success); +} + +TEST(ClientTelemetryTest, SerializesConcurrentCommandBatches) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + std::atomic calls{0}; + std::atomic release{false}; + manager.RegisterCommandHandler("custom", [&calls, &release](const milvus::TelemetryCommand& command) { + ++calls; + while (!release.load()) { + std::this_thread::yield(); + } + return milvus::TelemetryCommandReply{command.command_id, true, "", ""}; + }); + const milvus::TelemetryCommand command{"same", "custom", "", 1, false, ""}; + + std::thread first([&]() { manager.ProcessCommands({command}); }); + while (calls.load() == 0) { + std::this_thread::yield(); + } + std::thread second([&]() { manager.ProcessCommands({command}); }); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + release = true; + first.join(); + second.join(); + + EXPECT_EQ(calls.load(), 1); +} + +TEST(ClientTelemetryTest, UsesOneFixedPointSamplerAcrossOperations) { + milvus::TelemetryConfig config; + config.sampling_rate = 0.25; + milvus::ClientTelemetryManager manager(config); + milvus::proto::milvus::SearchRequest request; + request.set_collection_name("books"); + + for (int index = 0; index < 12; ++index) { + manager.RecordOperation(index % 2 == 0 ? "Search" : "Query", request, std::chrono::steady_clock::now(), true, + ""); + } + manager.Start(); + for (int retry = 0; retry < 100 && manager.MetricsSnapshots().empty(); ++retry) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + manager.Stop(); + + auto snapshots = manager.MetricsSnapshots(); + ASSERT_FALSE(snapshots.empty()); + int64_t sampled = 0; + for (const auto& operation : snapshots.back().metrics) { + sampled += operation.global.request_count; + } + EXPECT_EQ(sampled, 3); +} + +TEST(ClientTelemetryTest, StopAndRestartPreserveCommandState) { + milvus::TelemetryConfig config; + config.enabled = false; + milvus::ClientTelemetryManager manager(config); + int calls = 0; + manager.RegisterCommandHandler("custom", [&calls](const milvus::TelemetryCommand& command) { + ++calls; + return milvus::TelemetryCommandReply{command.command_id, true, "", ""}; + }); + const milvus::TelemetryCommand command{"custom", "custom", "", 2, false, ""}; + + manager.ProcessCommands({command}); + manager.Start(); + EXPECT_TRUE(manager.IsReady()); + manager.Stop(); + EXPECT_FALSE(manager.IsReady()); + manager.Start(); + EXPECT_TRUE(manager.IsReady()); + manager.ProcessCommands({command}); + + EXPECT_EQ(calls, 1); + EXPECT_EQ(manager.LastCommandTimestamp(), 2); + manager.Stop(); +} + +TEST(ClientRequestContextTest, GeneratesAndScopesTraceIds) { + auto request_id = milvus::ClientRequestContext::NewRequestId(); + EXPECT_EQ(request_id.size(), 32U); + EXPECT_EQ(request_id.find_first_not_of("0123456789abcdef"), std::string::npos); + EXPECT_NE(request_id, std::string(32, '0')); + EXPECT_TRUE(milvus::ClientRequestContext::IsValid(request_id)); + EXPECT_FALSE(milvus::ClientRequestContext::IsValid(std::string(32, '0'))); + EXPECT_FALSE(milvus::ClientRequestContext::IsValid("ABCDEF0123456789ABCDEF0123456789")); + EXPECT_FALSE(milvus::ClientRequestContext::IsValid("short")); + + milvus::ClientRequestContext::Set("outer"); + { + milvus::ScopedClientRequestId scoped("inner"); + EXPECT_EQ(milvus::ClientRequestContext::Get(), "inner"); + } + EXPECT_EQ(milvus::ClientRequestContext::Get(), "outer"); + milvus::ClientRequestContext::Clear(); +} diff --git a/test/ut/utils/TestGlobalCluster.cpp b/test/ut/utils/TestGlobalCluster.cpp index 267f9a2f..f4bd0634 100644 --- a/test/ut/utils/TestGlobalCluster.cpp +++ b/test/ut/utils/TestGlobalCluster.cpp @@ -14,26 +14,78 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include #include +#include +#include #include #include #include "cpp-httplib/httplib.h" +#include "milvus.grpc.pb.h" +#include "milvus.pb.h" +#include "milvus/ClientTelemetry.h" +#include "milvus/thirdparty/nlohmann/json.hpp" +#include "milvus/types/ConnectParam.h" #include "types/GlobalCluster.h" +#include "utils/ConnectionHandler.h" #include "utils/GlobalClusterUtils.h" #include "utils/TopologyRefresher.h" namespace { std::string -TopologyBody(int64_t version) { +TopologyBody(int64_t version, const std::string& endpoint = "a:19530") { return R"({"code":0,"data":{"version":)" + std::to_string(version) + - R"(,"clusters":[{"clusterId":"a","endpoint":"a:19530","capability":3}]}})"; + R"(,"clusters":[{"clusterId":"a","endpoint":")" + endpoint + R"(","capability":3}]}})"; } +class GlobalMilvusService final : public milvus::proto::milvus::MilvusService::Service { + public: + grpc::Status + Connect(grpc::ServerContext*, const milvus::proto::milvus::ConnectRequest*, + milvus::proto::milvus::ConnectResponse*) override { + ++connects; + return grpc::Status::OK; + } + + std::atomic connects{0}; +}; + +std::unique_ptr +StartGlobalMilvusServer(GlobalMilvusService& service, int& port) { + grpc::ServerBuilder builder; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &port); + builder.RegisterService(&service); + return builder.BuildAndStart(); +} + +class ScopedTopologyServerThread { + public: + explicit ScopedTopologyServerThread(httplib::Server& server) + : server_(server), thread_([&server]() { server.listen_after_bind(); }) { + } + + ~ScopedTopologyServerThread() { + server_.stop(); + if (thread_.joinable()) { + thread_.join(); + } + } + + ScopedTopologyServerThread(const ScopedTopologyServerThread&) = delete; + ScopedTopologyServerThread& + operator=(const ScopedTopologyServerThread&) = delete; + + private: + httplib::Server& server_; + std::thread thread_; +}; + } // namespace TEST(GlobalClusterUtilsTest, IsGlobalEndpoint) { @@ -265,3 +317,88 @@ TEST(TopologyRefresherTest, CallbackOnVersionChange) { server.stop(); srv.join(); } + +TEST(GlobalClusterTelemetryTest, FailoverAndUseDatabasePreserveLogicalTelemetryState) { + GlobalMilvusService first_service; + GlobalMilvusService second_service; + int first_port = 0; + int second_port = 0; + auto first_server = StartGlobalMilvusServer(first_service, first_port); + auto second_server = StartGlobalMilvusServer(second_service, second_port); + ASSERT_NE(first_server, nullptr); + ASSERT_NE(second_server, nullptr); + + httplib::Server topology_server; + std::atomic served_version{1}; + std::atomic served_primary_port{first_port}; + // Put the global-cluster marker in the path rather than a localhost subdomain. Windows does + // not consistently resolve arbitrary *.localhost names, while a numeric loopback address is + // portable across all CI runners. + topology_server.Get( + "/global-cluster-test/global-cluster/topology", [&](const httplib::Request&, httplib::Response& response) { + response.set_content( + TopologyBody(served_version.load(), "127.0.0.1:" + std::to_string(served_primary_port.load())), + "application/json"); + }); + const auto topology_port = topology_server.bind_to_any_port("127.0.0.1"); + ASSERT_TRUE(topology_server.is_valid()); + ScopedTopologyServerThread topology_thread(topology_server); + topology_server.wait_until_ready(); + + const auto logical_endpoint = "http://127.0.0.1:" + std::to_string(topology_port) + "/global-cluster-test"; + milvus::ConnectParam connect_param(logical_endpoint); + connect_param.SetDbName("primary_db"); + milvus::TelemetryConfig telemetry_config; + telemetry_config.enabled = false; + connect_param.SetTelemetryConfig(telemetry_config); + + milvus::ConnectionHandler handler; + ASSERT_TRUE(handler.Connect(connect_param).IsOk()); + auto manager = handler.GetTelemetry(); + ASSERT_NE(manager, nullptr); + const auto client_id = manager->ClientId(); + manager->ProcessCommands({ + {"config", "push_config", R"({"sampling_rate":0.25})", 1, true, ""}, + {"collections", "collection_metrics", R"({"enabled":true,"collections":["before_failover"]})", 2, false, ""}, + }); + const auto config_hash = manager->ConfigHash(); + const auto reply_count = manager->PendingCommandReplies().size(); + ASSERT_EQ(reply_count, 2U); + + served_primary_port = second_port; + served_version = 2; + handler.TriggerGlobalRefresh(); + const auto second_uri = "http://127.0.0.1:" + std::to_string(second_port); + for (int retry = 0; retry < 3000; ++retry) { + auto connection = handler.GetConnection(); + if (connection != nullptr && connection->GetConnectParam().Uri() == second_uri) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_NE(handler.GetConnection(), nullptr); + EXPECT_EQ(handler.GetConnection()->GetConnectParam().Uri(), second_uri); + EXPECT_EQ(handler.CurrentEndpoint(), logical_endpoint); + EXPECT_EQ(handler.GetTelemetry(), manager); + EXPECT_EQ(manager->ClientId(), client_id); + EXPECT_EQ(manager->ConfigHash(), config_hash); + EXPECT_EQ(manager->LastCommandTimestamp(), 2); + EXPECT_DOUBLE_EQ(manager->Config().sampling_rate, 0.25); + EXPECT_EQ(manager->PendingCommandReplies().size(), reply_count); + + ASSERT_TRUE(handler.UseDatabase("secondary_db").IsOk()); + ASSERT_EQ(handler.GetTelemetry(), manager); + manager->ProcessCommands({{"config-after-use-db", "get_config", "", 3, false, ""}}); + const auto replies = manager->PendingCommandReplies(); + ASSERT_FALSE(replies.empty()); + ASSERT_TRUE(replies.back().success) << replies.back().error_message; + const auto user_config = nlohmann::json::parse(replies.back().payload).at("user_config"); + EXPECT_EQ(user_config.at("address"), logical_endpoint); + EXPECT_EQ(user_config.at("db_name"), "secondary_db"); + EXPECT_EQ(manager->ClientId(), client_id); + EXPECT_EQ(manager->ConfigHash(), config_hash); + + EXPECT_TRUE(handler.Disconnect().IsOk()); + first_server->Shutdown(); + second_server->Shutdown(); +}