From a102d532274d05fe876d96dd7f5566f4e4a1fd12 Mon Sep 17 00:00:00 2001 From: maladetska Date: Sat, 7 Feb 2026 15:57:16 +0300 Subject: [PATCH 01/93] init commit --- cmake/external_libs.cmake | 1 + include/ydb-cpp-sdk/client/metrics/metrics.h | 81 ++++++++++ .../ydb-cpp-sdk/open_telemetry/extension.h | 34 ++++ include/ydb-cpp-sdk/open_telemetry/otel.h | 33 ++++ src/CMakeLists.txt | 3 +- src/client/CMakeLists.txt | 1 + .../grpc_connections/grpc_connections.h | 12 ++ src/client/metrics/CMakeLists.txt | 14 ++ src/client/metrics/metrics.cpp | 32 ++++ src/client/query/CMakeLists.txt | 1 + src/client/query/client.cpp | 55 ++++++- src/client/query/impl/CMakeLists.txt | 1 + src/client/query/impl/query_spans.cpp | 56 +++++++ src/client/query/impl/query_spans.h | 23 +++ src/open_telemetry/CMakeLists.txt | 17 ++ src/open_telemetry/otel.cpp | 153 ++++++++++++++++++ 16 files changed, 509 insertions(+), 8 deletions(-) create mode 100644 include/ydb-cpp-sdk/client/metrics/metrics.h create mode 100644 include/ydb-cpp-sdk/open_telemetry/extension.h create mode 100644 include/ydb-cpp-sdk/open_telemetry/otel.h create mode 100644 src/client/metrics/CMakeLists.txt create mode 100644 src/client/metrics/metrics.cpp create mode 100644 src/client/query/impl/query_spans.cpp create mode 100644 src/client/query/impl/query_spans.h create mode 100644 src/open_telemetry/CMakeLists.txt create mode 100644 src/open_telemetry/otel.cpp diff --git a/cmake/external_libs.cmake b/cmake/external_libs.cmake index dc46fdb1d5e..be645f332ae 100644 --- a/cmake/external_libs.cmake +++ b/cmake/external_libs.cmake @@ -13,6 +13,7 @@ find_package(base64 REQUIRED) find_package(Brotli 1.1.0 REQUIRED) find_package(jwt-cpp REQUIRED) find_package(double-conversion REQUIRED) +find_package(opentelemetry-cpp REQUIRED) # RapidJSON if (YDB_SDK_USE_RAPID_JSON) diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h new file mode 100644 index 00000000000..d3dc5eb2c19 --- /dev/null +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -0,0 +1,81 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace NYdb::inline V3::NMetrics { + +using TLabels = std::map; + +class ICounter { +public: + virtual ~ICounter() = default; + virtual void Inc() = 0; +}; + +class IGauge { +public: + virtual ~IGauge() = default; + virtual void Add(double delta) = 0; + virtual void Set(double value) = 0; +}; + +class IHistogram { +public: + virtual ~IHistogram() = default; + virtual void Record(double value) = 0; +}; + +class IMetricRegistry { +public: + virtual ~IMetricRegistry() = default; + + virtual std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) = 0; + virtual std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) = 0; + virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; +}; + +enum class ESpanKind { + INTERNAL, + SERVER, + CLIENT, + PRODUCER, + CONSUMER +}; + +class ISpan { +public: + virtual ~ISpan() = default; + virtual void End() = 0; + virtual void SetAttribute(const std::string& key, const std::string& value) = 0; + virtual void SetAttribute(const std::string& key, int64_t value) = 0; +}; + +class ITracer { +public: + virtual ~ITracer() = default; + virtual std::shared_ptr StartSpan(const std::string& name, ESpanKind kind = ESpanKind::INTERNAL) = 0; +}; + +class ITraceProvider { +public: + virtual ~ITraceProvider() = default; + virtual std::shared_ptr GetTracer(const std::string& name) = 0; +}; + +class IMetricsApi : public IExtensionApi { +public: + static IMetricsApi* Create(TDriver driver); +public: + virtual ~IMetricsApi() = default; + virtual void SetMetricRegistry(std::shared_ptr registry) = 0; + virtual void SetTraceProvider(std::shared_ptr provider) = 0; + virtual std::shared_ptr GetMetricRegistry() const = 0; + virtual std::shared_ptr GetTraceProvider() const = 0; +}; + +} // namespace NYdb::NMetrics diff --git a/include/ydb-cpp-sdk/open_telemetry/extension.h b/include/ydb-cpp-sdk/open_telemetry/extension.h new file mode 100644 index 00000000000..b5683d6d41b --- /dev/null +++ b/include/ydb-cpp-sdk/open_telemetry/extension.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelExtension : public IExtension { +public: + using IApi = IMetricsApi; + + struct TParams { + opentelemetry::nostd::shared_ptr MeterProvider; + opentelemetry::nostd::shared_ptr TracerProvider; + }; + + TOtelExtension(const TParams& params, IApi* api) { + if (params.MeterProvider) { + api->SetMetricRegistry(std::make_shared(params.MeterProvider)); + } + if (params.TracerProvider) { + api->SetTraceProvider(std::make_shared(params.TracerProvider)); + } + } +}; + +inline void AddOpenTelemetry(TDriver& driver + , opentelemetry::nostd::shared_ptr meterProvider + , opentelemetry::nostd::shared_ptr tracerProvider +) { + driver.AddExtension({meterProvider, tracerProvider}); +} + +} // namespace NYdb::NMetrics diff --git a/include/ydb-cpp-sdk/open_telemetry/otel.h b/include/ydb-cpp-sdk/open_telemetry/otel.h new file mode 100644 index 00000000000..b1d23c92fa6 --- /dev/null +++ b/include/ydb-cpp-sdk/open_telemetry/otel.h @@ -0,0 +1,33 @@ +#pragma once + +#include + +#include +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelMetricRegistry : public IMetricRegistry { +public: + TOtelMetricRegistry(opentelemetry::nostd::shared_ptr meterProvider); + + std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) override; + std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) override; + std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) override; + +private: + opentelemetry::nostd::shared_ptr MeterProvider_; + opentelemetry::nostd::shared_ptr Meter_; +}; + +class TOtelTraceProvider : public ITraceProvider { +public: + TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); + + std::shared_ptr GetTracer(const std::string& name) override; + +private: + opentelemetry::nostd::shared_ptr TracerProvider_; +}; + +} // namespace NYdb::NMetrics diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3dff7094058..c3fa4e733cd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,3 +1,4 @@ add_subdirectory(api) add_subdirectory(client) -add_subdirectory(library) \ No newline at end of file +add_subdirectory(library) +add_subdirectory(open_telemetry) diff --git a/src/client/CMakeLists.txt b/src/client/CMakeLists.txt index e7f448e8675..65167cddbd0 100644 --- a/src/client/CMakeLists.txt +++ b/src/client/CMakeLists.txt @@ -12,6 +12,7 @@ add_subdirectory(iam) add_subdirectory(iam_private) add_subdirectory(impl) add_subdirectory(import) +add_subdirectory(metrics) add_subdirectory(monitoring) add_subdirectory(operation) add_subdirectory(params) diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 756d2f0d957..15cddb4e6ec 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -581,6 +581,18 @@ class TGRpcConnectionsImpl ::NMonitoring::TMetricRegistry* GetMetricRegistry() override; void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); + + template + T* GetExtensionApi() { + std::lock_guard lock(ExtensionsLock_); + for (const auto& api : ExtensionApis_) { + if (auto ptr = dynamic_cast(api.get())) { + return ptr; + } + } + return nullptr; + } + void SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb); const TLog& GetLog() const override; diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt new file mode 100644 index 00000000000..94902b0f415 --- /dev/null +++ b/src/client/metrics/CMakeLists.txt @@ -0,0 +1,14 @@ +_ydb_sdk_add_library(client-metrics) + +target_sources(client-metrics PRIVATE + metrics.cpp +) + +target_include_directories(client-metrics PUBLIC + $ + $ +) + +target_link_libraries(client-metrics PUBLIC + client-extension_common +) diff --git a/src/client/metrics/metrics.cpp b/src/client/metrics/metrics.cpp new file mode 100644 index 00000000000..836d01f2071 --- /dev/null +++ b/src/client/metrics/metrics.cpp @@ -0,0 +1,32 @@ +#include + +namespace NYdb::inline V3::NMetrics { + +class TMetricsApiImpl : public IMetricsApi { +public: + void SetMetricRegistry(std::shared_ptr registry) override { + Registry_ = std::move(registry); + } + + void SetTraceProvider(std::shared_ptr provider) override { + TraceProvider_ = std::move(provider); + } + + std::shared_ptr GetMetricRegistry() const override { + return Registry_; + } + + std::shared_ptr GetTraceProvider() const override { + return TraceProvider_; + } + +private: + std::shared_ptr Registry_; + std::shared_ptr TraceProvider_; +}; + +IMetricsApi* IMetricsApi::Create(TDriver driver) { + return new TMetricsApiImpl(); +} + +} // namespace NYdb::NMetrics diff --git a/src/client/query/CMakeLists.txt b/src/client/query/CMakeLists.txt index 6677d402d4d..f1395ff107b 100644 --- a/src/client/query/CMakeLists.txt +++ b/src/client/query/CMakeLists.txt @@ -7,6 +7,7 @@ target_link_libraries(client-ydb_query PUBLIC impl-internal-make_request impl-session impl-internal-retry + client-metrics client-ydb_common_client client-ydb_driver client-ydb_query-impl diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index ccf90f1175c..ca1453e0faf 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -15,7 +15,9 @@ #include #include #include +#include #include +#include #include @@ -67,6 +69,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public { SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); + + if (auto metricsApi = Connections_->GetExtensionApi()) { + if (auto traceProvider = metricsApi->GetTraceProvider()) { + Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); + } + } } ~TImpl() { @@ -94,8 +102,21 @@ class TQueryClient::TImpl: public TClientImplCommon, public { CollectQuerySize(query); CollectParamsSize(params ? ¶ms->GetProtoMap() : nullptr); + + auto span = std::make_shared(Tracer_, "ExecuteQuery", DbDriverState_->DiscoveryEndpoint); + return TExecQueryImpl::ExecuteQuery( - Connections_, DbDriverState_, query, txControl, params, settings, session); + Connections_, DbDriverState_, query, txControl, params, settings, session) + .Apply([span](TAsyncExecuteQueryResult future) { + try { + auto result = future.GetValue(); + span->End(result.GetStatus()); + return result; + } catch (...) { + span->End(EStatus::CLIENT_INTERNAL_ERROR); + throw; + } + }); } NThreading::TFuture ExecuteScript(const std::string& script, const std::optional& params, const TExecuteScriptSettings& settings) { @@ -162,7 +183,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto promise = NThreading::NewPromise(); - auto responseCb = [promise, session] + auto span = std::make_shared(Tracer_, "Rollback", DbDriverState_->DiscoveryEndpoint); + + auto responseCb = [promise, session, span] (Ydb::Query::RollbackTransactionResponse* response, TPlainStatus status) mutable { try { if (response) { @@ -171,11 +194,15 @@ class TQueryClient::TImpl: public TClientImplCommon, public TStatus rollbackTxStatus(TPlainStatus{static_cast(response->status()), std::move(opIssues), status.Endpoint, std::move(status.Metadata)}); + span->End(rollbackTxStatus.GetStatus()); + promise.SetValue(std::move(rollbackTxStatus)); } else { + span->End(status.Status); promise.SetValue(TStatus(std::move(status))); } } catch (...) { + span->End(EStatus::CLIENT_INTERNAL_ERROR); promise.SetException(std::current_exception()); } }; @@ -203,7 +230,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto promise = NThreading::NewPromise(); - auto responseCb = [promise, session] + auto span = std::make_shared(Tracer_, "Commit", DbDriverState_->DiscoveryEndpoint); + + auto responseCb = [promise, session, span] (Ydb::Query::CommitTransactionResponse* response, TPlainStatus status) mutable { try { if (response) { @@ -212,12 +241,16 @@ class TQueryClient::TImpl: public TClientImplCommon, public TStatus commitTxStatus(TPlainStatus{static_cast(response->status()), std::move(opIssues), status.Endpoint, std::move(status.Metadata)}); + span->End(commitTxStatus.GetStatus()); + TCommitTransactionResult commitTxResult(std::move(commitTxStatus)); promise.SetValue(std::move(commitTxResult)); } else { + span->End(status.Status); promise.SetValue(TCommitTransactionResult(TStatus(std::move(status)))); } } catch (...) { + span->End(EStatus::CLIENT_INTERNAL_ERROR); promise.SetException(std::current_exception()); } }; @@ -425,10 +458,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public TAsyncCreateSessionResult GetSession(const TCreateSessionSettings& settings) { class TQueryClientGetSessionCtx : public NSessionPool::IGetSessionCtx { public: - TQueryClientGetSessionCtx(std::shared_ptr client, const TCreateSessionSettings& settings) + TQueryClientGetSessionCtx(std::shared_ptr client, const TCreateSessionSettings& settings, std::shared_ptr span) : Promise(NThreading::NewPromise()) , Client(client) , RpcSettings(TRpcRequestSettings::Make(settings)) + , Span(span) {} TAsyncCreateSessionResult GetFuture() { @@ -437,6 +471,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public void ReplyError(TStatus status) override { TSession session; + if (Span) Span->End(status.GetStatus()); ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -449,14 +484,17 @@ class TQueryClient::TImpl: public TClientImplCommon, public ) ); + if (Span) Span->End(EStatus::SUCCESS); ScheduleReply(std::move(val)); } void ReplyNewSession() override { Client->CreateAttachedSession(RpcSettings).Subscribe( - [promise{std::move(Promise)}](TAsyncCreateSessionResult future) mutable + [promise{std::move(Promise)}, span = Span](TAsyncCreateSessionResult future) mutable { - promise.SetValue(future.ExtractValue()); + auto val = future.ExtractValue(); + if (span) span->End(val.GetStatus()); + promise.SetValue(std::move(val)); }); } @@ -481,9 +519,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public NThreading::TPromise Promise; std::shared_ptr Client; const TRpcRequestSettings RpcSettings; + std::shared_ptr Span; }; - auto ctx = std::make_unique(shared_from_this(), settings); + auto span = std::make_shared(Tracer_, "CreateSession", DbDriverState_->DiscoveryEndpoint); + auto ctx = std::make_unique(shared_from_this(), settings, span); auto future = ctx->GetFuture(); SessionPool_.GetSession(std::move(ctx)); @@ -552,6 +592,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public } private: + std::shared_ptr Tracer_; NSdkStats::TStatCollector::TClientRetryOperationStatCollector RetryOperationStatCollector_; NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> QuerySizeHistogram_; NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> ParamsSizeHistogram_; diff --git a/src/client/query/impl/CMakeLists.txt b/src/client/query/impl/CMakeLists.txt index 76b112b2254..70f93b6d68d 100644 --- a/src/client/query/impl/CMakeLists.txt +++ b/src/client/query/impl/CMakeLists.txt @@ -12,6 +12,7 @@ target_link_libraries(client-ydb_query-impl PUBLIC target_sources(client-ydb_query-impl PRIVATE exec_query.cpp client_session.cpp + query_spans.cpp ) _ydb_sdk_install_targets(TARGETS client-ydb_query-impl) diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp new file mode 100644 index 00000000000..c612f442663 --- /dev/null +++ b/src/client/query/impl/query_spans.cpp @@ -0,0 +1,56 @@ +#include "query_spans.h" + +#include + +namespace NYdb::inline V3::NQuery { + +namespace { + +void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { + auto pos = endpoint.find(':'); + if (pos != std::string::npos) { + host = endpoint.substr(0, pos); + try { + port = std::stoi(endpoint.substr(pos + 1)); + } catch (...) { + port = 2135; + } + } else { + host = endpoint; + port = 2135; + } +} + +} // namespace + +TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { + if (!tracer) return; + + std::string host; + int port; + ParseEndpoint(endpoint, host, port); + + Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); + Span_->SetAttribute("db.system.name", "ydb"); + Span_->SetAttribute("server.address", host); + Span_->SetAttribute("server.port", static_cast(port)); +} + +TQuerySpan::~TQuerySpan() { + if (Span_) { + Span_->End(); + } +} + +void TQuerySpan::End(EStatus status) { + if (Span_) { + Span_->SetAttribute("db.response.status_code", static_cast(status)); + if (status != EStatus::SUCCESS) { + Span_->SetAttribute("error.type", ToString(status)); + } + Span_->End(); + Span_.reset(); + } +} + +} // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h new file mode 100644 index 00000000000..d34e263d4db --- /dev/null +++ b/src/client/query/impl/query_spans.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include +#include + +#include +#include + +namespace NYdb::inline V3::NQuery { + +class TQuerySpan { +public: + TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); + ~TQuerySpan(); + + void End(EStatus status); + +private: + std::shared_ptr Span_; +}; + +} // namespace NYdb::NQuery diff --git a/src/open_telemetry/CMakeLists.txt b/src/open_telemetry/CMakeLists.txt new file mode 100644 index 00000000000..4b0d7df8102 --- /dev/null +++ b/src/open_telemetry/CMakeLists.txt @@ -0,0 +1,17 @@ +_ydb_sdk_add_library(open_telemetry) + +target_sources(open_telemetry PRIVATE + otel.cpp +) + +target_include_directories(open_telemetry PUBLIC + $ + $ +) + +target_link_libraries(open_telemetry PUBLIC + client-metrics + opentelemetry-cpp::api + opentelemetry-cpp::metrics + opentelemetry-cpp::trace +) diff --git a/src/open_telemetry/otel.cpp b/src/open_telemetry/otel.cpp new file mode 100644 index 00000000000..b12a08c1b7b --- /dev/null +++ b/src/open_telemetry/otel.cpp @@ -0,0 +1,153 @@ +#include + +#include +#include +#include +#include +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +common::KeyValueIterableView MakeAttributes(const TLabels& labels) { + return common::KeyValueIterableView(labels); +} + +class TOtelCounter : public ICounter { +public: + TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Inc() override { + Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; +}; + +class TOtelUpDownCounterGauge : public IGauge { +public: + TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Add(double delta) override { + Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ += delta; + } + + void Set(double value) override { + Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ = value; + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; + double Value_ = 0; +}; + +class TOtelHistogram : public IHistogram { +public: + TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) + : Histogram_(std::move(histogram)) + , Labels_(labels) + {} + + void Record(double value) override { + Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Histogram_; + TLabels Labels_; +}; + +trace::SpanKind MapSpanKind(ESpanKind kind) { + switch (kind) { + case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; + case ESpanKind::SERVER: return trace::SpanKind::kServer; + case ESpanKind::CLIENT: return trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + } + return trace::SpanKind::kInternal; +} + +class TOtelSpan : public ISpan { +public: + TOtelSpan(nostd::shared_ptr span) + : Span_(std::move(span)) + {} + + void End() override { + Span_->End(); + } + + void SetAttribute(const std::string& key, const std::string& value) override { + Span_->SetAttribute(key, value); + } + + void SetAttribute(const std::string& key, int64_t value) override { + Span_->SetAttribute(key, value); + } + +private: + nostd::shared_ptr Span_; +}; + +class TOtelTracer : public ITracer { +public: + TOtelTracer(nostd::shared_ptr tracer) + : Tracer_(std::move(tracer)) + {} + + std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { + trace::StartSpanOptions options; + options.kind = MapSpanKind(kind); + return std::make_shared(Tracer_->StartSpan(name, options)); + } + +private: + nostd::shared_ptr Tracer_; +}; + +} // namespace + +TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) + : MeterProvider_(std::move(meterProvider)) + , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", "1.0.0")) +{} + +std::shared_ptr TOtelMetricRegistry::Counter(const std::string& name, const TLabels& labels) { + auto counter = Meter_->CreateUInt64Counter(name); + return std::make_shared(std::move(counter), labels); +} + +std::shared_ptr TOtelMetricRegistry::Gauge(const std::string& name, const TLabels& labels) { + auto counter = Meter_->CreateDoubleUpDownCounter(name); + return std::make_shared(std::move(counter), labels); +} + +std::shared_ptr TOtelMetricRegistry::Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) { + auto histogram = Meter_->CreateDoubleHistogram(name); + return std::make_shared(std::move(histogram), labels); +} + +TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) + : TracerProvider_(std::move(tracerProvider)) +{} + +std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { + return std::make_shared(TracerProvider_->GetTracer(name)); +} + +} // namespace NYdb::NMetrics From 90118388d09338f83871caf64fff96d5430507ca Mon Sep 17 00:00:00 2001 From: maladetska Date: Sat, 21 Feb 2026 15:57:37 +0300 Subject: [PATCH 02/93] . --- CMakeLists.txt | 4 ++++ cmake/external_libs.cmake | 5 ++++- {src/open_telemetry => open_telemetry}/CMakeLists.txt | 7 +++++-- .../include}/ydb-cpp-sdk/open_telemetry/extension.h | 0 .../include}/ydb-cpp-sdk/open_telemetry/otel.h | 0 {src/open_telemetry => open_telemetry/src}/otel.cpp | 0 src/CMakeLists.txt | 1 - 7 files changed, 13 insertions(+), 4 deletions(-) rename {src/open_telemetry => open_telemetry}/CMakeLists.txt (59%) rename {include => open_telemetry/include}/ydb-cpp-sdk/open_telemetry/extension.h (100%) rename {include => open_telemetry/include}/ydb-cpp-sdk/open_telemetry/otel.h (100%) rename {src/open_telemetry => open_telemetry/src}/otel.cpp (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 86eaef64720..991c12d8a93 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -59,6 +59,10 @@ add_subdirectory(include/ydb-cpp-sdk/client) add_subdirectory(src) add_subdirectory(util) +if (YDB_SDK_HAS_OPEN_TELEMETRY) + add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) +endif() + #_ydb_sdk_validate_public_headers() if (YDB_SDK_EXAMPLES) diff --git a/cmake/external_libs.cmake b/cmake/external_libs.cmake index be645f332ae..76666562f06 100644 --- a/cmake/external_libs.cmake +++ b/cmake/external_libs.cmake @@ -13,7 +13,10 @@ find_package(base64 REQUIRED) find_package(Brotli 1.1.0 REQUIRED) find_package(jwt-cpp REQUIRED) find_package(double-conversion REQUIRED) -find_package(opentelemetry-cpp REQUIRED) +find_package(opentelemetry-cpp QUIET) +if (opentelemetry-cpp_FOUND) + set(YDB_SDK_HAS_OPEN_TELEMETRY ON) +endif() # RapidJSON if (YDB_SDK_USE_RAPID_JSON) diff --git a/src/open_telemetry/CMakeLists.txt b/open_telemetry/CMakeLists.txt similarity index 59% rename from src/open_telemetry/CMakeLists.txt rename to open_telemetry/CMakeLists.txt index 4b0d7df8102..aa0cee855e1 100644 --- a/src/open_telemetry/CMakeLists.txt +++ b/open_telemetry/CMakeLists.txt @@ -1,11 +1,11 @@ _ydb_sdk_add_library(open_telemetry) target_sources(open_telemetry PRIVATE - otel.cpp + src/otel.cpp ) target_include_directories(open_telemetry PUBLIC - $ + $ $ ) @@ -15,3 +15,6 @@ target_link_libraries(open_telemetry PUBLIC opentelemetry-cpp::metrics opentelemetry-cpp::trace ) + +_ydb_sdk_make_client_component(OpenTelemetry open_telemetry) +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/include/ydb-cpp-sdk/open_telemetry/extension.h b/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h similarity index 100% rename from include/ydb-cpp-sdk/open_telemetry/extension.h rename to open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h diff --git a/include/ydb-cpp-sdk/open_telemetry/otel.h b/open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h similarity index 100% rename from include/ydb-cpp-sdk/open_telemetry/otel.h rename to open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h diff --git a/src/open_telemetry/otel.cpp b/open_telemetry/src/otel.cpp similarity index 100% rename from src/open_telemetry/otel.cpp rename to open_telemetry/src/otel.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c3fa4e733cd..b251a041380 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,4 +1,3 @@ add_subdirectory(api) add_subdirectory(client) add_subdirectory(library) -add_subdirectory(open_telemetry) From bbc803b1aa7e9261e0b32177bec9309886761cb6 Mon Sep 17 00:00:00 2001 From: maladetska Date: Tue, 24 Feb 2026 14:43:45 +0300 Subject: [PATCH 03/93] make metrics as client_component --- src/client/metrics/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index 94902b0f415..ce8526b5e7e 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -12,3 +12,5 @@ target_include_directories(client-metrics PUBLIC target_link_libraries(client-metrics PUBLIC client-extension_common ) + +_ydb_sdk_make_client_component(Metrics client-metrics) From feb3e8dc461d1b498cd03a719a66d08775d8bfbb Mon Sep 17 00:00:00 2001 From: maladetska Date: Sun, 1 Mar 2026 02:20:51 +0300 Subject: [PATCH 04/93] fixes --- CMakeLists.txt | 7 +- cmake/external_libs.cmake | 9 +- include/ydb-cpp-sdk/client/driver/driver.h | 13 +++ include/ydb-cpp-sdk/client/metrics/metrics.h | 16 --- open_telemetry/CMakeLists.txt | 20 ---- .../ydb-cpp-sdk/open_telemetry/extension.h | 34 ------ plugins/CMakeLists.txt | 3 + plugins/open_telemetry/CMakeLists.txt | 36 ++++++ .../ydb-cpp-sdk/open_telemetry/extension.h | 16 +++ .../ydb-cpp-sdk/open_telemetry/metrics.h | 23 ++-- .../ydb-cpp-sdk/open_telemetry/trace.h | 29 +++++ .../open_telemetry/src/metrics.cpp | 105 ++++++++---------- plugins/open_telemetry/src/trace.cpp | 70 ++++++++++++ src/client/driver/driver.cpp | 24 ++++ .../grpc_connections/grpc_connections.cpp | 10 ++ .../grpc_connections/grpc_connections.h | 9 ++ .../impl/internal/grpc_connections/params.h | 7 ++ src/client/metrics/CMakeLists.txt | 9 -- src/client/metrics/metrics.cpp | 31 ------ src/client/query/client.cpp | 18 +-- src/client/query/impl/query_spans.cpp | 54 +++++++-- src/client/query/impl/query_spans.h | 4 +- 22 files changed, 340 insertions(+), 207 deletions(-) delete mode 100644 open_telemetry/CMakeLists.txt delete mode 100644 open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h create mode 100644 plugins/CMakeLists.txt create mode 100644 plugins/open_telemetry/CMakeLists.txt create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h rename open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h => plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h (62%) create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h rename open_telemetry/src/otel.cpp => plugins/open_telemetry/src/metrics.cpp (61%) create mode 100644 plugins/open_telemetry/src/trace.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 991c12d8a93..6df450c510c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,6 +10,8 @@ project(YDB-CPP-SDK VERSION ${YDB_SDK_VERSION} LANGUAGES C CXX ASM) option(YDB_SDK_INSTALL "Install YDB C++ SDK" Off) option(YDB_SDK_TESTS "Build YDB C++ SDK tests" Off) option(YDB_SDK_EXAMPLES "Build YDB C++ SDK examples" On) +option(YDB_SDK_ENABLE_OTEL_METRICS "Build OpenTelemetry metrics plugin" Off) +option(YDB_SDK_ENABLE_OTEL_TRACE "Build OpenTelemetry trace plugin" Off) set(YDB_SDK_GOOGLE_COMMON_PROTOS_TARGET "" CACHE STRING "Name of cmake target preparing google common proto library") option(YDB_SDK_USE_RAPID_JSON "Search for rapid json library in system" ON) @@ -58,10 +60,7 @@ add_subdirectory(library/cpp) add_subdirectory(include/ydb-cpp-sdk/client) add_subdirectory(src) add_subdirectory(util) - -if (YDB_SDK_HAS_OPEN_TELEMETRY) - add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) -endif() +add_subdirectory(plugins) #_ydb_sdk_validate_public_headers() diff --git a/cmake/external_libs.cmake b/cmake/external_libs.cmake index 76666562f06..8445e4d2fc1 100644 --- a/cmake/external_libs.cmake +++ b/cmake/external_libs.cmake @@ -13,9 +13,12 @@ find_package(base64 REQUIRED) find_package(Brotli 1.1.0 REQUIRED) find_package(jwt-cpp REQUIRED) find_package(double-conversion REQUIRED) -find_package(opentelemetry-cpp QUIET) -if (opentelemetry-cpp_FOUND) - set(YDB_SDK_HAS_OPEN_TELEMETRY ON) + +if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) + find_package(opentelemetry-cpp QUIET) + if (NOT opentelemetry-cpp_FOUND) + message(FATAL_ERROR "Dependency 'opentelemetry-cpp' was not found.") + endif() endif() # RapidJSON diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index 72aa008ccca..8f373e9b25c 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -153,6 +154,18 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); + //! Set external metrics exporter implementation. + TDriverConfig& SetMetricExporter(std::shared_ptr exporter); + + //! Set external tracing exporter implementation. + TDriverConfig& SetTraceExporter(std::shared_ptr exporter); + + //! Get configured metrics exporter implementation. + std::shared_ptr GetMetricExporter() const; + + //! Get configured tracing exporter implementation. + std::shared_ptr GetTraceExporter() const; + private: class TImpl; std::shared_ptr Impl_; diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index d3dc5eb2c19..0fce1f081bf 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -1,11 +1,6 @@ #pragma once -#include - #include -#include -#include -#include namespace NYdb::inline V3::NMetrics { @@ -67,15 +62,4 @@ class ITraceProvider { virtual std::shared_ptr GetTracer(const std::string& name) = 0; }; -class IMetricsApi : public IExtensionApi { -public: - static IMetricsApi* Create(TDriver driver); -public: - virtual ~IMetricsApi() = default; - virtual void SetMetricRegistry(std::shared_ptr registry) = 0; - virtual void SetTraceProvider(std::shared_ptr provider) = 0; - virtual std::shared_ptr GetMetricRegistry() const = 0; - virtual std::shared_ptr GetTraceProvider() const = 0; -}; - } // namespace NYdb::NMetrics diff --git a/open_telemetry/CMakeLists.txt b/open_telemetry/CMakeLists.txt deleted file mode 100644 index aa0cee855e1..00000000000 --- a/open_telemetry/CMakeLists.txt +++ /dev/null @@ -1,20 +0,0 @@ -_ydb_sdk_add_library(open_telemetry) - -target_sources(open_telemetry PRIVATE - src/otel.cpp -) - -target_include_directories(open_telemetry PUBLIC - $ - $ -) - -target_link_libraries(open_telemetry PUBLIC - client-metrics - opentelemetry-cpp::api - opentelemetry-cpp::metrics - opentelemetry-cpp::trace -) - -_ydb_sdk_make_client_component(OpenTelemetry open_telemetry) -_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h deleted file mode 100644 index b5683d6d41b..00000000000 --- a/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelExtension : public IExtension { -public: - using IApi = IMetricsApi; - - struct TParams { - opentelemetry::nostd::shared_ptr MeterProvider; - opentelemetry::nostd::shared_ptr TracerProvider; - }; - - TOtelExtension(const TParams& params, IApi* api) { - if (params.MeterProvider) { - api->SetMetricRegistry(std::make_shared(params.MeterProvider)); - } - if (params.TracerProvider) { - api->SetTraceProvider(std::make_shared(params.TracerProvider)); - } - } -}; - -inline void AddOpenTelemetry(TDriver& driver - , opentelemetry::nostd::shared_ptr meterProvider - , opentelemetry::nostd::shared_ptr tracerProvider -) { - driver.AddExtension({meterProvider, tracerProvider}); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt new file mode 100644 index 00000000000..836ed903825 --- /dev/null +++ b/plugins/CMakeLists.txt @@ -0,0 +1,3 @@ +if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) + add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) +endif() diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt new file mode 100644 index 00000000000..d005708d422 --- /dev/null +++ b/plugins/open_telemetry/CMakeLists.txt @@ -0,0 +1,36 @@ +if (YDB_SDK_ENABLE_OTEL_METRICS) + _ydb_sdk_add_library(open_telemetry_metrics) + target_sources(open_telemetry_metrics PRIVATE + src/metrics.cpp + ) + target_include_directories(open_telemetry_metrics PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_metrics PUBLIC + client-metrics + client-resources + opentelemetry-cpp::api + opentelemetry-cpp::metrics + ) + _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) +endif() + +if (YDB_SDK_ENABLE_OTEL_TRACE) + _ydb_sdk_add_library(open_telemetry_trace) + target_sources(open_telemetry_trace PRIVATE + src/trace.cpp + ) + target_include_directories(open_telemetry_trace PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_trace PUBLIC + client-metrics + opentelemetry-cpp::api + opentelemetry-cpp::trace + ) + _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) +endif() + +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h new file mode 100644 index 00000000000..b0db9ea7d7c --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace NYdb::inline V3::NMetrics { + +inline void AddOpenTelemetry(TDriverConfig& config + , opentelemetry::nostd::shared_ptr meterProvider + , opentelemetry::nostd::shared_ptr tracerProvider +) { + AddOpenTelemetryMetrics(config, std::move(meterProvider)); + AddOpenTelemetryTrace(config, std::move(tracerProvider)); +} + +} // namespace NYdb::NMetrics diff --git a/open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h similarity index 62% rename from open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h rename to plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h index b1d23c92fa6..5e9e9e77dea 100644 --- a/open_telemetry/include/ydb-cpp-sdk/open_telemetry/otel.h +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -1,9 +1,9 @@ #pragma once +#include #include #include -#include namespace NYdb::inline V3::NMetrics { @@ -16,18 +16,21 @@ class TOtelMetricRegistry : public IMetricRegistry { std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) override; private: + void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets); + opentelemetry::nostd::shared_ptr MeterProvider_; opentelemetry::nostd::shared_ptr Meter_; + std::mutex HistogramViewsLock_; + std::unordered_set HistogramViews_; }; -class TOtelTraceProvider : public ITraceProvider { -public: - TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); - - std::shared_ptr GetTracer(const std::string& name) override; - -private: - opentelemetry::nostd::shared_ptr TracerProvider_; -}; +inline void AddOpenTelemetryMetrics( + TDriverConfig& config, + opentelemetry::nostd::shared_ptr meterProvider) +{ + if (meterProvider) { + config.SetMetricExporter(std::make_shared(std::move(meterProvider))); + } +} } // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h new file mode 100644 index 00000000000..3ba2e146fd9 --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -0,0 +1,29 @@ +#pragma once + +#include +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelTraceProvider : public ITraceProvider { +public: + TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); + + std::shared_ptr GetTracer(const std::string& name) override; + +private: + opentelemetry::nostd::shared_ptr TracerProvider_; +}; + +inline void AddOpenTelemetryTrace( + TDriverConfig& config, + opentelemetry::nostd::shared_ptr tracerProvider) +{ + if (tracerProvider) { + config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); + } +} + +} // namespace NYdb::NMetrics diff --git a/open_telemetry/src/otel.cpp b/plugins/open_telemetry/src/metrics.cpp similarity index 61% rename from open_telemetry/src/otel.cpp rename to plugins/open_telemetry/src/metrics.cpp index b12a08c1b7b..65850fd08b3 100644 --- a/open_telemetry/src/otel.cpp +++ b/plugins/open_telemetry/src/metrics.cpp @@ -1,10 +1,12 @@ -#include +#include +#include -#include -#include #include -#include #include +#include +#include +#include +#include namespace NYdb::inline V3::NMetrics { @@ -71,61 +73,53 @@ class TOtelHistogram : public IHistogram { TLabels Labels_; }; -trace::SpanKind MapSpanKind(ESpanKind kind) { - switch (kind) { - case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; - case ESpanKind::SERVER: return trace::SpanKind::kServer; - case ESpanKind::CLIENT: return trace::SpanKind::kClient; - case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; - case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; - } - return trace::SpanKind::kInternal; -} +} // namespace -class TOtelSpan : public ISpan { -public: - TOtelSpan(nostd::shared_ptr span) - : Span_(std::move(span)) - {} +TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) + : MeterProvider_(std::move(meterProvider)) + , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) +{} - void End() override { - Span_->End(); +void TOtelMetricRegistry::ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { + if (buckets.empty()) { + return; } - void SetAttribute(const std::string& key, const std::string& value) override { - Span_->SetAttribute(key, value); + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); + if (!sdkProvider) { + return; } - void SetAttribute(const std::string& key, int64_t value) override { - Span_->SetAttribute(key, value); + { + std::lock_guard lock(HistogramViewsLock_); + if (!HistogramViews_.insert(name).second) { + return; + } } -private: - nostd::shared_ptr Span_; -}; - -class TOtelTracer : public ITracer { -public: - TOtelTracer(nostd::shared_ptr tracer) - : Tracer_(std::move(tracer)) - {} - - std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { - trace::StartSpanOptions options; - options.kind = MapSpanKind(kind); - return std::make_shared(Tracer_->StartSpan(name, options)); - } - -private: - nostd::shared_ptr Tracer_; -}; - -} // namespace - -TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) - : MeterProvider_(std::move(meterProvider)) - , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", "1.0.0")) -{} + auto selector = std::make_unique( + sdk::metrics::InstrumentType::kHistogram, + name, + "" + ); + auto meterSelector = std::make_unique( + "ydb-cpp-sdk", + GetSdkSemver(), + {} + ); + + auto histogramConfig = std::make_shared(); + histogramConfig->boundaries_ = buckets; + + auto view = std::make_unique( + {}, + {}, + sdk::metrics::AggregationType::kHistogram, + histogramConfig + ); + + sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); +} std::shared_ptr TOtelMetricRegistry::Counter(const std::string& name, const TLabels& labels) { auto counter = Meter_->CreateUInt64Counter(name); @@ -138,16 +132,9 @@ std::shared_ptr TOtelMetricRegistry::Gauge(const std::string& name, cons } std::shared_ptr TOtelMetricRegistry::Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) { + ConfigureHistogramBuckets(name, buckets); auto histogram = Meter_->CreateDoubleHistogram(name); return std::make_shared(std::move(histogram), labels); } -TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) - : TracerProvider_(std::move(tracerProvider)) -{} - -std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { - return std::make_shared(TracerProvider_->GetTracer(name)); -} - } // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/open_telemetry/src/trace.cpp new file mode 100644 index 00000000000..54f04cb84df --- /dev/null +++ b/plugins/open_telemetry/src/trace.cpp @@ -0,0 +1,70 @@ +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +trace::SpanKind MapSpanKind(ESpanKind kind) { + switch (kind) { + case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; + case ESpanKind::SERVER: return trace::SpanKind::kServer; + case ESpanKind::CLIENT: return trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + } + return trace::SpanKind::kInternal; +} + +class TOtelSpan : public ISpan { +public: + TOtelSpan(nostd::shared_ptr span) + : Span_(std::move(span)) + {} + + void End() override { + Span_->End(); + } + + void SetAttribute(const std::string& key, const std::string& value) override { + Span_->SetAttribute(key, value); + } + + void SetAttribute(const std::string& key, int64_t value) override { + Span_->SetAttribute(key, value); + } + +private: + nostd::shared_ptr Span_; +}; + +class TOtelTracer : public ITracer { +public: + TOtelTracer(nostd::shared_ptr tracer) + : Tracer_(std::move(tracer)) + {} + + std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { + trace::StartSpanOptions options; + options.kind = MapSpanKind(kind); + return std::make_shared(Tracer_->StartSpan(name, options)); + } + +private: + nostd::shared_ptr Tracer_; +}; + +} // namespace + +TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) + : TracerProvider_(std::move(tracerProvider)) +{} + +std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { + return std::make_shared(TracerProvider_->GetTracer(name)); +} + +} // namespace NYdb::NMetrics diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 207c67b6d5f..63aeac8aea5 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -51,6 +51,8 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t GetMaxMessageSize() const override { return MaxMessageSize; } const TLog& GetLog() const override { return Log; } std::shared_ptr GetExecutor() const override { return Executor; } + std::shared_ptr GetMetricExporter() const override { return MetricExporter; } + std::shared_ptr GetTraceExporter() const override { return TraceExporter; } std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -80,6 +82,8 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t MaxMessageSize = 0; TLog Log; // Null by default. std::shared_ptr Executor; + std::shared_ptr MetricExporter; + std::shared_ptr TraceExporter; }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -229,6 +233,24 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { return *this; } +TDriverConfig& TDriverConfig::SetMetricExporter(std::shared_ptr exporter) { + Impl_->MetricExporter = std::move(exporter); + return *this; +} + +TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { + Impl_->TraceExporter = std::move(exporter); + return *this; +} + +std::shared_ptr TDriverConfig::GetMetricExporter() const { + return Impl_->MetricExporter; +} + +std::shared_ptr TDriverConfig::GetTraceExporter() const { + return Impl_->TraceExporter; +} + //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -280,6 +302,8 @@ TDriverConfig TDriver::GetConfig() const { config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); config.SetMaxMessageSize(Impl_->MaxMessageSize_); config.Impl_->Log = Impl_->Log; + config.SetMetricExporter(Impl_->GetMetricExporter()); + config.SetTraceExporter(Impl_->GetTraceExporter()); return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index 48e170d28c6..32645964933 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -167,6 +167,8 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout()) #endif + , MetricExporter_(params->GetMetricExporter()) + , TraceExporter_(params->GetTraceExporter()) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -434,6 +436,14 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { ExtensionApis_.emplace_back(api); } +std::shared_ptr TGRpcConnectionsImpl::GetMetricExporter() const { + return MetricExporter_; +} + +std::shared_ptr TGRpcConnectionsImpl::GetTraceExporter() const { + return TraceExporter_; +} + void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { std::lock_guard lock(ExtensionsLock_); DiscoveryMutatorCb = std::move(cb); diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 15cddb4e6ec..14db28028d2 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -18,6 +18,11 @@ namespace NYdb::inline V3 { +namespace NMetrics { + class IMetricRegistry; + class ITraceProvider; +} // namespace NMetrics + constexpr TDeadline::Duration GRPC_KEEP_ALIVE_TIMEOUT_FOR_DISCOVERY = std::chrono::seconds(10); constexpr TDeadline::Duration INITIAL_DEFERRED_CALL_DELAY = std::chrono::milliseconds(10); // The delay before first deferred service call constexpr TDeadline::Duration GET_ENDPOINTS_TIMEOUT = std::chrono::seconds(10); // Time wait for ListEndpoints request, after this time we pass error to client @@ -581,6 +586,8 @@ class TGRpcConnectionsImpl ::NMonitoring::TMetricRegistry* GetMetricRegistry() override; void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); + std::shared_ptr GetMetricExporter() const; + std::shared_ptr GetTraceExporter() const; template T* GetExtensionApi() { @@ -726,6 +733,8 @@ class TGRpcConnectionsImpl std::vector> Extensions_; std::vector> ExtensionApis_; + std::shared_ptr MetricExporter_; + std::shared_ptr TraceExporter_; IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 2bc9f4567c5..e6aa5e87dde 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -11,6 +11,11 @@ namespace NYdb::inline V3 { +namespace NMetrics { + class IMetricRegistry; + class ITraceProvider; +} // namespace NMetrics + class IConnectionsParams { public: virtual ~IConnectionsParams() = default; @@ -36,6 +41,8 @@ class IConnectionsParams { virtual uint64_t GetMaxOutboundMessageSize() const = 0; virtual uint64_t GetMaxMessageSize() const = 0; virtual std::shared_ptr GetExecutor() const = 0; + virtual std::shared_ptr GetMetricExporter() const = 0; + virtual std::shared_ptr GetTraceExporter() const = 0; }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index ce8526b5e7e..03915de6b5e 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -4,13 +4,4 @@ target_sources(client-metrics PRIVATE metrics.cpp ) -target_include_directories(client-metrics PUBLIC - $ - $ -) - -target_link_libraries(client-metrics PUBLIC - client-extension_common -) - _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/metrics/metrics.cpp b/src/client/metrics/metrics.cpp index 836d01f2071..341917291bb 100644 --- a/src/client/metrics/metrics.cpp +++ b/src/client/metrics/metrics.cpp @@ -1,32 +1 @@ #include - -namespace NYdb::inline V3::NMetrics { - -class TMetricsApiImpl : public IMetricsApi { -public: - void SetMetricRegistry(std::shared_ptr registry) override { - Registry_ = std::move(registry); - } - - void SetTraceProvider(std::shared_ptr provider) override { - TraceProvider_ = std::move(provider); - } - - std::shared_ptr GetMetricRegistry() const override { - return Registry_; - } - - std::shared_ptr GetTraceProvider() const override { - return TraceProvider_; - } - -private: - std::shared_ptr Registry_; - std::shared_ptr TraceProvider_; -}; - -IMetricsApi* IMetricsApi::Create(TDriver driver) { - return new TMetricsApiImpl(); -} - -} // namespace NYdb::NMetrics diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index ca1453e0faf..80a3d6ed00f 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -70,10 +70,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); - if (auto metricsApi = Connections_->GetExtensionApi()) { - if (auto traceProvider = metricsApi->GetTraceProvider()) { - Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); - } + if (auto traceProvider = Connections_->GetTraceExporter()) { + Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } } @@ -471,7 +469,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public void ReplyError(TStatus status) override { TSession session; - if (Span) Span->End(status.GetStatus()); + if (Span) { + Span->End(status.GetStatus()); + } ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -484,7 +484,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public ) ); - if (Span) Span->End(EStatus::SUCCESS); + if (Span) { + Span->End(EStatus::SUCCESS); + } ScheduleReply(std::move(val)); } @@ -493,7 +495,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public [promise{std::move(Promise)}, span = Span](TAsyncCreateSessionResult future) mutable { auto val = future.ExtractValue(); - if (span) span->End(val.GetStatus()); + if (span) { + span->End(val.GetStatus()); + } promise.SetValue(std::move(val)); }); } diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index c612f442663..fc1329aacd0 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -21,34 +21,64 @@ void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { } } +void SafeLogSpanError(const char* message) noexcept { + try { + try { + Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + return; + } catch (...) { + } + Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + } catch (...) { + } +} + } // namespace TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { - if (!tracer) return; + if (!tracer) { + return; + } std::string host; int port; ParseEndpoint(endpoint, host, port); - Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); - Span_->SetAttribute("db.system.name", "ydb"); - Span_->SetAttribute("server.address", host); - Span_->SetAttribute("server.port", static_cast(port)); + try { + Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); + if (!Span_) { + return; + } + Span_->SetAttribute("db.system.name", "ydb"); + Span_->SetAttribute("server.address", host); + Span_->SetAttribute("server.port", static_cast(port)); + } catch (...) { + SafeLogSpanError("failed to initialize span"); + Span_.reset(); + } } -TQuerySpan::~TQuerySpan() { +TQuerySpan::~TQuerySpan() noexcept { if (Span_) { - Span_->End(); + try { + Span_->End(); + } catch (...) { + SafeLogSpanError("failed to end span"); + } } } -void TQuerySpan::End(EStatus status) { +void TQuerySpan::End(EStatus status) noexcept { if (Span_) { - Span_->SetAttribute("db.response.status_code", static_cast(status)); - if (status != EStatus::SUCCESS) { - Span_->SetAttribute("error.type", ToString(status)); + try { + Span_->SetAttribute("db.response.status_code", static_cast(status)); + if (status != EStatus::SUCCESS) { + Span_->SetAttribute("error.type", ToString(status)); + } + Span_->End(); + } catch (...) { + SafeLogSpanError("failed to finalize span"); } - Span_->End(); Span_.reset(); } } diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index d34e263d4db..ca0b6853954 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -12,9 +12,9 @@ namespace NYdb::inline V3::NQuery { class TQuerySpan { public: TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); - ~TQuerySpan(); + ~TQuerySpan() noexcept; - void End(EStatus status); + void End(EStatus status) noexcept; private: std::shared_ptr Span_; From 0e31bc16c1158e5598a44b58e79926a84b2142a4 Mon Sep 17 00:00:00 2001 From: maladetska Date: Mon, 2 Mar 2026 15:12:38 +0300 Subject: [PATCH 05/93] add includes --- include/ydb-cpp-sdk/client/metrics/metrics.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index 0fce1f081bf..ea06c28f9a0 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -1,6 +1,10 @@ #pragma once +#include #include +#include +#include +#include namespace NYdb::inline V3::NMetrics { From 23414c1339f8ef5373687b27c4c0e0260a059a9a Mon Sep 17 00:00:00 2001 From: maladetska Date: Mon, 16 Mar 2026 01:55:33 +0800 Subject: [PATCH 06/93] fixes and add metric tests --- include/ydb-cpp-sdk/client/driver/driver.h | 17 +- include/ydb-cpp-sdk/client/metrics/metrics.h | 28 -- include/ydb-cpp-sdk/client/trace/trace.h | 39 +++ plugins/CMakeLists.txt | 5 +- plugins/metrics/CMakeLists.txt | 3 + plugins/metrics/otel/CMakeLists.txt | 17 ++ .../ydb-cpp-sdk/open_telemetry/metrics.h | 16 + plugins/metrics/otel/src/metrics.cpp | 156 ++++++++++ plugins/open_telemetry/CMakeLists.txt | 36 --- .../ydb-cpp-sdk/open_telemetry/extension.h | 16 - .../ydb-cpp-sdk/open_telemetry/metrics.h | 36 --- .../ydb-cpp-sdk/open_telemetry/trace.h | 29 -- plugins/open_telemetry/src/metrics.cpp | 140 --------- plugins/trace/CMakeLists.txt | 3 + plugins/trace/otel/CMakeLists.txt | 16 + .../ydb-cpp-sdk/open_telemetry/trace.h | 16 + .../otel}/src/trace.cpp | 39 ++- src/client/CMakeLists.txt | 1 + src/client/driver/driver.cpp | 28 +- .../grpc_connections/grpc_connections.cpp | 12 +- .../grpc_connections/grpc_connections.h | 19 +- .../impl/internal/grpc_connections/params.h | 4 +- src/client/metrics/CMakeLists.txt | 2 +- src/client/query/CMakeLists.txt | 3 +- src/client/query/client.cpp | 48 ++- src/client/query/impl/CMakeLists.txt | 1 + src/client/query/impl/query_metrics.cpp | 71 +++++ src/client/query/impl/query_metrics.h | 28 ++ src/client/query/impl/query_spans.cpp | 71 ++++- src/client/query/impl/query_spans.h | 7 +- src/client/trace/CMakeLists.txt | 7 + src/client/trace/trace.cpp | 1 + tests/common/fake_metric_registry.h | 122 ++++++++ tests/integration/CMakeLists.txt | 1 + tests/integration/metrics/CMakeLists.txt | 12 + tests/integration/metrics/main.cpp | 273 ++++++++++++++++++ tests/unit/client/CMakeLists.txt | 13 + tests/unit/client/query/query_metrics_ut.cpp | 190 ++++++++++++ 38 files changed, 1162 insertions(+), 364 deletions(-) create mode 100644 include/ydb-cpp-sdk/client/trace/trace.h create mode 100644 plugins/metrics/CMakeLists.txt create mode 100644 plugins/metrics/otel/CMakeLists.txt create mode 100644 plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h create mode 100644 plugins/metrics/otel/src/metrics.cpp delete mode 100644 plugins/open_telemetry/CMakeLists.txt delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h delete mode 100644 plugins/open_telemetry/src/metrics.cpp create mode 100644 plugins/trace/CMakeLists.txt create mode 100644 plugins/trace/otel/CMakeLists.txt create mode 100644 plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h rename plugins/{open_telemetry => trace/otel}/src/trace.cpp (55%) create mode 100644 src/client/query/impl/query_metrics.cpp create mode 100644 src/client/query/impl/query_metrics.h create mode 100644 src/client/trace/CMakeLists.txt create mode 100644 src/client/trace/trace.cpp create mode 100644 tests/common/fake_metric_registry.h create mode 100644 tests/integration/metrics/CMakeLists.txt create mode 100644 tests/integration/metrics/main.cpp create mode 100644 tests/unit/client/query/query_metrics_ut.cpp diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index 8f373e9b25c..20fa52d5e60 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -3,10 +3,11 @@ #include "fwd.h" #include +#include +#include #include #include #include -#include #include #include #include @@ -154,17 +155,11 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); - //! Set external metrics exporter implementation. - TDriverConfig& SetMetricExporter(std::shared_ptr exporter); - - //! Set external tracing exporter implementation. - TDriverConfig& SetTraceExporter(std::shared_ptr exporter); - - //! Get configured metrics exporter implementation. - std::shared_ptr GetMetricExporter() const; + //! Set external metrics registry implementation. + TDriverConfig& SetMetricRegistry(std::shared_ptr registry); - //! Get configured tracing exporter implementation. - std::shared_ptr GetTraceExporter() const; + //! Set external trace provider implementation. + TDriverConfig& SetTraceProvider(std::shared_ptr provider); private: class TImpl; diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index ea06c28f9a0..7e2b0b903dd 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -38,32 +38,4 @@ class IMetricRegistry { virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; }; -enum class ESpanKind { - INTERNAL, - SERVER, - CLIENT, - PRODUCER, - CONSUMER -}; - -class ISpan { -public: - virtual ~ISpan() = default; - virtual void End() = 0; - virtual void SetAttribute(const std::string& key, const std::string& value) = 0; - virtual void SetAttribute(const std::string& key, int64_t value) = 0; -}; - -class ITracer { -public: - virtual ~ITracer() = default; - virtual std::shared_ptr StartSpan(const std::string& name, ESpanKind kind = ESpanKind::INTERNAL) = 0; -}; - -class ITraceProvider { -public: - virtual ~ITraceProvider() = default; - virtual std::shared_ptr GetTracer(const std::string& name) = 0; -}; - } // namespace NYdb::NMetrics diff --git a/include/ydb-cpp-sdk/client/trace/trace.h b/include/ydb-cpp-sdk/client/trace/trace.h new file mode 100644 index 00000000000..b86297146a9 --- /dev/null +++ b/include/ydb-cpp-sdk/client/trace/trace.h @@ -0,0 +1,39 @@ +#pragma once + +#include +#include +#include +#include + +namespace NYdb::inline V3::NMetrics { + +enum class ESpanKind { + INTERNAL, + SERVER, + CLIENT, + PRODUCER, + CONSUMER +}; + +class ISpan { +public: + virtual ~ISpan() = default; + virtual void End() = 0; + virtual void SetAttribute(const std::string& key, const std::string& value) = 0; + virtual void SetAttribute(const std::string& key, int64_t value) = 0; + virtual void AddEvent(const std::string& name, const std::map& attributes = {}) = 0; +}; + +class ITracer { +public: + virtual ~ITracer() = default; + virtual std::shared_ptr StartSpan(const std::string& name, ESpanKind kind = ESpanKind::INTERNAL) = 0; +}; + +class ITraceProvider { +public: + virtual ~ITraceProvider() = default; + virtual std::shared_ptr GetTracer(const std::string& name) = 0; +}; + +} // namespace NYdb::NMetrics diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 836ed903825..0d232800455 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -1,3 +1,2 @@ -if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) - add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) -endif() +add_subdirectory(metrics) +add_subdirectory(trace) diff --git a/plugins/metrics/CMakeLists.txt b/plugins/metrics/CMakeLists.txt new file mode 100644 index 00000000000..6d50a5111e7 --- /dev/null +++ b/plugins/metrics/CMakeLists.txt @@ -0,0 +1,3 @@ +if (YDB_SDK_ENABLE_OTEL_METRICS) + add_subdirectory(otel EXCLUDE_FROM_ALL) +endif() diff --git a/plugins/metrics/otel/CMakeLists.txt b/plugins/metrics/otel/CMakeLists.txt new file mode 100644 index 00000000000..e26b1931984 --- /dev/null +++ b/plugins/metrics/otel/CMakeLists.txt @@ -0,0 +1,17 @@ +_ydb_sdk_add_library(open_telemetry_metrics) +target_sources(open_telemetry_metrics PRIVATE + src/metrics.cpp +) +target_include_directories(open_telemetry_metrics PUBLIC + $ + $ +) +target_link_libraries(open_telemetry_metrics PUBLIC + client-metrics + client-resources + opentelemetry-cpp::api + opentelemetry-cpp::metrics +) +_ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) + +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h new file mode 100644 index 00000000000..054d040f97d --- /dev/null +++ b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +#include + +namespace opentelemetry::metrics { +class MeterProvider; +} + +namespace NYdb::inline V3::NMetrics { + +std::shared_ptr CreateOtelMetricRegistry( + opentelemetry::nostd::shared_ptr meterProvider); + +} // namespace NYdb::NMetrics diff --git a/plugins/metrics/otel/src/metrics.cpp b/plugins/metrics/otel/src/metrics.cpp new file mode 100644 index 00000000000..6b9f14be362 --- /dev/null +++ b/plugins/metrics/otel/src/metrics.cpp @@ -0,0 +1,156 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +common::KeyValueIterableView MakeAttributes(const TLabels& labels) { + return common::KeyValueIterableView(labels); +} + +class TOtelCounter : public ICounter { +public: + TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Inc() override { + Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; +}; + +class TOtelUpDownCounterGauge : public IGauge { +public: + TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Add(double delta) override { + Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ += delta; + } + + void Set(double value) override { + Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ = value; + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; + double Value_ = 0; +}; + +class TOtelHistogram : public IHistogram { +public: + TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) + : Histogram_(std::move(histogram)) + , Labels_(labels) + {} + + void Record(double value) override { + Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Histogram_; + TLabels Labels_; +}; + +class TOtelMetricRegistry : public IMetricRegistry { +public: + TOtelMetricRegistry(nostd::shared_ptr meterProvider) + : MeterProvider_(std::move(meterProvider)) + , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) + {} + + std::shared_ptr Counter(const std::string& name, const TLabels& labels) override { + auto counter = Meter_->CreateUInt64Counter(name); + return std::make_shared(std::move(counter), labels); + } + + std::shared_ptr Gauge(const std::string& name, const TLabels& labels) override { + auto counter = Meter_->CreateDoubleUpDownCounter(name); + return std::make_shared(std::move(counter), labels); + } + + std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) override { + ConfigureHistogramBuckets(name, buckets); + auto histogram = Meter_->CreateDoubleHistogram(name); + return std::make_shared(std::move(histogram), labels); + } + +private: + void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { + if (buckets.empty()) { + return; + } + + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); + if (!sdkProvider) { + return; + } + + { + std::lock_guard lock(HistogramViewsLock_); + if (!HistogramViews_.insert(name).second) { + return; + } + } + + auto selector = std::make_unique( + sdk::metrics::InstrumentType::kHistogram, + name, + "" + ); + auto meterSelector = std::make_unique( + "ydb-cpp-sdk", + GetSdkSemver(), + {} + ); + + auto histogramConfig = std::make_shared(); + histogramConfig->boundaries_ = buckets; + + auto view = std::make_unique( + {}, + {}, + sdk::metrics::AggregationType::kHistogram, + histogramConfig + ); + + sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); + } + + nostd::shared_ptr MeterProvider_; + nostd::shared_ptr Meter_; + std::mutex HistogramViewsLock_; + std::unordered_set HistogramViews_; +}; + +} // namespace + +std::shared_ptr CreateOtelMetricRegistry( + opentelemetry::nostd::shared_ptr meterProvider) +{ + return std::make_shared(std::move(meterProvider)); +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt deleted file mode 100644 index d005708d422..00000000000 --- a/plugins/open_telemetry/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -if (YDB_SDK_ENABLE_OTEL_METRICS) - _ydb_sdk_add_library(open_telemetry_metrics) - target_sources(open_telemetry_metrics PRIVATE - src/metrics.cpp - ) - target_include_directories(open_telemetry_metrics PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_metrics PUBLIC - client-metrics - client-resources - opentelemetry-cpp::api - opentelemetry-cpp::metrics - ) - _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) -endif() - -if (YDB_SDK_ENABLE_OTEL_TRACE) - _ydb_sdk_add_library(open_telemetry_trace) - target_sources(open_telemetry_trace PRIVATE - src/trace.cpp - ) - target_include_directories(open_telemetry_trace PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_trace PUBLIC - client-metrics - opentelemetry-cpp::api - opentelemetry-cpp::trace - ) - _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) -endif() - -_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h deleted file mode 100644 index b0db9ea7d7c..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include -#include - -namespace NYdb::inline V3::NMetrics { - -inline void AddOpenTelemetry(TDriverConfig& config - , opentelemetry::nostd::shared_ptr meterProvider - , opentelemetry::nostd::shared_ptr tracerProvider -) { - AddOpenTelemetryMetrics(config, std::move(meterProvider)); - AddOpenTelemetryTrace(config, std::move(tracerProvider)); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h deleted file mode 100644 index 5e9e9e77dea..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelMetricRegistry : public IMetricRegistry { -public: - TOtelMetricRegistry(opentelemetry::nostd::shared_ptr meterProvider); - - std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) override; - std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) override; - std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) override; - -private: - void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets); - - opentelemetry::nostd::shared_ptr MeterProvider_; - opentelemetry::nostd::shared_ptr Meter_; - std::mutex HistogramViewsLock_; - std::unordered_set HistogramViews_; -}; - -inline void AddOpenTelemetryMetrics( - TDriverConfig& config, - opentelemetry::nostd::shared_ptr meterProvider) -{ - if (meterProvider) { - config.SetMetricExporter(std::make_shared(std::move(meterProvider))); - } -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h deleted file mode 100644 index 3ba2e146fd9..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelTraceProvider : public ITraceProvider { -public: - TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); - - std::shared_ptr GetTracer(const std::string& name) override; - -private: - opentelemetry::nostd::shared_ptr TracerProvider_; -}; - -inline void AddOpenTelemetryTrace( - TDriverConfig& config, - opentelemetry::nostd::shared_ptr tracerProvider) -{ - if (tracerProvider) { - config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); - } -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/metrics.cpp b/plugins/open_telemetry/src/metrics.cpp deleted file mode 100644 index 65850fd08b3..00000000000 --- a/plugins/open_telemetry/src/metrics.cpp +++ /dev/null @@ -1,140 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include -#include - -namespace NYdb::inline V3::NMetrics { - -namespace { - -using namespace opentelemetry; - -common::KeyValueIterableView MakeAttributes(const TLabels& labels) { - return common::KeyValueIterableView(labels); -} - -class TOtelCounter : public ICounter { -public: - TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) - : Counter_(std::move(counter)) - , Labels_(labels) - {} - - void Inc() override { - Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - } - -private: - nostd::shared_ptr> Counter_; - TLabels Labels_; -}; - -class TOtelUpDownCounterGauge : public IGauge { -public: - TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) - : Counter_(std::move(counter)) - , Labels_(labels) - {} - - void Add(double delta) override { - Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - Value_ += delta; - } - - void Set(double value) override { - Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - Value_ = value; - } - -private: - nostd::shared_ptr> Counter_; - TLabels Labels_; - double Value_ = 0; -}; - -class TOtelHistogram : public IHistogram { -public: - TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) - : Histogram_(std::move(histogram)) - , Labels_(labels) - {} - - void Record(double value) override { - Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - } - -private: - nostd::shared_ptr> Histogram_; - TLabels Labels_; -}; - -} // namespace - -TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) - : MeterProvider_(std::move(meterProvider)) - , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) -{} - -void TOtelMetricRegistry::ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { - if (buckets.empty()) { - return; - } - - auto* sdkProvider = dynamic_cast(MeterProvider_.get()); - if (!sdkProvider) { - return; - } - - { - std::lock_guard lock(HistogramViewsLock_); - if (!HistogramViews_.insert(name).second) { - return; - } - } - - auto selector = std::make_unique( - sdk::metrics::InstrumentType::kHistogram, - name, - "" - ); - auto meterSelector = std::make_unique( - "ydb-cpp-sdk", - GetSdkSemver(), - {} - ); - - auto histogramConfig = std::make_shared(); - histogramConfig->boundaries_ = buckets; - - auto view = std::make_unique( - {}, - {}, - sdk::metrics::AggregationType::kHistogram, - histogramConfig - ); - - sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); -} - -std::shared_ptr TOtelMetricRegistry::Counter(const std::string& name, const TLabels& labels) { - auto counter = Meter_->CreateUInt64Counter(name); - return std::make_shared(std::move(counter), labels); -} - -std::shared_ptr TOtelMetricRegistry::Gauge(const std::string& name, const TLabels& labels) { - auto counter = Meter_->CreateDoubleUpDownCounter(name); - return std::make_shared(std::move(counter), labels); -} - -std::shared_ptr TOtelMetricRegistry::Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) { - ConfigureHistogramBuckets(name, buckets); - auto histogram = Meter_->CreateDoubleHistogram(name); - return std::make_shared(std::move(histogram), labels); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/trace/CMakeLists.txt b/plugins/trace/CMakeLists.txt new file mode 100644 index 00000000000..ef231ab7103 --- /dev/null +++ b/plugins/trace/CMakeLists.txt @@ -0,0 +1,3 @@ +if (YDB_SDK_ENABLE_OTEL_TRACE) + add_subdirectory(otel EXCLUDE_FROM_ALL) +endif() diff --git a/plugins/trace/otel/CMakeLists.txt b/plugins/trace/otel/CMakeLists.txt new file mode 100644 index 00000000000..6816d8ff7c6 --- /dev/null +++ b/plugins/trace/otel/CMakeLists.txt @@ -0,0 +1,16 @@ +_ydb_sdk_add_library(open_telemetry_trace) +target_sources(open_telemetry_trace PRIVATE + src/trace.cpp +) +target_include_directories(open_telemetry_trace PUBLIC + $ + $ +) +target_link_libraries(open_telemetry_trace PUBLIC + client-trace + opentelemetry-cpp::api + opentelemetry-cpp::trace +) +_ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) + +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h new file mode 100644 index 00000000000..9bdc12fb25f --- /dev/null +++ b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +#include + +namespace opentelemetry::trace { +class TracerProvider; +} + +namespace NYdb::inline V3::NMetrics { + +std::shared_ptr CreateOtelTraceProvider( + opentelemetry::nostd::shared_ptr tracerProvider); + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/trace/otel/src/trace.cpp similarity index 55% rename from plugins/open_telemetry/src/trace.cpp rename to plugins/trace/otel/src/trace.cpp index 54f04cb84df..7cac3f4c1cb 100644 --- a/plugins/open_telemetry/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -1,6 +1,8 @@ #include +#include #include +#include namespace NYdb::inline V3::NMetrics { @@ -37,6 +39,19 @@ class TOtelSpan : public ISpan { Span_->SetAttribute(key, value); } + void AddEvent(const std::string& name, const std::map& attributes) override { + if (attributes.empty()) { + Span_->AddEvent(name); + } else { + std::vector> attrs; + attrs.reserve(attributes.size()); + for (const auto& [k, v] : attributes) { + attrs.emplace_back(nostd::string_view(k), common::AttributeValue(nostd::string_view(v))); + } + Span_->AddEvent(name, attrs); + } + } + private: nostd::shared_ptr Span_; }; @@ -57,14 +72,26 @@ class TOtelTracer : public ITracer { nostd::shared_ptr Tracer_; }; -} // namespace +class TOtelTraceProvider : public ITraceProvider { +public: + TOtelTraceProvider(nostd::shared_ptr tracerProvider) + : TracerProvider_(std::move(tracerProvider)) + {} -TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) - : TracerProvider_(std::move(tracerProvider)) -{} + std::shared_ptr GetTracer(const std::string& name) override { + return std::make_shared(TracerProvider_->GetTracer(name)); + } + +private: + nostd::shared_ptr TracerProvider_; +}; + +} // namespace -std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { - return std::make_shared(TracerProvider_->GetTracer(name)); +std::shared_ptr CreateOtelTraceProvider( + opentelemetry::nostd::shared_ptr tracerProvider) +{ + return std::make_shared(std::move(tracerProvider)); } } // namespace NYdb::NMetrics diff --git a/src/client/CMakeLists.txt b/src/client/CMakeLists.txt index 65167cddbd0..ce5e4938058 100644 --- a/src/client/CMakeLists.txt +++ b/src/client/CMakeLists.txt @@ -26,5 +26,6 @@ add_subdirectory(scheme) add_subdirectory(ss_tasks) add_subdirectory(table) add_subdirectory(topic) +add_subdirectory(trace) add_subdirectory(types) add_subdirectory(value) diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 63aeac8aea5..c0ef98756fe 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -51,8 +51,8 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t GetMaxMessageSize() const override { return MaxMessageSize; } const TLog& GetLog() const override { return Log; } std::shared_ptr GetExecutor() const override { return Executor; } - std::shared_ptr GetMetricExporter() const override { return MetricExporter; } - std::shared_ptr GetTraceExporter() const override { return TraceExporter; } + std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } + std::shared_ptr GetTraceProvider() const override { return TraceProvider; } std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -82,8 +82,8 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t MaxMessageSize = 0; TLog Log; // Null by default. std::shared_ptr Executor; - std::shared_ptr MetricExporter; - std::shared_ptr TraceExporter; + std::shared_ptr MetricRegistry; + std::shared_ptr TraceProvider; }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -233,24 +233,16 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { return *this; } -TDriverConfig& TDriverConfig::SetMetricExporter(std::shared_ptr exporter) { - Impl_->MetricExporter = std::move(exporter); +TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { + Impl_->MetricRegistry = std::move(registry); return *this; } -TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { - Impl_->TraceExporter = std::move(exporter); +TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr provider) { + Impl_->TraceProvider = std::move(provider); return *this; } -std::shared_ptr TDriverConfig::GetMetricExporter() const { - return Impl_->MetricExporter; -} - -std::shared_ptr TDriverConfig::GetTraceExporter() const { - return Impl_->TraceExporter; -} - //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -302,8 +294,8 @@ TDriverConfig TDriver::GetConfig() const { config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); config.SetMaxMessageSize(Impl_->MaxMessageSize_); config.Impl_->Log = Impl_->Log; - config.SetMetricExporter(Impl_->GetMetricExporter()); - config.SetTraceExporter(Impl_->GetTraceExporter()); + config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); + config.SetTraceProvider(Impl_->GetTraceProvider()); return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index 32645964933..757d5b777b7 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -167,8 +167,8 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout()) #endif - , MetricExporter_(params->GetMetricExporter()) - , TraceExporter_(params->GetTraceExporter()) + , MetricRegistry_(params->GetExternalMetricRegistry()) + , TraceProvider_(params->GetTraceProvider()) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -436,12 +436,12 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { ExtensionApis_.emplace_back(api); } -std::shared_ptr TGRpcConnectionsImpl::GetMetricExporter() const { - return MetricExporter_; +std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { + return MetricRegistry_; } -std::shared_ptr TGRpcConnectionsImpl::GetTraceExporter() const { - return TraceExporter_; +std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { + return TraceProvider_; } void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 14db28028d2..84a25162912 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -586,19 +586,8 @@ class TGRpcConnectionsImpl ::NMonitoring::TMetricRegistry* GetMetricRegistry() override; void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); - std::shared_ptr GetMetricExporter() const; - std::shared_ptr GetTraceExporter() const; - - template - T* GetExtensionApi() { - std::lock_guard lock(ExtensionsLock_); - for (const auto& api : ExtensionApis_) { - if (auto ptr = dynamic_cast(api.get())) { - return ptr; - } - } - return nullptr; - } + std::shared_ptr GetExternalMetricRegistry() const; + std::shared_ptr GetTraceProvider() const; void SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb); const TLog& GetLog() const override; @@ -733,8 +722,8 @@ class TGRpcConnectionsImpl std::vector> Extensions_; std::vector> ExtensionApis_; - std::shared_ptr MetricExporter_; - std::shared_ptr TraceExporter_; + std::shared_ptr MetricRegistry_; + std::shared_ptr TraceProvider_; IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index e6aa5e87dde..1e827d3343b 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -41,8 +41,8 @@ class IConnectionsParams { virtual uint64_t GetMaxOutboundMessageSize() const = 0; virtual uint64_t GetMaxMessageSize() const = 0; virtual std::shared_ptr GetExecutor() const = 0; - virtual std::shared_ptr GetMetricExporter() const = 0; - virtual std::shared_ptr GetTraceExporter() const = 0; + virtual std::shared_ptr GetExternalMetricRegistry() const = 0; + virtual std::shared_ptr GetTraceProvider() const = 0; }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index 03915de6b5e..e681a846b26 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -1,7 +1,7 @@ _ydb_sdk_add_library(client-metrics) target_sources(client-metrics PRIVATE - metrics.cpp + metrics.cpp ) _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/query/CMakeLists.txt b/src/client/query/CMakeLists.txt index f1395ff107b..3cc7401200b 100644 --- a/src/client/query/CMakeLists.txt +++ b/src/client/query/CMakeLists.txt @@ -7,11 +7,12 @@ target_link_libraries(client-ydb_query PUBLIC impl-internal-make_request impl-session impl-internal-retry - client-metrics client-ydb_common_client client-ydb_driver client-ydb_query-impl client-ydb_result + client-metrics + client-trace client-types-operation api-protos api-grpc diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 80a3d6ed00f..0fc3f75ec88 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -15,9 +15,10 @@ #include #include #include +#include #include #include -#include +#include #include @@ -70,9 +71,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); - if (auto traceProvider = Connections_->GetTraceExporter()) { + if (auto traceProvider = Connections_->GetTraceProvider()) { Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } + MetricRegistry_ = Connections_->GetExternalMetricRegistry(); } ~TImpl() { @@ -102,16 +104,21 @@ class TQueryClient::TImpl: public TClientImplCommon, public CollectParamsSize(params ? ¶ms->GetProtoMap() : nullptr); auto span = std::make_shared(Tracer_, "ExecuteQuery", DbDriverState_->DiscoveryEndpoint); + span->SetQueryText(query); + auto metrics = std::make_shared(MetricRegistry_, "ExecuteQuery"); return TExecQueryImpl::ExecuteQuery( Connections_, DbDriverState_, query, txControl, params, settings, session) - .Apply([span](TAsyncExecuteQueryResult future) { + .Apply([span, metrics](TAsyncExecuteQueryResult future) { try { auto result = future.GetValue(); + span->SetPeerEndpoint(result.GetEndpoint()); span->End(result.GetStatus()); + metrics->End(result.GetStatus()); return result; } catch (...) { span->End(EStatus::CLIENT_INTERNAL_ERROR); + metrics->End(EStatus::CLIENT_INTERNAL_ERROR); throw; } }); @@ -182,10 +189,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto promise = NThreading::NewPromise(); auto span = std::make_shared(Tracer_, "Rollback", DbDriverState_->DiscoveryEndpoint); + auto metrics = std::make_shared(MetricRegistry_, "Rollback"); - auto responseCb = [promise, session, span] + auto responseCb = [promise, session, span, metrics] (Ydb::Query::RollbackTransactionResponse* response, TPlainStatus status) mutable { try { + span->SetPeerEndpoint(status.Endpoint); if (response) { NYdb::NIssue::TIssues opIssues; NYdb::NIssue::IssuesFromMessage(response->issues(), opIssues); @@ -193,14 +202,17 @@ class TQueryClient::TImpl: public TClientImplCommon, public status.Endpoint, std::move(status.Metadata)}); span->End(rollbackTxStatus.GetStatus()); + metrics->End(rollbackTxStatus.GetStatus()); promise.SetValue(std::move(rollbackTxStatus)); } else { span->End(status.Status); + metrics->End(status.Status); promise.SetValue(TStatus(std::move(status))); } } catch (...) { span->End(EStatus::CLIENT_INTERNAL_ERROR); + metrics->End(EStatus::CLIENT_INTERNAL_ERROR); promise.SetException(std::current_exception()); } }; @@ -229,10 +241,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto promise = NThreading::NewPromise(); auto span = std::make_shared(Tracer_, "Commit", DbDriverState_->DiscoveryEndpoint); + auto metrics = std::make_shared(MetricRegistry_, "Commit"); - auto responseCb = [promise, session, span] + auto responseCb = [promise, session, span, metrics] (Ydb::Query::CommitTransactionResponse* response, TPlainStatus status) mutable { try { + span->SetPeerEndpoint(status.Endpoint); if (response) { NYdb::NIssue::TIssues opIssues; NYdb::NIssue::IssuesFromMessage(response->issues(), opIssues); @@ -240,15 +254,18 @@ class TQueryClient::TImpl: public TClientImplCommon, public status.Endpoint, std::move(status.Metadata)}); span->End(commitTxStatus.GetStatus()); + metrics->End(commitTxStatus.GetStatus()); TCommitTransactionResult commitTxResult(std::move(commitTxStatus)); promise.SetValue(std::move(commitTxResult)); } else { span->End(status.Status); + metrics->End(status.Status); promise.SetValue(TCommitTransactionResult(TStatus(std::move(status)))); } } catch (...) { span->End(EStatus::CLIENT_INTERNAL_ERROR); + metrics->End(EStatus::CLIENT_INTERNAL_ERROR); promise.SetException(std::current_exception()); } }; @@ -456,11 +473,13 @@ class TQueryClient::TImpl: public TClientImplCommon, public TAsyncCreateSessionResult GetSession(const TCreateSessionSettings& settings) { class TQueryClientGetSessionCtx : public NSessionPool::IGetSessionCtx { public: - TQueryClientGetSessionCtx(std::shared_ptr client, const TCreateSessionSettings& settings, std::shared_ptr span) + TQueryClientGetSessionCtx(std::shared_ptr client, const TCreateSessionSettings& settings, + std::shared_ptr span, std::shared_ptr metrics) : Promise(NThreading::NewPromise()) , Client(client) , RpcSettings(TRpcRequestSettings::Make(settings)) , Span(span) + , Metrics(metrics) {} TAsyncCreateSessionResult GetFuture() { @@ -472,6 +491,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(status.GetStatus()); } + if (Metrics) { + Metrics->End(status.GetStatus()); + } ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -487,17 +509,24 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(EStatus::SUCCESS); } + if (Metrics) { + Metrics->End(EStatus::SUCCESS); + } ScheduleReply(std::move(val)); } void ReplyNewSession() override { Client->CreateAttachedSession(RpcSettings).Subscribe( - [promise{std::move(Promise)}, span = Span](TAsyncCreateSessionResult future) mutable + [promise{std::move(Promise)}, span = Span, metrics = Metrics](TAsyncCreateSessionResult future) mutable { auto val = future.ExtractValue(); if (span) { + span->SetPeerEndpoint(val.GetEndpoint()); span->End(val.GetStatus()); } + if (metrics) { + metrics->End(val.GetStatus()); + } promise.SetValue(std::move(val)); }); } @@ -524,10 +553,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public std::shared_ptr Client; const TRpcRequestSettings RpcSettings; std::shared_ptr Span; + std::shared_ptr Metrics; }; auto span = std::make_shared(Tracer_, "CreateSession", DbDriverState_->DiscoveryEndpoint); - auto ctx = std::make_unique(shared_from_this(), settings, span); + auto metrics = std::make_shared(MetricRegistry_, "CreateSession"); + auto ctx = std::make_unique(shared_from_this(), settings, span, metrics); auto future = ctx->GetFuture(); SessionPool_.GetSession(std::move(ctx)); @@ -597,6 +628,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public private: std::shared_ptr Tracer_; + std::shared_ptr MetricRegistry_; NSdkStats::TStatCollector::TClientRetryOperationStatCollector RetryOperationStatCollector_; NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> QuerySizeHistogram_; NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> ParamsSizeHistogram_; diff --git a/src/client/query/impl/CMakeLists.txt b/src/client/query/impl/CMakeLists.txt index 70f93b6d68d..c6c290795fe 100644 --- a/src/client/query/impl/CMakeLists.txt +++ b/src/client/query/impl/CMakeLists.txt @@ -12,6 +12,7 @@ target_link_libraries(client-ydb_query-impl PUBLIC target_sources(client-ydb_query-impl PRIVATE exec_query.cpp client_session.cpp + query_metrics.cpp query_spans.cpp ) diff --git a/src/client/query/impl/query_metrics.cpp b/src/client/query/impl/query_metrics.cpp new file mode 100644 index 00000000000..f314f3d8b96 --- /dev/null +++ b/src/client/query/impl/query_metrics.cpp @@ -0,0 +1,71 @@ +#include "query_metrics.h" + +#include + +namespace NYdb::inline V3::NQuery { + +namespace { + +void SafeLogMetricsError(const char* message) noexcept { + try { + try { + std::cerr << "TQueryMetrics: " << message << ": " << CurrentExceptionMessage() << std::endl; + return; + } catch (...) { + } + std::cerr << "TQueryMetrics: " << message << ": (unknown)" << std::endl; + } catch (...) { + } +} + +} // namespace + +static const std::vector LatencyBuckets = { + 1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000 +}; + +TQueryMetrics::TQueryMetrics(std::shared_ptr registry, const std::string& operationName) { + if (!registry) { + return; + } + + try { + NMetrics::TLabels labels = {{"operation", operationName}}; + RequestCounter_ = registry->Counter("ydb.query.requests", labels); + ErrorCounter_ = registry->Counter("ydb.query.errors", labels); + LatencyHistogram_ = registry->Histogram("ydb.query.latency_ms", LatencyBuckets, labels); + + RequestCounter_->Inc(); + StartTime_ = TInstant::Now(); + } catch (...) { + SafeLogMetricsError("failed to initialize metrics"); + RequestCounter_.reset(); + ErrorCounter_.reset(); + LatencyHistogram_.reset(); + } +} + +TQueryMetrics::~TQueryMetrics() noexcept { + End(EStatus::CLIENT_INTERNAL_ERROR); +} + +void TQueryMetrics::End(EStatus status) noexcept { + if (Ended_) { + return; + } + Ended_ = true; + + try { + if (LatencyHistogram_) { + auto durationMs = (TInstant::Now() - StartTime_).MilliSeconds(); + LatencyHistogram_->Record(static_cast(durationMs)); + } + if (status != EStatus::SUCCESS && ErrorCounter_) { + ErrorCounter_->Inc(); + } + } catch (...) { + SafeLogMetricsError("failed to record metrics"); + } +} + +} // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_metrics.h b/src/client/query/impl/query_metrics.h new file mode 100644 index 00000000000..807472e2434 --- /dev/null +++ b/src/client/query/impl/query_metrics.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +#include + +#include +#include + +namespace NYdb::inline V3::NQuery { + +class TQueryMetrics { +public: + TQueryMetrics(std::shared_ptr registry, const std::string& operationName); + ~TQueryMetrics() noexcept; + + void End(EStatus status) noexcept; + +private: + std::shared_ptr RequestCounter_; + std::shared_ptr ErrorCounter_; + std::shared_ptr LatencyHistogram_; + TInstant StartTime_; + bool Ended_ = false; +}; + +} // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index fc1329aacd0..4bbd4d2250b 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -6,29 +6,49 @@ namespace NYdb::inline V3::NQuery { namespace { +constexpr int DefaultGrpcPort = 2135; + void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { - auto pos = endpoint.find(':'); + port = DefaultGrpcPort; + + if (endpoint.empty()) { + host = endpoint; + return; + } + + // IPv6 bracket notation: [addr]:port + if (endpoint.front() == '[') { + auto bracketEnd = endpoint.find(']'); + if (bracketEnd != std::string::npos) { + host = endpoint.substr(1, bracketEnd - 1); + if (bracketEnd + 2 < endpoint.size() && endpoint[bracketEnd + 1] == ':') { + try { + port = std::stoi(endpoint.substr(bracketEnd + 2)); + } catch (...) {} + } + return; + } + } + + auto pos = endpoint.rfind(':'); if (pos != std::string::npos) { host = endpoint.substr(0, pos); try { port = std::stoi(endpoint.substr(pos + 1)); - } catch (...) { - port = 2135; - } + } catch (...) {} } else { host = endpoint; - port = 2135; } } void SafeLogSpanError(const char* message) noexcept { try { try { - Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + std::cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << std::endl; return; } catch (...) { } - Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + std::cerr << "TQuerySpan: " << message << ": (unknown)" << std::endl; } catch (...) { } } @@ -68,6 +88,43 @@ TQuerySpan::~TQuerySpan() noexcept { } } +void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { + if (!Span_ || endpoint.empty()) { + return; + } + try { + std::string host; + int port; + ParseEndpoint(endpoint, host, port); + Span_->SetAttribute("network.peer.address", host); + Span_->SetAttribute("network.peer.port", static_cast(port)); + } catch (...) { + SafeLogSpanError("failed to set peer endpoint"); + } +} + +void TQuerySpan::SetQueryText(const std::string& query) noexcept { + if (!Span_ || query.empty()) { + return; + } + try { + Span_->SetAttribute("db.query.text", query); + } catch (...) { + SafeLogSpanError("failed to set query text"); + } +} + +void TQuerySpan::AddEvent(const std::string& name, const std::map& attributes) noexcept { + if (!Span_) { + return; + } + try { + Span_->AddEvent(name, attributes); + } catch (...) { + SafeLogSpanError("failed to add event"); + } +} + void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index ca0b6853954..75fd0fa830e 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -1,9 +1,10 @@ #pragma once -#include +#include #include #include +#include #include #include @@ -14,6 +15,10 @@ class TQuerySpan { TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); ~TQuerySpan() noexcept; + void SetPeerEndpoint(const std::string& endpoint) noexcept; + void SetQueryText(const std::string& query) noexcept; + void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; + void End(EStatus status) noexcept; private: diff --git a/src/client/trace/CMakeLists.txt b/src/client/trace/CMakeLists.txt new file mode 100644 index 00000000000..86a8f8d4208 --- /dev/null +++ b/src/client/trace/CMakeLists.txt @@ -0,0 +1,7 @@ +_ydb_sdk_add_library(client-trace) + +target_sources(client-trace PRIVATE + trace.cpp +) + +_ydb_sdk_make_client_component(Trace client-trace) diff --git a/src/client/trace/trace.cpp b/src/client/trace/trace.cpp new file mode 100644 index 00000000000..6bf5bc664f0 --- /dev/null +++ b/src/client/trace/trace.cpp @@ -0,0 +1 @@ +#include diff --git a/tests/common/fake_metric_registry.h b/tests/common/fake_metric_registry.h new file mode 100644 index 00000000000..60ff1414633 --- /dev/null +++ b/tests/common/fake_metric_registry.h @@ -0,0 +1,122 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace NYdb::NTests { + +class TFakeCounter : public NMetrics::ICounter { +public: + void Inc() override { + Count_.fetch_add(1, std::memory_order_relaxed); + } + + int64_t Get() const { + return Count_.load(std::memory_order_relaxed); + } + +private: + std::atomic Count_{0}; +}; + +class TFakeHistogram : public NMetrics::IHistogram { +public: + void Record(double value) override { + std::lock_guard lock(Mutex_); + Values_.push_back(value); + } + + std::vector GetValues() const { + std::lock_guard lock(Mutex_); + return Values_; + } + + size_t Count() const { + std::lock_guard lock(Mutex_); + return Values_.size(); + } + +private: + mutable std::mutex Mutex_; + std::vector Values_; +}; + +class TFakeGauge : public NMetrics::IGauge { +public: + void Add(double delta) override { Value_ += delta; } + void Set(double value) override { Value_ = value; } + double Get() const { return Value_; } + +private: + double Value_ = 0.0; +}; + +struct TMetricKey { + std::string Name; + NMetrics::TLabels Labels; + + bool operator==(const TMetricKey& other) const = default; + bool operator<(const TMetricKey& other) const { + if (Name != other.Name) return Name < other.Name; + return Labels < other.Labels; + } +}; + +class TFakeMetricRegistry : public NMetrics::IMetricRegistry { +public: + std::shared_ptr Counter(const std::string& name, const NMetrics::TLabels& labels) override { + std::lock_guard lock(Mutex_); + auto key = TMetricKey{name, labels}; + auto it = Counters_.find(key); + if (it != Counters_.end()) { + return it->second; + } + auto counter = std::make_shared(); + Counters_[key] = counter; + return counter; + } + + std::shared_ptr Gauge(const std::string& name, const NMetrics::TLabels& labels) override { + std::lock_guard lock(Mutex_); + auto key = TMetricKey{name, labels}; + auto gauge = std::make_shared(); + Gauges_[key] = gauge; + return gauge; + } + + std::shared_ptr Histogram(const std::string& name, const std::vector& /*buckets*/, const NMetrics::TLabels& labels) override { + std::lock_guard lock(Mutex_); + auto key = TMetricKey{name, labels}; + auto it = Histograms_.find(key); + if (it != Histograms_.end()) { + return it->second; + } + auto histogram = std::make_shared(); + Histograms_[key] = histogram; + return histogram; + } + + std::shared_ptr GetCounter(const std::string& name, const NMetrics::TLabels& labels = {}) const { + std::lock_guard lock(Mutex_); + auto it = Counters_.find(TMetricKey{name, labels}); + return it != Counters_.end() ? it->second : nullptr; + } + + std::shared_ptr GetHistogram(const std::string& name, const NMetrics::TLabels& labels = {}) const { + std::lock_guard lock(Mutex_); + auto it = Histograms_.find(TMetricKey{name, labels}); + return it != Histograms_.end() ? it->second : nullptr; + } + +private: + mutable std::mutex Mutex_; + std::map> Counters_; + std::map> Gauges_; + std::map> Histograms_; +}; + +} // namespace NYdb::NTests diff --git a/tests/integration/CMakeLists.txt b/tests/integration/CMakeLists.txt index d5a1d709245..8aa28839a63 100644 --- a/tests/integration/CMakeLists.txt +++ b/tests/integration/CMakeLists.txt @@ -1,6 +1,7 @@ add_subdirectory(auth) add_subdirectory(basic_example) add_subdirectory(bulk_upsert) +add_subdirectory(metrics) add_subdirectory(server_restart) add_subdirectory(sessions) add_subdirectory(sessions_pool) diff --git a/tests/integration/metrics/CMakeLists.txt b/tests/integration/metrics/CMakeLists.txt new file mode 100644 index 00000000000..6c9bb8b3abd --- /dev/null +++ b/tests/integration/metrics/CMakeLists.txt @@ -0,0 +1,12 @@ +add_ydb_test(NAME metrics_it GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + main.cpp + LINK_LIBRARIES + yutil + YDB-CPP-SDK::Query + client-metrics + LABELS + integration +) diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp new file mode 100644 index 00000000000..cab8f71cdc4 --- /dev/null +++ b/tests/integration/metrics/main.cpp @@ -0,0 +1,273 @@ +#include +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NQuery; +using namespace NYdb::NTests; + +namespace { + +struct TRunArgs { + TDriver Driver; + std::shared_ptr Registry; +}; + +TRunArgs MakeRunArgs() { + std::string endpoint = std::getenv("YDB_ENDPOINT"); + std::string database = std::getenv("YDB_DATABASE"); + + auto registry = std::make_shared(); + + auto driverConfig = TDriverConfig() + .SetEndpoint(endpoint) + .SetDatabase(database) + .SetAuthToken(std::getenv("YDB_TOKEN") ? std::getenv("YDB_TOKEN") : "") + .SetMetricRegistry(registry); + + TDriver driver(driverConfig); + return {driver, registry}; +} + +std::shared_ptr GetCounter( + const std::shared_ptr& registry, + const std::string& name, + const std::string& operation) +{ + return registry->GetCounter(name, {{"operation", operation}}); +} + +std::shared_ptr GetHistogram( + const std::shared_ptr& registry, + const std::string& name, + const std::string& operation) +{ + return registry->GetHistogram(name, {{"operation", operation}}); +} + +} // namespace + +TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto session = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); + + auto result = session.GetSession().ExecuteQuery( + "SELECT 1;", + TTxControl::BeginTx().CommitTx() + ).ExtractValueSync(); + ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); + + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; + EXPECT_GE(requests->Get(), 1); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr) << "ExecuteQuery latency histogram not created"; + EXPECT_GE(latency->Count(), 1u); + for (double v : latency->GetValues()) { + EXPECT_GE(v, 0.0); + } + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto session = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); + + auto result = session.GetSession().ExecuteQuery( + "INVALID SQL QUERY !!!", + TTxControl::BeginTx().CommitTx() + ).ExtractValueSync(); + EXPECT_NE(result.GetStatus(), EStatus::SUCCESS); + + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_GE(requests->Get(), 1); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_GE(errors->Get(), 1); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + EXPECT_GE(latency->Count(), 1u); + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto session = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); + + auto requests = GetCounter(registry, "ydb.query.requests", "CreateSession"); + ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; + EXPECT_GE(requests->Get(), 1); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "CreateSession"); + ASSERT_NE(latency, nullptr) << "CreateSession latency histogram not created"; + EXPECT_GE(latency->Count(), 1u); + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto sessionResult = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(sessionResult.IsSuccess()) << sessionResult.GetIssues().ToString(); + auto session = sessionResult.GetSession(); + + auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).ExtractValueSync(); + ASSERT_TRUE(beginResult.IsSuccess()) << beginResult.GetIssues().ToString(); + auto tx = beginResult.GetTransaction(); + + auto execResult = tx.GetSession().ExecuteQuery( + "SELECT 1;", + TTxControl::Tx(tx) + ).ExtractValueSync(); + ASSERT_EQ(execResult.GetStatus(), EStatus::SUCCESS) << execResult.GetIssues().ToString(); + + if (execResult.GetTransaction()) { + auto commitResult = execResult.GetTransaction()->Commit().ExtractValueSync(); + ASSERT_TRUE(commitResult.IsSuccess()) << commitResult.GetIssues().ToString(); + + auto commitRequests = GetCounter(registry, "ydb.query.requests", "Commit"); + ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; + EXPECT_GE(commitRequests->Get(), 1); + + auto commitLatency = GetHistogram(registry, "ydb.query.latency_ms", "Commit"); + ASSERT_NE(commitLatency, nullptr); + EXPECT_GE(commitLatency->Count(), 1u); + } + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto sessionResult = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(sessionResult.IsSuccess()) << sessionResult.GetIssues().ToString(); + auto session = sessionResult.GetSession(); + + auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).ExtractValueSync(); + ASSERT_TRUE(beginResult.IsSuccess()) << beginResult.GetIssues().ToString(); + auto tx = beginResult.GetTransaction(); + + auto rollbackResult = tx.Rollback().ExtractValueSync(); + ASSERT_TRUE(rollbackResult.IsSuccess()) << rollbackResult.GetIssues().ToString(); + + auto rollbackRequests = GetCounter(registry, "ydb.query.requests", "Rollback"); + ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; + EXPECT_GE(rollbackRequests->Get(), 1); + + auto rollbackErrors = GetCounter(registry, "ydb.query.errors", "Rollback"); + ASSERT_NE(rollbackErrors, nullptr); + EXPECT_EQ(rollbackErrors->Get(), 0); + + auto rollbackLatency = GetHistogram(registry, "ydb.query.latency_ms", "Rollback"); + ASSERT_NE(rollbackLatency, nullptr); + EXPECT_GE(rollbackLatency->Count(), 1u); + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto sessionResult = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(sessionResult.IsSuccess()) << sessionResult.GetIssues().ToString(); + auto session = sessionResult.GetSession(); + + const int numQueries = 5; + for (int i = 0; i < numQueries; ++i) { + auto result = session.ExecuteQuery( + "SELECT 1;", + TTxControl::BeginTx().CommitTx() + ).ExtractValueSync(); + ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); + } + + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), numQueries); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + EXPECT_EQ(latency->Count(), static_cast(numQueries)); + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, NoRegistryDoesNotBreakOperations) { + std::string endpoint = std::getenv("YDB_ENDPOINT"); + std::string database = std::getenv("YDB_DATABASE"); + + auto driverConfig = TDriverConfig() + .SetEndpoint(endpoint) + .SetDatabase(database) + .SetAuthToken(std::getenv("YDB_TOKEN") ? std::getenv("YDB_TOKEN") : ""); + + TDriver driver(driverConfig); + TQueryClient client(driver); + + auto session = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); + + auto result = session.GetSession().ExecuteQuery( + "SELECT 1;", + TTxControl::BeginTx().CommitTx() + ).ExtractValueSync(); + EXPECT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); + + driver.Stop(true); +} + +TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { + auto [driver, registry] = MakeRunArgs(); + TQueryClient client(driver); + + auto sessionResult = client.GetSession().ExtractValueSync(); + ASSERT_TRUE(sessionResult.IsSuccess()) << sessionResult.GetIssues().ToString(); + auto session = sessionResult.GetSession(); + + auto result = session.ExecuteQuery( + "SELECT 1;", + TTxControl::BeginTx().CommitTx() + ).ExtractValueSync(); + ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + ASSERT_GE(latency->Count(), 1u); + + for (double v : latency->GetValues()) { + EXPECT_GE(v, 0.0) << "Latency must be non-negative"; + EXPECT_LT(v, 30000.0) << "Latency > 30s is unrealistic for SELECT 1"; + } + + driver.Stop(true); +} diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index 03b0a17c386..2ad5a38e01f 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -100,3 +100,16 @@ add_ydb_test(NAME client-ydb_value_ut GTEST LABELS unit ) + +add_ydb_test(NAME client-ydb_query_metrics_ut GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + query/query_metrics_ut.cpp + LINK_LIBRARIES + yutil + client-ydb_query-impl + client-metrics + LABELS + unit +) diff --git a/tests/unit/client/query/query_metrics_ut.cpp b/tests/unit/client/query/query_metrics_ut.cpp new file mode 100644 index 00000000000..20c681b7eca --- /dev/null +++ b/tests/unit/client/query/query_metrics_ut.cpp @@ -0,0 +1,190 @@ +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NQuery; +using namespace NYdb::NMetrics; +using namespace NYdb::NTests; + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +class QueryMetricsTest : public ::testing::Test { +protected: + void SetUp() override { + Registry = std::make_shared(); + } + + std::shared_ptr RequestCounter(const std::string& op) { + return Registry->GetCounter("ydb.query.requests", {{"operation", op}}); + } + + std::shared_ptr ErrorCounter(const std::string& op) { + return Registry->GetCounter("ydb.query.errors", {{"operation", op}}); + } + + std::shared_ptr LatencyHistogram(const std::string& op) { + return Registry->GetHistogram("ydb.query.latency_ms", {{"operation", op}}); + } + + std::shared_ptr Registry; +}; + +TEST_F(QueryMetricsTest, RequestCounterIncrementedOnConstruction) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + + auto counter = RequestCounter("ExecuteQuery"); + ASSERT_NE(counter, nullptr); + EXPECT_EQ(counter->Get(), 1); +} + +TEST_F(QueryMetricsTest, SuccessDoesNotIncrementErrorCounter) { + { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + } + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); +} + +TEST_F(QueryMetricsTest, FailureIncrementsErrorCounter) { + { + TQueryMetrics metrics(Registry, "Commit"); + metrics.End(EStatus::UNAVAILABLE); + } + + auto errors = ErrorCounter("Commit"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); +} + +TEST_F(QueryMetricsTest, LatencyRecordedOnEnd) { + { + TQueryMetrics metrics(Registry, "Rollback"); + metrics.End(EStatus::SUCCESS); + } + + auto hist = LatencyHistogram("Rollback"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); + EXPECT_GE(hist->GetValues()[0], 0.0); +} + +TEST_F(QueryMetricsTest, DoubleEndIsIdempotent) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + metrics.End(EStatus::INTERNAL_ERROR); + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto hist = LatencyHistogram("ExecuteQuery"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(QueryMetricsTest, DestructorCallsEndWithClientInternalError) { + { + TQueryMetrics metrics(Registry, "CreateSession"); + } + + auto requests = RequestCounter("CreateSession"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), 1); + + auto errors = ErrorCounter("CreateSession"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); + + auto hist = LatencyHistogram("CreateSession"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(QueryMetricsTest, NullRegistryDoesNotCrash) { + EXPECT_NO_THROW({ + TQueryMetrics metrics(nullptr, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + }); +} + +TEST_F(QueryMetricsTest, CorrectMetricNamesAndLabels) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + + EXPECT_NE(Registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(Registry->GetCounter("ydb.query.errors", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(Registry->GetHistogram("ydb.query.latency_ms", {{"operation", "ExecuteQuery"}}), nullptr); + + EXPECT_EQ(Registry->GetCounter("ydb.query.requests", {{"operation", "Commit"}}), nullptr); +} + +TEST_F(QueryMetricsTest, DifferentOperationsHaveSeparateMetrics) { + { + TQueryMetrics m1(Registry, "ExecuteQuery"); + m1.End(EStatus::SUCCESS); + } + { + TQueryMetrics m2(Registry, "Commit"); + m2.End(EStatus::OVERLOADED); + } + + auto execRequests = RequestCounter("ExecuteQuery"); + auto commitRequests = RequestCounter("Commit"); + ASSERT_NE(execRequests, nullptr); + ASSERT_NE(commitRequests, nullptr); + EXPECT_EQ(execRequests->Get(), 1); + EXPECT_EQ(commitRequests->Get(), 1); + + auto execErrors = ErrorCounter("ExecuteQuery"); + auto commitErrors = ErrorCounter("Commit"); + EXPECT_EQ(execErrors->Get(), 0); + EXPECT_EQ(commitErrors->Get(), 1); +} + +TEST_F(QueryMetricsTest, MultipleRequestsAccumulate) { + for (int i = 0; i < 5; ++i) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(i % 2 == 0 ? EStatus::SUCCESS : EStatus::TIMEOUT); + } + + auto requests = RequestCounter("ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), 5); + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 2); + + auto hist = LatencyHistogram("ExecuteQuery"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 5u); +} + +TEST_F(QueryMetricsTest, AllErrorStatusesIncrementErrorCounter) { + std::vector errorStatuses = { + EStatus::BAD_REQUEST, + EStatus::UNAUTHORIZED, + EStatus::INTERNAL_ERROR, + EStatus::UNAVAILABLE, + EStatus::OVERLOADED, + EStatus::TIMEOUT, + EStatus::NOT_FOUND, + EStatus::CLIENT_INTERNAL_ERROR, + }; + + for (auto status : errorStatuses) { + TQueryMetrics metrics(Registry, "Rollback"); + metrics.End(status); + } + + auto errors = ErrorCounter("Rollback"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), static_cast(errorStatuses.size())); +} From aa80faab259c06902eca1787038f047dc5889eb7 Mon Sep 17 00:00:00 2001 From: maladetska Date: Fri, 27 Mar 2026 22:20:48 +0300 Subject: [PATCH 07/93] add table metrics --- src/client/impl/CMakeLists.txt | 1 + src/client/impl/observability/CMakeLists.txt | 12 + .../impl/observability/client_metrics.cpp | 116 +++++++++ .../impl/observability/client_metrics.h | 28 +++ src/client/query/impl/CMakeLists.txt | 5 +- src/client/query/impl/query_metrics.cpp | 71 ------ src/client/query/impl/query_metrics.h | 24 +- src/client/table/impl/CMakeLists.txt | 2 + src/client/table/impl/table_metrics.h | 14 ++ tests/integration/metrics/main.cpp | 60 ++--- tests/unit/client/CMakeLists.txt | 6 +- .../observability/client_metrics_ut.cpp | 229 ++++++++++++++++++ tests/unit/client/query/query_metrics_ut.cpp | 190 --------------- 13 files changed, 446 insertions(+), 312 deletions(-) create mode 100644 src/client/impl/observability/CMakeLists.txt create mode 100644 src/client/impl/observability/client_metrics.cpp create mode 100644 src/client/impl/observability/client_metrics.h delete mode 100644 src/client/query/impl/query_metrics.cpp create mode 100644 src/client/table/impl/table_metrics.h create mode 100644 tests/unit/client/observability/client_metrics_ut.cpp delete mode 100644 tests/unit/client/query/query_metrics_ut.cpp diff --git a/src/client/impl/CMakeLists.txt b/src/client/impl/CMakeLists.txt index 9e04f134b37..8dfc3fa865b 100644 --- a/src/client/impl/CMakeLists.txt +++ b/src/client/impl/CMakeLists.txt @@ -1,5 +1,6 @@ add_subdirectory(endpoints) add_subdirectory(executor) add_subdirectory(internal) +add_subdirectory(observability) add_subdirectory(session) add_subdirectory(stats) diff --git a/src/client/impl/observability/CMakeLists.txt b/src/client/impl/observability/CMakeLists.txt new file mode 100644 index 00000000000..961d2821559 --- /dev/null +++ b/src/client/impl/observability/CMakeLists.txt @@ -0,0 +1,12 @@ +_ydb_sdk_add_library(impl-observability) + +target_link_libraries(impl-observability PUBLIC + yutil + client-metrics +) + +target_sources(impl-observability PRIVATE + client_metrics.cpp +) + +_ydb_sdk_install_targets(TARGETS impl-observability) diff --git a/src/client/impl/observability/client_metrics.cpp b/src/client/impl/observability/client_metrics.cpp new file mode 100644 index 00000000000..efa9b739517 --- /dev/null +++ b/src/client/impl/observability/client_metrics.cpp @@ -0,0 +1,116 @@ +#include "client_metrics.h" + +#include + +namespace NYdb::inline V3::NObservability { + +namespace { + +void SafeLogMetricsError(const char* /*message*/) noexcept { + try { + try { + std::rethrow_exception(std::current_exception()); + } catch (const std::exception&) { + return; + } catch (...) { + } + } catch (...) { + } +} + +std::string StatusToString(EStatus status) { + switch (status) { + case EStatus::SUCCESS: return "SUCCESS"; + case EStatus::BAD_REQUEST: return "BAD_REQUEST"; + case EStatus::UNAUTHORIZED: return "UNAUTHORIZED"; + case EStatus::INTERNAL_ERROR: return "INTERNAL_ERROR"; + case EStatus::ABORTED: return "ABORTED"; + case EStatus::UNAVAILABLE: return "UNAVAILABLE"; + case EStatus::OVERLOADED: return "OVERLOADED"; + case EStatus::SCHEME_ERROR: return "SCHEME_ERROR"; + case EStatus::GENERIC_ERROR: return "GENERIC_ERROR"; + case EStatus::TIMEOUT: return "TIMEOUT"; + case EStatus::BAD_SESSION: return "BAD_SESSION"; + case EStatus::PRECONDITION_FAILED: return "PRECONDITION_FAILED"; + case EStatus::ALREADY_EXISTS: return "ALREADY_EXISTS"; + case EStatus::NOT_FOUND: return "NOT_FOUND"; + case EStatus::SESSION_EXPIRED: return "SESSION_EXPIRED"; + case EStatus::CANCELLED: return "CANCELLED"; + case EStatus::UNDETERMINED: return "UNDETERMINED"; + case EStatus::UNSUPPORTED: return "UNSUPPORTED"; + case EStatus::SESSION_BUSY: return "SESSION_BUSY"; + case EStatus::EXTERNAL_ERROR: return "EXTERNAL_ERROR"; + case EStatus::TRANSPORT_UNAVAILABLE: return "TRANSPORT_UNAVAILABLE"; + case EStatus::CLIENT_RESOURCE_EXHAUSTED:return "CLIENT_RESOURCE_EXHAUSTED"; + case EStatus::CLIENT_DEADLINE_EXCEEDED: return "CLIENT_DEADLINE_EXCEEDED"; + case EStatus::CLIENT_INTERNAL_ERROR: return "CLIENT_INTERNAL_ERROR"; + case EStatus::CLIENT_CANCELLED: return "CLIENT_CANCELLED"; + case EStatus::CLIENT_UNAUTHENTICATED: return "CLIENT_UNAUTHENTICATED"; + case EStatus::CLIENT_CALL_UNIMPLEMENTED:return "CLIENT_CALL_UNIMPLEMENTED"; + case EStatus::CLIENT_OUT_OF_RANGE: return "CLIENT_OUT_OF_RANGE"; + case EStatus::CLIENT_DISCOVERY_FAILED: return "CLIENT_DISCOVERY_FAILED"; + case EStatus::CLIENT_LIMITS_REACHED: return "CLIENT_LIMITS_REACHED"; + default: return "STATUS_UNDEFINED"; + } +} + +} // namespace + +static const std::vector DurationBucketsSec = { + 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10 +}; + +TClientMetrics::TClientMetrics(std::shared_ptr registry, + const std::string& prefix, const std::string& operationName) +{ + if (!registry) { + return; + } + + try { + NMetrics::TLabels labels = {{"operation", operationName}}; + RequestCounter_ = registry->Counter(prefix + ".requests", labels); + ErrorCounter_ = registry->Counter(prefix + ".errors", labels); + + NMetrics::TLabels durationLabels = { + {"db.system.name", "ydb"}, + {"db.operation.name", operationName}, + }; + DurationHistogram_ = registry->Histogram("db.client.operation.duration", DurationBucketsSec, durationLabels); + + RequestCounter_->Inc(); + StartTime_ = std::chrono::steady_clock::now(); + } catch (...) { + SafeLogMetricsError("failed to initialize metrics"); + RequestCounter_.reset(); + ErrorCounter_.reset(); + DurationHistogram_.reset(); + } +} + +TClientMetrics::~TClientMetrics() noexcept { + End(EStatus::CLIENT_INTERNAL_ERROR); +} + +void TClientMetrics::End(EStatus status) noexcept { + if (Ended_) { + return; + } + Ended_ = true; + + try { + if (DurationHistogram_) { + auto elapsed = std::chrono::steady_clock::now() - StartTime_; + double durationSec = std::chrono::duration(elapsed).count(); + DurationHistogram_->Record(durationSec); + } + + if (status != EStatus::SUCCESS && ErrorCounter_) { + ErrorCounter_->Inc(); + } + } catch (...) { + SafeLogMetricsError("failed to record metrics"); + } +} + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/client_metrics.h b/src/client/impl/observability/client_metrics.h new file mode 100644 index 00000000000..bce81a958f2 --- /dev/null +++ b/src/client/impl/observability/client_metrics.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace NYdb::inline V3::NObservability { + +class TClientMetrics { +public: + TClientMetrics(std::shared_ptr registry, + const std::string& prefix, const std::string& operationName); + ~TClientMetrics() noexcept; + + void End(EStatus status) noexcept; + +private: + std::shared_ptr RequestCounter_; + std::shared_ptr ErrorCounter_; + std::shared_ptr DurationHistogram_; + std::chrono::steady_clock::time_point StartTime_; + bool Ended_ = false; +}; + +} // namespace NYdb::NObservability diff --git a/src/client/query/impl/CMakeLists.txt b/src/client/query/impl/CMakeLists.txt index c6c290795fe..fdc46b50f95 100644 --- a/src/client/query/impl/CMakeLists.txt +++ b/src/client/query/impl/CMakeLists.txt @@ -9,10 +9,13 @@ target_link_libraries(client-ydb_query-impl PUBLIC client-ydb_result ) +target_link_libraries(client-ydb_query-impl PUBLIC + impl-observability +) + target_sources(client-ydb_query-impl PRIVATE exec_query.cpp client_session.cpp - query_metrics.cpp query_spans.cpp ) diff --git a/src/client/query/impl/query_metrics.cpp b/src/client/query/impl/query_metrics.cpp deleted file mode 100644 index f314f3d8b96..00000000000 --- a/src/client/query/impl/query_metrics.cpp +++ /dev/null @@ -1,71 +0,0 @@ -#include "query_metrics.h" - -#include - -namespace NYdb::inline V3::NQuery { - -namespace { - -void SafeLogMetricsError(const char* message) noexcept { - try { - try { - std::cerr << "TQueryMetrics: " << message << ": " << CurrentExceptionMessage() << std::endl; - return; - } catch (...) { - } - std::cerr << "TQueryMetrics: " << message << ": (unknown)" << std::endl; - } catch (...) { - } -} - -} // namespace - -static const std::vector LatencyBuckets = { - 1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000 -}; - -TQueryMetrics::TQueryMetrics(std::shared_ptr registry, const std::string& operationName) { - if (!registry) { - return; - } - - try { - NMetrics::TLabels labels = {{"operation", operationName}}; - RequestCounter_ = registry->Counter("ydb.query.requests", labels); - ErrorCounter_ = registry->Counter("ydb.query.errors", labels); - LatencyHistogram_ = registry->Histogram("ydb.query.latency_ms", LatencyBuckets, labels); - - RequestCounter_->Inc(); - StartTime_ = TInstant::Now(); - } catch (...) { - SafeLogMetricsError("failed to initialize metrics"); - RequestCounter_.reset(); - ErrorCounter_.reset(); - LatencyHistogram_.reset(); - } -} - -TQueryMetrics::~TQueryMetrics() noexcept { - End(EStatus::CLIENT_INTERNAL_ERROR); -} - -void TQueryMetrics::End(EStatus status) noexcept { - if (Ended_) { - return; - } - Ended_ = true; - - try { - if (LatencyHistogram_) { - auto durationMs = (TInstant::Now() - StartTime_).MilliSeconds(); - LatencyHistogram_->Record(static_cast(durationMs)); - } - if (status != EStatus::SUCCESS && ErrorCounter_) { - ErrorCounter_->Inc(); - } - } catch (...) { - SafeLogMetricsError("failed to record metrics"); - } -} - -} // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_metrics.h b/src/client/query/impl/query_metrics.h index 807472e2434..841e3212f14 100644 --- a/src/client/query/impl/query_metrics.h +++ b/src/client/query/impl/query_metrics.h @@ -1,28 +1,14 @@ #pragma once -#include -#include - -#include - -#include -#include +#include namespace NYdb::inline V3::NQuery { -class TQueryMetrics { +class TQueryMetrics : public NObservability::TClientMetrics { public: - TQueryMetrics(std::shared_ptr registry, const std::string& operationName); - ~TQueryMetrics() noexcept; - - void End(EStatus status) noexcept; - -private: - std::shared_ptr RequestCounter_; - std::shared_ptr ErrorCounter_; - std::shared_ptr LatencyHistogram_; - TInstant StartTime_; - bool Ended_ = false; + TQueryMetrics(std::shared_ptr registry, const std::string& operationName) + : TClientMetrics(std::move(registry), "ydb.query", operationName) + {} }; } // namespace NYdb::NQuery diff --git a/src/client/table/impl/CMakeLists.txt b/src/client/table/impl/CMakeLists.txt index 8f53d386fc6..8ecfe4ead87 100644 --- a/src/client/table/impl/CMakeLists.txt +++ b/src/client/table/impl/CMakeLists.txt @@ -10,6 +10,8 @@ target_link_libraries(client-ydb_table-impl client-impl-ydb_endpoints impl-session client-ydb_table-query_stats + client-metrics + impl-observability PRIVATE OpenSSL::SSL ) diff --git a/src/client/table/impl/table_metrics.h b/src/client/table/impl/table_metrics.h new file mode 100644 index 00000000000..5bf6128a6ea --- /dev/null +++ b/src/client/table/impl/table_metrics.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +namespace NYdb::inline V3::NTable { + +class TTableMetrics : public NObservability::TClientMetrics { +public: + TTableMetrics(std::shared_ptr registry, const std::string& operationName) + : TClientMetrics(std::move(registry), "ydb.table", operationName) + {} +}; + +} // namespace NYdb::NTable diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index cab8f71cdc4..21c2398954c 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -39,12 +39,14 @@ std::shared_ptr GetCounter( return registry->GetCounter(name, {{"operation", operation}}); } -std::shared_ptr GetHistogram( +std::shared_ptr GetDuration( const std::shared_ptr& registry, - const std::string& name, const std::string& operation) { - return registry->GetHistogram(name, {{"operation", operation}}); + return registry->GetHistogram("db.client.operation.duration", { + {"db.system.name", "ydb"}, + {"db.operation.name", operation}, + }); } } // namespace @@ -70,10 +72,10 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); - ASSERT_NE(latency, nullptr) << "ExecuteQuery latency histogram not created"; - EXPECT_GE(latency->Count(), 1u); - for (double v : latency->GetValues()) { + auto duration = GetDuration(registry, "ExecuteQuery"); + ASSERT_NE(duration, nullptr) << "ExecuteQuery duration histogram not created"; + EXPECT_GE(duration->Count(), 1u); + for (double v : duration->GetValues()) { EXPECT_GE(v, 0.0); } @@ -101,9 +103,9 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ASSERT_NE(errors, nullptr); EXPECT_GE(errors->Get(), 1); - auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); - ASSERT_NE(latency, nullptr); - EXPECT_GE(latency->Count(), 1u); + auto duration = GetDuration(registry, "ExecuteQuery"); + ASSERT_NE(duration, nullptr); + EXPECT_GE(duration->Count(), 1u); driver.Stop(true); } @@ -119,9 +121,9 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); - auto latency = GetHistogram(registry, "ydb.query.latency_ms", "CreateSession"); - ASSERT_NE(latency, nullptr) << "CreateSession latency histogram not created"; - EXPECT_GE(latency->Count(), 1u); + auto duration = GetDuration(registry, "CreateSession"); + ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; + EXPECT_GE(duration->Count(), 1u); driver.Stop(true); } @@ -152,9 +154,9 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; EXPECT_GE(commitRequests->Get(), 1); - auto commitLatency = GetHistogram(registry, "ydb.query.latency_ms", "Commit"); - ASSERT_NE(commitLatency, nullptr); - EXPECT_GE(commitLatency->Count(), 1u); + auto commitDuration = GetDuration(registry, "Commit"); + ASSERT_NE(commitDuration, nullptr); + EXPECT_GE(commitDuration->Count(), 1u); } driver.Stop(true); @@ -183,9 +185,9 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { ASSERT_NE(rollbackErrors, nullptr); EXPECT_EQ(rollbackErrors->Get(), 0); - auto rollbackLatency = GetHistogram(registry, "ydb.query.latency_ms", "Rollback"); - ASSERT_NE(rollbackLatency, nullptr); - EXPECT_GE(rollbackLatency->Count(), 1u); + auto rollbackDuration = GetDuration(registry, "Rollback"); + ASSERT_NE(rollbackDuration, nullptr); + EXPECT_GE(rollbackDuration->Count(), 1u); driver.Stop(true); } @@ -215,9 +217,9 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); - ASSERT_NE(latency, nullptr); - EXPECT_EQ(latency->Count(), static_cast(numQueries)); + auto duration = GetDuration(registry, "ExecuteQuery"); + ASSERT_NE(duration, nullptr); + EXPECT_EQ(duration->Count(), static_cast(numQueries)); driver.Stop(true); } @@ -246,7 +248,7 @@ TEST(QueryMetricsIntegration, NoRegistryDoesNotBreakOperations) { driver.Stop(true); } -TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { +TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { auto [driver, registry] = MakeRunArgs(); TQueryClient client(driver); @@ -260,13 +262,13 @@ TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); - auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); - ASSERT_NE(latency, nullptr); - ASSERT_GE(latency->Count(), 1u); + auto duration = GetDuration(registry, "ExecuteQuery"); + ASSERT_NE(duration, nullptr); + ASSERT_GE(duration->Count(), 1u); - for (double v : latency->GetValues()) { - EXPECT_GE(v, 0.0) << "Latency must be non-negative"; - EXPECT_LT(v, 30000.0) << "Latency > 30s is unrealistic for SELECT 1"; + for (double v : duration->GetValues()) { + EXPECT_GE(v, 0.0) << "Duration must be non-negative"; + EXPECT_LT(v, 30.0) << "Duration > 30s is unrealistic for SELECT 1"; } driver.Stop(true); diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index 2ad5a38e01f..de86c3fe274 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -101,14 +101,16 @@ add_ydb_test(NAME client-ydb_value_ut GTEST unit ) -add_ydb_test(NAME client-ydb_query_metrics_ut GTEST +add_ydb_test(NAME client-ydb_metrics_ut GTEST INCLUDE_DIRS ${YDB_SDK_SOURCE_DIR} SOURCES - query/query_metrics_ut.cpp + observability/client_metrics_ut.cpp LINK_LIBRARIES yutil + impl-observability client-ydb_query-impl + client-ydb_table-impl client-metrics LABELS unit diff --git a/tests/unit/client/observability/client_metrics_ut.cpp b/tests/unit/client/observability/client_metrics_ut.cpp new file mode 100644 index 00000000000..3dbedfa801d --- /dev/null +++ b/tests/unit/client/observability/client_metrics_ut.cpp @@ -0,0 +1,229 @@ +#include +#include +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NObservability; +using namespace NYdb::NMetrics; +using namespace NYdb::NTests; + +// --------------------------------------------------------------------------- +// TClientMetrics (shared logic) +// --------------------------------------------------------------------------- + +class ClientMetricsTest : public ::testing::Test { +protected: + void SetUp() override { + Registry = std::make_shared(); + } + + std::shared_ptr RequestCounter(const std::string& op) { + return Registry->GetCounter(Prefix + ".requests", {{"operation", op}}); + } + + std::shared_ptr ErrorCounter(const std::string& op) { + return Registry->GetCounter(Prefix + ".errors", {{"operation", op}}); + } + + std::shared_ptr DurationHistogram(const std::string& op) { + return Registry->GetHistogram("db.client.operation.duration", { + {"db.system.name", "ydb"}, + {"db.operation.name", op}, + }); + } + + const std::string Prefix = "ydb.test"; + std::shared_ptr Registry; +}; + +TEST_F(ClientMetricsTest, RequestCounterIncrementedOnConstruction) { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + + auto counter = RequestCounter("DoSomething"); + ASSERT_NE(counter, nullptr); + EXPECT_EQ(counter->Get(), 1); +} + +TEST_F(ClientMetricsTest, SuccessDoesNotIncrementErrorCounter) { + { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + metrics.End(EStatus::SUCCESS); + } + + auto errors = ErrorCounter("DoSomething"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); +} + +TEST_F(ClientMetricsTest, FailureIncrementsErrorCounter) { + { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + metrics.End(EStatus::UNAVAILABLE); + } + + auto errors = ErrorCounter("DoSomething"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); +} + +TEST_F(ClientMetricsTest, DurationRecordedOnEnd) { + { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + metrics.End(EStatus::SUCCESS); + } + + auto hist = DurationHistogram("DoSomething"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); + EXPECT_GE(hist->GetValues()[0], 0.0); +} + +TEST_F(ClientMetricsTest, DurationIsInSeconds) { + { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + metrics.End(EStatus::SUCCESS); + } + + auto hist = DurationHistogram("DoSomething"); + ASSERT_NE(hist, nullptr); + EXPECT_LT(hist->GetValues()[0], 1.0); +} + +TEST_F(ClientMetricsTest, DoubleEndIsIdempotent) { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + metrics.End(EStatus::SUCCESS); + metrics.End(EStatus::INTERNAL_ERROR); + + auto errors = ErrorCounter("DoSomething"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto hist = DurationHistogram("DoSomething"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(ClientMetricsTest, DestructorCallsEndWithClientInternalError) { + { + TClientMetrics metrics(Registry, Prefix, "DoSomething"); + } + + auto requests = RequestCounter("DoSomething"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), 1); + + auto errors = ErrorCounter("DoSomething"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); + + auto hist = DurationHistogram("DoSomething"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(ClientMetricsTest, NullRegistryDoesNotCrash) { + EXPECT_NO_THROW({ + TClientMetrics metrics(nullptr, Prefix, "DoSomething"); + metrics.End(EStatus::SUCCESS); + }); +} + +TEST_F(ClientMetricsTest, DifferentOperationsHaveSeparateMetrics) { + { + TClientMetrics m1(Registry, Prefix, "OpA"); + m1.End(EStatus::SUCCESS); + } + { + TClientMetrics m2(Registry, Prefix, "OpB"); + m2.End(EStatus::OVERLOADED); + } + + EXPECT_EQ(RequestCounter("OpA")->Get(), 1); + EXPECT_EQ(RequestCounter("OpB")->Get(), 1); + EXPECT_EQ(ErrorCounter("OpA")->Get(), 0); + EXPECT_EQ(ErrorCounter("OpB")->Get(), 1); + EXPECT_EQ(DurationHistogram("OpA")->Count(), 1u); + EXPECT_EQ(DurationHistogram("OpB")->Count(), 1u); +} + +TEST_F(ClientMetricsTest, MultipleRequestsAccumulate) { + for (int i = 0; i < 5; ++i) { + TClientMetrics metrics(Registry, Prefix, "Op"); + metrics.End(i % 2 == 0 ? EStatus::SUCCESS : EStatus::TIMEOUT); + } + + EXPECT_EQ(RequestCounter("Op")->Get(), 5); + EXPECT_EQ(ErrorCounter("Op")->Get(), 2); + EXPECT_EQ(DurationHistogram("Op")->Count(), 5u); +} + +TEST_F(ClientMetricsTest, AllErrorStatusesIncrementErrorCounter) { + std::vector errorStatuses = { + EStatus::BAD_REQUEST, + EStatus::UNAUTHORIZED, + EStatus::INTERNAL_ERROR, + EStatus::UNAVAILABLE, + EStatus::OVERLOADED, + EStatus::TIMEOUT, + EStatus::NOT_FOUND, + EStatus::CLIENT_INTERNAL_ERROR, + }; + + for (auto status : errorStatuses) { + TClientMetrics metrics(Registry, Prefix, "Op"); + metrics.End(status); + } + + auto errors = ErrorCounter("Op"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), static_cast(errorStatuses.size())); +} + +TEST_F(ClientMetricsTest, PrefixAppliedToCounterNames) { + TClientMetrics metrics(Registry, "ydb.custom", "Op"); + metrics.End(EStatus::SUCCESS); + + EXPECT_NE(Registry->GetCounter("ydb.custom.requests", {{"operation", "Op"}}), nullptr); + EXPECT_NE(Registry->GetCounter("ydb.custom.errors", {{"operation", "Op"}}), nullptr); + + EXPECT_EQ(Registry->GetCounter("ydb.test.requests", {{"operation", "Op"}}), nullptr); +} + +// --------------------------------------------------------------------------- +// TQueryMetrics prefix +// --------------------------------------------------------------------------- + +TEST(QueryMetricsTest, UsesQueryPrefix) { + auto registry = std::make_shared(); + + NQuery::TQueryMetrics metrics(registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + + EXPECT_NE(registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(registry->GetCounter("ydb.query.errors", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(registry->GetHistogram("db.client.operation.duration", { + {"db.system.name", "ydb"}, {"db.operation.name", "ExecuteQuery"}}), nullptr); + + EXPECT_EQ(registry->GetCounter("ydb.table.requests", {{"operation", "ExecuteQuery"}}), nullptr); +} + +// --------------------------------------------------------------------------- +// TTableMetrics prefix +// --------------------------------------------------------------------------- + +TEST(TableMetricsTest, UsesTablePrefix) { + auto registry = std::make_shared(); + + NTable::TTableMetrics metrics(registry, "ExecuteDataQuery"); + metrics.End(EStatus::SUCCESS); + + EXPECT_NE(registry->GetCounter("ydb.table.requests", {{"operation", "ExecuteDataQuery"}}), nullptr); + EXPECT_NE(registry->GetCounter("ydb.table.errors", {{"operation", "ExecuteDataQuery"}}), nullptr); + EXPECT_NE(registry->GetHistogram("db.client.operation.duration", { + {"db.system.name", "ydb"}, {"db.operation.name", "ExecuteDataQuery"}}), nullptr); + + EXPECT_EQ(registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteDataQuery"}}), nullptr); +} diff --git a/tests/unit/client/query/query_metrics_ut.cpp b/tests/unit/client/query/query_metrics_ut.cpp deleted file mode 100644 index 20c681b7eca..00000000000 --- a/tests/unit/client/query/query_metrics_ut.cpp +++ /dev/null @@ -1,190 +0,0 @@ -#include -#include - -#include - -using namespace NYdb; -using namespace NYdb::NQuery; -using namespace NYdb::NMetrics; -using namespace NYdb::NTests; - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -class QueryMetricsTest : public ::testing::Test { -protected: - void SetUp() override { - Registry = std::make_shared(); - } - - std::shared_ptr RequestCounter(const std::string& op) { - return Registry->GetCounter("ydb.query.requests", {{"operation", op}}); - } - - std::shared_ptr ErrorCounter(const std::string& op) { - return Registry->GetCounter("ydb.query.errors", {{"operation", op}}); - } - - std::shared_ptr LatencyHistogram(const std::string& op) { - return Registry->GetHistogram("ydb.query.latency_ms", {{"operation", op}}); - } - - std::shared_ptr Registry; -}; - -TEST_F(QueryMetricsTest, RequestCounterIncrementedOnConstruction) { - TQueryMetrics metrics(Registry, "ExecuteQuery"); - - auto counter = RequestCounter("ExecuteQuery"); - ASSERT_NE(counter, nullptr); - EXPECT_EQ(counter->Get(), 1); -} - -TEST_F(QueryMetricsTest, SuccessDoesNotIncrementErrorCounter) { - { - TQueryMetrics metrics(Registry, "ExecuteQuery"); - metrics.End(EStatus::SUCCESS); - } - - auto errors = ErrorCounter("ExecuteQuery"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), 0); -} - -TEST_F(QueryMetricsTest, FailureIncrementsErrorCounter) { - { - TQueryMetrics metrics(Registry, "Commit"); - metrics.End(EStatus::UNAVAILABLE); - } - - auto errors = ErrorCounter("Commit"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), 1); -} - -TEST_F(QueryMetricsTest, LatencyRecordedOnEnd) { - { - TQueryMetrics metrics(Registry, "Rollback"); - metrics.End(EStatus::SUCCESS); - } - - auto hist = LatencyHistogram("Rollback"); - ASSERT_NE(hist, nullptr); - EXPECT_EQ(hist->Count(), 1u); - EXPECT_GE(hist->GetValues()[0], 0.0); -} - -TEST_F(QueryMetricsTest, DoubleEndIsIdempotent) { - TQueryMetrics metrics(Registry, "ExecuteQuery"); - metrics.End(EStatus::SUCCESS); - metrics.End(EStatus::INTERNAL_ERROR); - - auto errors = ErrorCounter("ExecuteQuery"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), 0); - - auto hist = LatencyHistogram("ExecuteQuery"); - ASSERT_NE(hist, nullptr); - EXPECT_EQ(hist->Count(), 1u); -} - -TEST_F(QueryMetricsTest, DestructorCallsEndWithClientInternalError) { - { - TQueryMetrics metrics(Registry, "CreateSession"); - } - - auto requests = RequestCounter("CreateSession"); - ASSERT_NE(requests, nullptr); - EXPECT_EQ(requests->Get(), 1); - - auto errors = ErrorCounter("CreateSession"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), 1); - - auto hist = LatencyHistogram("CreateSession"); - ASSERT_NE(hist, nullptr); - EXPECT_EQ(hist->Count(), 1u); -} - -TEST_F(QueryMetricsTest, NullRegistryDoesNotCrash) { - EXPECT_NO_THROW({ - TQueryMetrics metrics(nullptr, "ExecuteQuery"); - metrics.End(EStatus::SUCCESS); - }); -} - -TEST_F(QueryMetricsTest, CorrectMetricNamesAndLabels) { - TQueryMetrics metrics(Registry, "ExecuteQuery"); - metrics.End(EStatus::SUCCESS); - - EXPECT_NE(Registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteQuery"}}), nullptr); - EXPECT_NE(Registry->GetCounter("ydb.query.errors", {{"operation", "ExecuteQuery"}}), nullptr); - EXPECT_NE(Registry->GetHistogram("ydb.query.latency_ms", {{"operation", "ExecuteQuery"}}), nullptr); - - EXPECT_EQ(Registry->GetCounter("ydb.query.requests", {{"operation", "Commit"}}), nullptr); -} - -TEST_F(QueryMetricsTest, DifferentOperationsHaveSeparateMetrics) { - { - TQueryMetrics m1(Registry, "ExecuteQuery"); - m1.End(EStatus::SUCCESS); - } - { - TQueryMetrics m2(Registry, "Commit"); - m2.End(EStatus::OVERLOADED); - } - - auto execRequests = RequestCounter("ExecuteQuery"); - auto commitRequests = RequestCounter("Commit"); - ASSERT_NE(execRequests, nullptr); - ASSERT_NE(commitRequests, nullptr); - EXPECT_EQ(execRequests->Get(), 1); - EXPECT_EQ(commitRequests->Get(), 1); - - auto execErrors = ErrorCounter("ExecuteQuery"); - auto commitErrors = ErrorCounter("Commit"); - EXPECT_EQ(execErrors->Get(), 0); - EXPECT_EQ(commitErrors->Get(), 1); -} - -TEST_F(QueryMetricsTest, MultipleRequestsAccumulate) { - for (int i = 0; i < 5; ++i) { - TQueryMetrics metrics(Registry, "ExecuteQuery"); - metrics.End(i % 2 == 0 ? EStatus::SUCCESS : EStatus::TIMEOUT); - } - - auto requests = RequestCounter("ExecuteQuery"); - ASSERT_NE(requests, nullptr); - EXPECT_EQ(requests->Get(), 5); - - auto errors = ErrorCounter("ExecuteQuery"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), 2); - - auto hist = LatencyHistogram("ExecuteQuery"); - ASSERT_NE(hist, nullptr); - EXPECT_EQ(hist->Count(), 5u); -} - -TEST_F(QueryMetricsTest, AllErrorStatusesIncrementErrorCounter) { - std::vector errorStatuses = { - EStatus::BAD_REQUEST, - EStatus::UNAUTHORIZED, - EStatus::INTERNAL_ERROR, - EStatus::UNAVAILABLE, - EStatus::OVERLOADED, - EStatus::TIMEOUT, - EStatus::NOT_FOUND, - EStatus::CLIENT_INTERNAL_ERROR, - }; - - for (auto status : errorStatuses) { - TQueryMetrics metrics(Registry, "Rollback"); - metrics.End(status); - } - - auto errors = ErrorCounter("Rollback"); - ASSERT_NE(errors, nullptr); - EXPECT_EQ(errors->Get(), static_cast(errorStatuses.size())); -} From c7e65b5c3c7b774c590f6d1ae45cff008c7cdd1f Mon Sep 17 00:00:00 2001 From: maladetska Date: Sun, 29 Mar 2026 23:43:25 +0300 Subject: [PATCH 08/93] fix semconv --- include/ydb-cpp-sdk/client/metrics/metrics.h | 22 ++- plugins/metrics/otel/src/metrics.cpp | 31 ++-- .../impl/observability/client_metrics.cpp | 83 +++++----- .../impl/observability/client_metrics.h | 8 +- src/client/query/impl/query_metrics.h | 2 +- src/client/query/impl/query_spans.cpp | 7 +- src/client/table/impl/table_metrics.h | 2 +- tests/common/fake_metric_registry.h | 19 ++- tests/integration/metrics/main.cpp | 54 ++++--- .../observability/client_metrics_ut.cpp | 144 +++++++++++------- 10 files changed, 226 insertions(+), 146 deletions(-) diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index 7e2b0b903dd..5faa930ed50 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -33,9 +33,25 @@ class IMetricRegistry { public: virtual ~IMetricRegistry() = default; - virtual std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) = 0; - virtual std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) = 0; - virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; + virtual std::shared_ptr Counter( + const std::string& name, + const TLabels& labels = {}, + const std::string& description = {}, + const std::string& unit = {} + ) = 0; + virtual std::shared_ptr Gauge( + const std::string& name, + const TLabels& labels = {}, + const std::string& description = {}, + const std::string& unit = {} + ) = 0; + virtual std::shared_ptr Histogram( + const std::string& name, + const std::vector& buckets, + const TLabels& labels = {}, + const std::string& description = {}, + const std::string& unit = {} + ) = 0; }; } // namespace NYdb::NMetrics diff --git a/plugins/metrics/otel/src/metrics.cpp b/plugins/metrics/otel/src/metrics.cpp index 6b9f14be362..af07af89bb8 100644 --- a/plugins/metrics/otel/src/metrics.cpp +++ b/plugins/metrics/otel/src/metrics.cpp @@ -81,24 +81,37 @@ class TOtelMetricRegistry : public IMetricRegistry { , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) {} - std::shared_ptr Counter(const std::string& name, const TLabels& labels) override { - auto counter = Meter_->CreateUInt64Counter(name); + std::shared_ptr Counter(const std::string& name + , const TLabels& labels + , const std::string& description + , const std::string& unit + ) override { + auto counter = Meter_->CreateUInt64Counter(name, description, unit); return std::make_shared(std::move(counter), labels); } - std::shared_ptr Gauge(const std::string& name, const TLabels& labels) override { - auto counter = Meter_->CreateDoubleUpDownCounter(name); + std::shared_ptr Gauge(const std::string& name + , const TLabels& labels + , const std::string& description + , const std::string& unit + ) override { + auto counter = Meter_->CreateDoubleUpDownCounter(name, description, unit); return std::make_shared(std::move(counter), labels); } - std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) override { - ConfigureHistogramBuckets(name, buckets); - auto histogram = Meter_->CreateDoubleHistogram(name); + std::shared_ptr Histogram(const std::string& name + , const std::vector& buckets + , const TLabels& labels + , const std::string& description + , const std::string& unit + ) override { + ConfigureHistogramBuckets(name, unit, buckets); + auto histogram = Meter_->CreateDoubleHistogram(name, description, unit); return std::make_shared(std::move(histogram), labels); } private: - void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { + void ConfigureHistogramBuckets(const std::string& name, const std::string& unit, const std::vector& buckets) { if (buckets.empty()) { return; } @@ -118,7 +131,7 @@ class TOtelMetricRegistry : public IMetricRegistry { auto selector = std::make_unique( sdk::metrics::InstrumentType::kHistogram, name, - "" + unit ); auto meterSelector = std::make_unique( "ydb-cpp-sdk", diff --git a/src/client/impl/observability/client_metrics.cpp b/src/client/impl/observability/client_metrics.cpp index efa9b739517..f66605bc772 100644 --- a/src/client/impl/observability/client_metrics.cpp +++ b/src/client/impl/observability/client_metrics.cpp @@ -1,6 +1,7 @@ #include "client_metrics.h" #include +#include namespace NYdb::inline V3::NObservability { @@ -18,65 +19,32 @@ void SafeLogMetricsError(const char* /*message*/) noexcept { } } -std::string StatusToString(EStatus status) { - switch (status) { - case EStatus::SUCCESS: return "SUCCESS"; - case EStatus::BAD_REQUEST: return "BAD_REQUEST"; - case EStatus::UNAUTHORIZED: return "UNAUTHORIZED"; - case EStatus::INTERNAL_ERROR: return "INTERNAL_ERROR"; - case EStatus::ABORTED: return "ABORTED"; - case EStatus::UNAVAILABLE: return "UNAVAILABLE"; - case EStatus::OVERLOADED: return "OVERLOADED"; - case EStatus::SCHEME_ERROR: return "SCHEME_ERROR"; - case EStatus::GENERIC_ERROR: return "GENERIC_ERROR"; - case EStatus::TIMEOUT: return "TIMEOUT"; - case EStatus::BAD_SESSION: return "BAD_SESSION"; - case EStatus::PRECONDITION_FAILED: return "PRECONDITION_FAILED"; - case EStatus::ALREADY_EXISTS: return "ALREADY_EXISTS"; - case EStatus::NOT_FOUND: return "NOT_FOUND"; - case EStatus::SESSION_EXPIRED: return "SESSION_EXPIRED"; - case EStatus::CANCELLED: return "CANCELLED"; - case EStatus::UNDETERMINED: return "UNDETERMINED"; - case EStatus::UNSUPPORTED: return "UNSUPPORTED"; - case EStatus::SESSION_BUSY: return "SESSION_BUSY"; - case EStatus::EXTERNAL_ERROR: return "EXTERNAL_ERROR"; - case EStatus::TRANSPORT_UNAVAILABLE: return "TRANSPORT_UNAVAILABLE"; - case EStatus::CLIENT_RESOURCE_EXHAUSTED:return "CLIENT_RESOURCE_EXHAUSTED"; - case EStatus::CLIENT_DEADLINE_EXCEEDED: return "CLIENT_DEADLINE_EXCEEDED"; - case EStatus::CLIENT_INTERNAL_ERROR: return "CLIENT_INTERNAL_ERROR"; - case EStatus::CLIENT_CANCELLED: return "CLIENT_CANCELLED"; - case EStatus::CLIENT_UNAUTHENTICATED: return "CLIENT_UNAUTHENTICATED"; - case EStatus::CLIENT_CALL_UNIMPLEMENTED:return "CLIENT_CALL_UNIMPLEMENTED"; - case EStatus::CLIENT_OUT_OF_RANGE: return "CLIENT_OUT_OF_RANGE"; - case EStatus::CLIENT_DISCOVERY_FAILED: return "CLIENT_DISCOVERY_FAILED"; - case EStatus::CLIENT_LIMITS_REACHED: return "CLIENT_LIMITS_REACHED"; - default: return "STATUS_UNDEFINED"; - } -} - } // namespace static const std::vector DurationBucketsSec = { 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10 }; -TClientMetrics::TClientMetrics(std::shared_ptr registry, - const std::string& prefix, const std::string& operationName) +static constexpr const char* RequestsDescription = "Number of database client operations started."; +static constexpr const char* ErrorsDescription = "Number of database client operations that failed."; +static constexpr const char* DurationDescription = "Duration of database client operations."; + +TClientMetrics::TClientMetrics(std::shared_ptr registry + , const std::string& operationName +) : Registry_(std::move(registry)) + , OperationName_(operationName) { - if (!registry) { + if (!Registry_) { return; } try { - NMetrics::TLabels labels = {{"operation", operationName}}; - RequestCounter_ = registry->Counter(prefix + ".requests", labels); - ErrorCounter_ = registry->Counter(prefix + ".errors", labels); - - NMetrics::TLabels durationLabels = { - {"db.system.name", "ydb"}, + NMetrics::TLabels labels = { + {"db.system.name", "other_sql"}, {"db.operation.name", operationName}, }; - DurationHistogram_ = registry->Histogram("db.client.operation.duration", DurationBucketsSec, durationLabels); + RequestCounter_ = Registry_->Counter("db.client.operation.requests", labels, RequestsDescription, "{operation}"); + ErrorCounter_ = Registry_->Counter("db.client.operation.errors", labels, ErrorsDescription, "{error}"); RequestCounter_->Inc(); StartTime_ = std::chrono::steady_clock::now(); @@ -84,7 +52,7 @@ TClientMetrics::TClientMetrics(std::shared_ptr regist SafeLogMetricsError("failed to initialize metrics"); RequestCounter_.reset(); ErrorCounter_.reset(); - DurationHistogram_.reset(); + Registry_.reset(); } } @@ -99,10 +67,27 @@ void TClientMetrics::End(EStatus status) noexcept { Ended_ = true; try { - if (DurationHistogram_) { + const std::string statusCode = ToString(status); + if (Registry_) { auto elapsed = std::chrono::steady_clock::now() - StartTime_; double durationSec = std::chrono::duration(elapsed).count(); - DurationHistogram_->Record(durationSec); + NMetrics::TLabels durationLabels = { + {"db.system.name", "other_sql"}, + {"db.operation.name", OperationName_}, + {"db.response.status_code", statusCode}, + }; + if (status != EStatus::SUCCESS) { + durationLabels["error.type"] = statusCode; + } + auto durationHistogram = Registry_->Histogram( + "db.client.operation.duration", + DurationBucketsSec, + durationLabels, + DurationDescription, + "s"); + if (durationHistogram) { + durationHistogram->Record(durationSec); + } } if (status != EStatus::SUCCESS && ErrorCounter_) { diff --git a/src/client/impl/observability/client_metrics.h b/src/client/impl/observability/client_metrics.h index bce81a958f2..e12a30cad60 100644 --- a/src/client/impl/observability/client_metrics.h +++ b/src/client/impl/observability/client_metrics.h @@ -11,16 +11,18 @@ namespace NYdb::inline V3::NObservability { class TClientMetrics { public: - TClientMetrics(std::shared_ptr registry, - const std::string& prefix, const std::string& operationName); + TClientMetrics(std::shared_ptr registry + , const std::string& operationName + ); ~TClientMetrics() noexcept; void End(EStatus status) noexcept; private: + std::shared_ptr Registry_; + std::string OperationName_; std::shared_ptr RequestCounter_; std::shared_ptr ErrorCounter_; - std::shared_ptr DurationHistogram_; std::chrono::steady_clock::time_point StartTime_; bool Ended_ = false; }; diff --git a/src/client/query/impl/query_metrics.h b/src/client/query/impl/query_metrics.h index 841e3212f14..2bd284f76f3 100644 --- a/src/client/query/impl/query_metrics.h +++ b/src/client/query/impl/query_metrics.h @@ -7,7 +7,7 @@ namespace NYdb::inline V3::NQuery { class TQueryMetrics : public NObservability::TClientMetrics { public: TQueryMetrics(std::shared_ptr registry, const std::string& operationName) - : TClientMetrics(std::move(registry), "ydb.query", operationName) + : TClientMetrics(std::move(registry), operationName) {} }; diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index 4bbd4d2250b..df4a5718771 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -65,11 +65,12 @@ TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::str ParseEndpoint(endpoint, host, port); try { - Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); + Span_ = tracer->StartSpan(operationName, NMetrics::ESpanKind::CLIENT); if (!Span_) { return; } - Span_->SetAttribute("db.system.name", "ydb"); + Span_->SetAttribute("db.system.name", "other_sql"); + Span_->SetAttribute("db.operation.name", operationName); Span_->SetAttribute("server.address", host); Span_->SetAttribute("server.port", static_cast(port)); } catch (...) { @@ -128,7 +129,7 @@ void TQuerySpan::AddEvent(const std::string& name, const std::mapSetAttribute("db.response.status_code", static_cast(status)); + Span_->SetAttribute("db.response.status_code", ToString(status)); if (status != EStatus::SUCCESS) { Span_->SetAttribute("error.type", ToString(status)); } diff --git a/src/client/table/impl/table_metrics.h b/src/client/table/impl/table_metrics.h index 5bf6128a6ea..83f9deafdb4 100644 --- a/src/client/table/impl/table_metrics.h +++ b/src/client/table/impl/table_metrics.h @@ -7,7 +7,7 @@ namespace NYdb::inline V3::NTable { class TTableMetrics : public NObservability::TClientMetrics { public: TTableMetrics(std::shared_ptr registry, const std::string& operationName) - : TClientMetrics(std::move(registry), "ydb.table", operationName) + : TClientMetrics(std::move(registry), operationName) {} }; diff --git a/tests/common/fake_metric_registry.h b/tests/common/fake_metric_registry.h index 60ff1414633..032234f080f 100644 --- a/tests/common/fake_metric_registry.h +++ b/tests/common/fake_metric_registry.h @@ -68,7 +68,11 @@ struct TMetricKey { class TFakeMetricRegistry : public NMetrics::IMetricRegistry { public: - std::shared_ptr Counter(const std::string& name, const NMetrics::TLabels& labels) override { + std::shared_ptr Counter(const std::string& name + , const NMetrics::TLabels& labels + , const std::string& /*description*/ + , const std::string& /*unit*/ + ) override { std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Counters_.find(key); @@ -80,7 +84,11 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return counter; } - std::shared_ptr Gauge(const std::string& name, const NMetrics::TLabels& labels) override { + std::shared_ptr Gauge(const std::string& name + , const NMetrics::TLabels& labels + , const std::string& /*description*/ + , const std::string& /*unit*/ + ) override { std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto gauge = std::make_shared(); @@ -88,7 +96,12 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return gauge; } - std::shared_ptr Histogram(const std::string& name, const std::vector& /*buckets*/, const NMetrics::TLabels& labels) override { + std::shared_ptr Histogram(const std::string& name + , const std::vector& /*buckets*/ + , const NMetrics::TLabels& labels + , const std::string& /*description*/ + , const std::string& /*unit*/ + ) override { std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Histograms_.find(key); diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index 21c2398954c..fec3aab583b 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include @@ -36,17 +37,26 @@ std::shared_ptr GetCounter( const std::string& name, const std::string& operation) { - return registry->GetCounter(name, {{"operation", operation}}); + return registry->GetCounter(name, { + {"db.system.name", "other_sql"}, + {"db.operation.name", operation}, + }); } std::shared_ptr GetDuration( const std::shared_ptr& registry, - const std::string& operation) + const std::string& operation, + EStatus status) { - return registry->GetHistogram("db.client.operation.duration", { - {"db.system.name", "ydb"}, + NMetrics::TLabels labels = { + {"db.system.name", "other_sql"}, {"db.operation.name", operation}, - }); + {"db.response.status_code", ToString(status)}, + }; + if (status != EStatus::SUCCESS) { + labels["error.type"] = ToString(status); + } + return registry->GetHistogram("db.client.operation.duration", labels); } } // namespace @@ -64,15 +74,15 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); - auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; EXPECT_GE(requests->Get(), 1); - auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, "db.client.operation.errors", "ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto duration = GetDuration(registry, "ExecuteQuery"); + auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "ExecuteQuery duration histogram not created"; EXPECT_GE(duration->Count(), 1u); for (double v : duration->GetValues()) { @@ -95,15 +105,15 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ).ExtractValueSync(); EXPECT_NE(result.GetStatus(), EStatus::SUCCESS); - auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_GE(requests->Get(), 1); - auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, "db.client.operation.errors", "ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_GE(errors->Get(), 1); - auto duration = GetDuration(registry, "ExecuteQuery"); + auto duration = GetDuration(registry, "ExecuteQuery", result.GetStatus()); ASSERT_NE(duration, nullptr); EXPECT_GE(duration->Count(), 1u); @@ -117,11 +127,11 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto session = client.GetSession().ExtractValueSync(); ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); - auto requests = GetCounter(registry, "ydb.query.requests", "CreateSession"); + auto requests = GetCounter(registry, "db.client.operation.requests", "CreateSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); - auto duration = GetDuration(registry, "CreateSession"); + auto duration = GetDuration(registry, "CreateSession", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; EXPECT_GE(duration->Count(), 1u); @@ -150,11 +160,11 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitResult = execResult.GetTransaction()->Commit().ExtractValueSync(); ASSERT_TRUE(commitResult.IsSuccess()) << commitResult.GetIssues().ToString(); - auto commitRequests = GetCounter(registry, "ydb.query.requests", "Commit"); + auto commitRequests = GetCounter(registry, "db.client.operation.requests", "Commit"); ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; EXPECT_GE(commitRequests->Get(), 1); - auto commitDuration = GetDuration(registry, "Commit"); + auto commitDuration = GetDuration(registry, "Commit", EStatus::SUCCESS); ASSERT_NE(commitDuration, nullptr); EXPECT_GE(commitDuration->Count(), 1u); } @@ -177,15 +187,15 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackResult = tx.Rollback().ExtractValueSync(); ASSERT_TRUE(rollbackResult.IsSuccess()) << rollbackResult.GetIssues().ToString(); - auto rollbackRequests = GetCounter(registry, "ydb.query.requests", "Rollback"); + auto rollbackRequests = GetCounter(registry, "db.client.operation.requests", "Rollback"); ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; EXPECT_GE(rollbackRequests->Get(), 1); - auto rollbackErrors = GetCounter(registry, "ydb.query.errors", "Rollback"); + auto rollbackErrors = GetCounter(registry, "db.client.operation.errors", "Rollback"); ASSERT_NE(rollbackErrors, nullptr); EXPECT_EQ(rollbackErrors->Get(), 0); - auto rollbackDuration = GetDuration(registry, "Rollback"); + auto rollbackDuration = GetDuration(registry, "Rollback", EStatus::SUCCESS); ASSERT_NE(rollbackDuration, nullptr); EXPECT_GE(rollbackDuration->Count(), 1u); @@ -209,15 +219,15 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); } - auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_EQ(requests->Get(), numQueries); - auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, "db.client.operation.errors", "ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto duration = GetDuration(registry, "ExecuteQuery"); + auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); EXPECT_EQ(duration->Count(), static_cast(numQueries)); @@ -262,7 +272,7 @@ TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); - auto duration = GetDuration(registry, "ExecuteQuery"); + auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); ASSERT_GE(duration->Count(), 1u); diff --git a/tests/unit/client/observability/client_metrics_ut.cpp b/tests/unit/client/observability/client_metrics_ut.cpp index 3dbedfa801d..209ce1db0d3 100644 --- a/tests/unit/client/observability/client_metrics_ut.cpp +++ b/tests/unit/client/observability/client_metrics_ut.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include @@ -21,26 +22,36 @@ class ClientMetricsTest : public ::testing::Test { } std::shared_ptr RequestCounter(const std::string& op) { - return Registry->GetCounter(Prefix + ".requests", {{"operation", op}}); + return Registry->GetCounter("db.client.operation.requests", { + {"db.system.name", "other_sql"}, + {"db.operation.name", op}, + }); } std::shared_ptr ErrorCounter(const std::string& op) { - return Registry->GetCounter(Prefix + ".errors", {{"operation", op}}); + return Registry->GetCounter("db.client.operation.errors", { + {"db.system.name", "other_sql"}, + {"db.operation.name", op}, + }); } - std::shared_ptr DurationHistogram(const std::string& op) { - return Registry->GetHistogram("db.client.operation.duration", { - {"db.system.name", "ydb"}, + std::shared_ptr DurationHistogram(const std::string& op, EStatus status) { + TLabels labels = { + {"db.system.name", "other_sql"}, {"db.operation.name", op}, - }); + {"db.response.status_code", ToString(status)}, + }; + if (status != EStatus::SUCCESS) { + labels["error.type"] = ToString(status); + } + return Registry->GetHistogram("db.client.operation.duration", labels); } - const std::string Prefix = "ydb.test"; std::shared_ptr Registry; }; TEST_F(ClientMetricsTest, RequestCounterIncrementedOnConstruction) { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); auto counter = RequestCounter("DoSomething"); ASSERT_NE(counter, nullptr); @@ -49,7 +60,7 @@ TEST_F(ClientMetricsTest, RequestCounterIncrementedOnConstruction) { TEST_F(ClientMetricsTest, SuccessDoesNotIncrementErrorCounter) { { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); metrics.End(EStatus::SUCCESS); } @@ -60,7 +71,7 @@ TEST_F(ClientMetricsTest, SuccessDoesNotIncrementErrorCounter) { TEST_F(ClientMetricsTest, FailureIncrementsErrorCounter) { { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); metrics.End(EStatus::UNAVAILABLE); } @@ -71,11 +82,11 @@ TEST_F(ClientMetricsTest, FailureIncrementsErrorCounter) { TEST_F(ClientMetricsTest, DurationRecordedOnEnd) { { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); metrics.End(EStatus::SUCCESS); } - auto hist = DurationHistogram("DoSomething"); + auto hist = DurationHistogram("DoSomething", EStatus::SUCCESS); ASSERT_NE(hist, nullptr); EXPECT_EQ(hist->Count(), 1u); EXPECT_GE(hist->GetValues()[0], 0.0); @@ -83,17 +94,17 @@ TEST_F(ClientMetricsTest, DurationRecordedOnEnd) { TEST_F(ClientMetricsTest, DurationIsInSeconds) { { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); metrics.End(EStatus::SUCCESS); } - auto hist = DurationHistogram("DoSomething"); + auto hist = DurationHistogram("DoSomething", EStatus::SUCCESS); ASSERT_NE(hist, nullptr); EXPECT_LT(hist->GetValues()[0], 1.0); } TEST_F(ClientMetricsTest, DoubleEndIsIdempotent) { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); metrics.End(EStatus::SUCCESS); metrics.End(EStatus::INTERNAL_ERROR); @@ -101,14 +112,14 @@ TEST_F(ClientMetricsTest, DoubleEndIsIdempotent) { ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto hist = DurationHistogram("DoSomething"); + auto hist = DurationHistogram("DoSomething", EStatus::SUCCESS); ASSERT_NE(hist, nullptr); EXPECT_EQ(hist->Count(), 1u); } TEST_F(ClientMetricsTest, DestructorCallsEndWithClientInternalError) { { - TClientMetrics metrics(Registry, Prefix, "DoSomething"); + TClientMetrics metrics(Registry, "DoSomething"); } auto requests = RequestCounter("DoSomething"); @@ -119,25 +130,25 @@ TEST_F(ClientMetricsTest, DestructorCallsEndWithClientInternalError) { ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 1); - auto hist = DurationHistogram("DoSomething"); + auto hist = DurationHistogram("DoSomething", EStatus::CLIENT_INTERNAL_ERROR); ASSERT_NE(hist, nullptr); EXPECT_EQ(hist->Count(), 1u); } TEST_F(ClientMetricsTest, NullRegistryDoesNotCrash) { EXPECT_NO_THROW({ - TClientMetrics metrics(nullptr, Prefix, "DoSomething"); + TClientMetrics metrics(nullptr, "DoSomething"); metrics.End(EStatus::SUCCESS); }); } TEST_F(ClientMetricsTest, DifferentOperationsHaveSeparateMetrics) { { - TClientMetrics m1(Registry, Prefix, "OpA"); + TClientMetrics m1(Registry, "OpA"); m1.End(EStatus::SUCCESS); } { - TClientMetrics m2(Registry, Prefix, "OpB"); + TClientMetrics m2(Registry, "OpB"); m2.End(EStatus::OVERLOADED); } @@ -145,19 +156,20 @@ TEST_F(ClientMetricsTest, DifferentOperationsHaveSeparateMetrics) { EXPECT_EQ(RequestCounter("OpB")->Get(), 1); EXPECT_EQ(ErrorCounter("OpA")->Get(), 0); EXPECT_EQ(ErrorCounter("OpB")->Get(), 1); - EXPECT_EQ(DurationHistogram("OpA")->Count(), 1u); - EXPECT_EQ(DurationHistogram("OpB")->Count(), 1u); + EXPECT_EQ(DurationHistogram("OpA", EStatus::SUCCESS)->Count(), 1u); + EXPECT_EQ(DurationHistogram("OpB", EStatus::OVERLOADED)->Count(), 1u); } TEST_F(ClientMetricsTest, MultipleRequestsAccumulate) { for (int i = 0; i < 5; ++i) { - TClientMetrics metrics(Registry, Prefix, "Op"); + TClientMetrics metrics(Registry, "Op"); metrics.End(i % 2 == 0 ? EStatus::SUCCESS : EStatus::TIMEOUT); } EXPECT_EQ(RequestCounter("Op")->Get(), 5); EXPECT_EQ(ErrorCounter("Op")->Get(), 2); - EXPECT_EQ(DurationHistogram("Op")->Count(), 5u); + EXPECT_EQ(DurationHistogram("Op", EStatus::SUCCESS)->Count(), 3u); + EXPECT_EQ(DurationHistogram("Op", EStatus::TIMEOUT)->Count(), 2u); } TEST_F(ClientMetricsTest, AllErrorStatusesIncrementErrorCounter) { @@ -173,7 +185,7 @@ TEST_F(ClientMetricsTest, AllErrorStatusesIncrementErrorCounter) { }; for (auto status : errorStatuses) { - TClientMetrics metrics(Registry, Prefix, "Op"); + TClientMetrics metrics(Registry, "Op"); metrics.End(status); } @@ -182,48 +194,76 @@ TEST_F(ClientMetricsTest, AllErrorStatusesIncrementErrorCounter) { EXPECT_EQ(errors->Get(), static_cast(errorStatuses.size())); } -TEST_F(ClientMetricsTest, PrefixAppliedToCounterNames) { - TClientMetrics metrics(Registry, "ydb.custom", "Op"); - metrics.End(EStatus::SUCCESS); - - EXPECT_NE(Registry->GetCounter("ydb.custom.requests", {{"operation", "Op"}}), nullptr); - EXPECT_NE(Registry->GetCounter("ydb.custom.errors", {{"operation", "Op"}}), nullptr); - - EXPECT_EQ(Registry->GetCounter("ydb.test.requests", {{"operation", "Op"}}), nullptr); -} - // --------------------------------------------------------------------------- -// TQueryMetrics prefix +// TQueryMetrics // --------------------------------------------------------------------------- -TEST(QueryMetricsTest, UsesQueryPrefix) { +TEST(QueryMetricsTest, UsesOtelStandardMetrics) { auto registry = std::make_shared(); NQuery::TQueryMetrics metrics(registry, "ExecuteQuery"); metrics.End(EStatus::SUCCESS); - EXPECT_NE(registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteQuery"}}), nullptr); - EXPECT_NE(registry->GetCounter("ydb.query.errors", {{"operation", "ExecuteQuery"}}), nullptr); - EXPECT_NE(registry->GetHistogram("db.client.operation.duration", { - {"db.system.name", "ydb"}, {"db.operation.name", "ExecuteQuery"}}), nullptr); - - EXPECT_EQ(registry->GetCounter("ydb.table.requests", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE( + registry->GetCounter( + "db.client.operation.requests", + {{"db.system.name", "other_sql"}, {"db.operation.name", "ExecuteQuery"}} + ), + nullptr + ); + EXPECT_NE( + registry->GetCounter( + "db.client.operation.errors", + {{"db.system.name", "other_sql"}, {"db.operation.name", "ExecuteQuery"}} + ), + nullptr + ); + EXPECT_NE( + registry->GetHistogram( + "db.client.operation.duration", + { + {"db.system.name", "other_sql"}, + {"db.operation.name", "ExecuteQuery"}, + {"db.response.status_code", ToString(EStatus::SUCCESS)}, + } + ), + nullptr + ); } // --------------------------------------------------------------------------- -// TTableMetrics prefix +// TTableMetrics // --------------------------------------------------------------------------- -TEST(TableMetricsTest, UsesTablePrefix) { +TEST(TableMetricsTest, UsesOtelStandardMetrics) { auto registry = std::make_shared(); NTable::TTableMetrics metrics(registry, "ExecuteDataQuery"); metrics.End(EStatus::SUCCESS); - EXPECT_NE(registry->GetCounter("ydb.table.requests", {{"operation", "ExecuteDataQuery"}}), nullptr); - EXPECT_NE(registry->GetCounter("ydb.table.errors", {{"operation", "ExecuteDataQuery"}}), nullptr); - EXPECT_NE(registry->GetHistogram("db.client.operation.duration", { - {"db.system.name", "ydb"}, {"db.operation.name", "ExecuteDataQuery"}}), nullptr); - - EXPECT_EQ(registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteDataQuery"}}), nullptr); + EXPECT_NE( + registry->GetCounter( + "db.client.operation.requests", + {{"db.system.name", "other_sql"}, {"db.operation.name", "ExecuteDataQuery"}} + ), + nullptr + ); + EXPECT_NE( + registry->GetCounter( + "db.client.operation.errors", + {{"db.system.name", "other_sql"}, {"db.operation.name", "ExecuteDataQuery"}} + ), + nullptr + ); + EXPECT_NE( + registry->GetHistogram( + "db.client.operation.duration", + { + {"db.system.name", "other_sql"}, + {"db.operation.name", "ExecuteDataQuery"}, + {"db.response.status_code", ToString(EStatus::SUCCESS)}, + } + ), + nullptr + ); } From 51757d69e3a94f537208e84ad542077a181247ca Mon Sep 17 00:00:00 2001 From: maladetska Date: Wed, 18 Mar 2026 01:47:28 +0800 Subject: [PATCH 09/93] add tracing implementation and its demo launch --- examples/CMakeLists.txt | 4 + examples/otel_tracing/CMakeLists.txt | 40 +++ examples/otel_tracing/README.md | 108 ++++++++ examples/otel_tracing/docker-compose.yml | 69 +++++ .../grafana/dashboards/ydb-query-service.json | 129 ++++++++++ .../provisioning/dashboards/dashboards.yml | 12 + .../provisioning/datasources/datasources.yml | 16 ++ examples/otel_tracing/main.cpp | 232 +++++++++++++++++ .../otel_tracing/otel-collector/config.yml | 32 +++ .../otel_tracing/prometheus/prometheus.yml | 8 + .../ydb-cpp-sdk/open_telemetry/metrics.h | 5 +- plugins/metrics/otel/src/metrics.cpp | 62 +++-- .../ydb-cpp-sdk/open_telemetry/trace.h | 5 +- plugins/trace/otel/src/trace.cpp | 36 +-- src/client/query/client.cpp | 1 - src/client/query/impl/query_spans.cpp | 11 - src/client/query/impl/query_spans.h | 1 - tests/common/fake_trace_provider.h | 141 ++++++++++ tests/unit/client/CMakeLists.txt | 13 + tests/unit/client/query/query_spans_ut.cpp | 242 ++++++++++++++++++ 20 files changed, 1101 insertions(+), 66 deletions(-) create mode 100644 examples/otel_tracing/CMakeLists.txt create mode 100644 examples/otel_tracing/README.md create mode 100644 examples/otel_tracing/docker-compose.yml create mode 100644 examples/otel_tracing/grafana/dashboards/ydb-query-service.json create mode 100644 examples/otel_tracing/grafana/provisioning/dashboards/dashboards.yml create mode 100644 examples/otel_tracing/grafana/provisioning/datasources/datasources.yml create mode 100644 examples/otel_tracing/main.cpp create mode 100644 examples/otel_tracing/otel-collector/config.yml create mode 100644 examples/otel_tracing/prometheus/prometheus.yml create mode 100644 tests/common/fake_trace_provider.h create mode 100644 tests/unit/client/query/query_spans_ut.cpp diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 9607701b4f6..380b49e8e99 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -9,3 +9,7 @@ add_subdirectory(topic_writer/transaction) add_subdirectory(ttl) add_subdirectory(vector_index) add_subdirectory(vector_index_builtin) + +if (YDB_SDK_ENABLE_OTEL_TRACE AND YDB_SDK_ENABLE_OTEL_METRICS) + add_subdirectory(otel_tracing) +endif() diff --git a/examples/otel_tracing/CMakeLists.txt b/examples/otel_tracing/CMakeLists.txt new file mode 100644 index 00000000000..e98ad9d3751 --- /dev/null +++ b/examples/otel_tracing/CMakeLists.txt @@ -0,0 +1,40 @@ +add_executable(otel_tracing_example) + +target_link_libraries(otel_tracing_example PUBLIC + yutil + getopt + YDB-CPP-SDK::Query + YDB-CPP-SDK::Params + YDB-CPP-SDK::Driver + YDB-CPP-SDK::OpenTelemetryTrace + YDB-CPP-SDK::OpenTelemetryMetrics + opentelemetry-cpp::otlp_http_exporter + opentelemetry-cpp::otlp_http_metric_exporter +) + +target_sources(otel_tracing_example PRIVATE + main.cpp +) + +vcs_info(otel_tracing_example) + +if (CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64" OR CMAKE_SYSTEM_PROCESSOR STREQUAL "AMD64") + target_link_libraries(otel_tracing_example PUBLIC + cpuid_check + ) +endif() + +if (CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_link_options(otel_tracing_example PRIVATE + -ldl + -lrt + -Wl,--no-as-needed + -lpthread + ) +elseif (CMAKE_SYSTEM_NAME STREQUAL "Darwin") + target_link_options(otel_tracing_example PRIVATE + -Wl,-platform_version,macos,11.0,11.0 + -framework + CoreFoundation + ) +endif() diff --git a/examples/otel_tracing/README.md b/examples/otel_tracing/README.md new file mode 100644 index 00000000000..412ea575ad8 --- /dev/null +++ b/examples/otel_tracing/README.md @@ -0,0 +1,108 @@ +# YDB C++ SDK — OpenTelemetry Tracing Demo + +Демонстрация трассировки операций QueryService (CreateSession, ExecuteQuery, Commit, Rollback) +с визуализацией в **Grafana**, **Jaeger** и **Prometheus**. + +## Архитектура + +``` +┌──────────────┐ OTLP/HTTP ┌──────────────────┐ +│ C++ demo │ ──────────────────> │ OTel Collector │ +│ application │ │ :4328 (HTTP) │ +└──────────────┘ └────────┬──────────┘ + │ │ + traces │ │ metrics + ▼ ▼ + ┌──────────┐ ┌────────────┐ + │ Jaeger │ │ Prometheus │ + │ :16686 │ │ :9090 │ + └─────┬─────┘ └──────┬──────┘ + │ │ + └───────┬───────┘ + ▼ + ┌──────────┐ + │ Grafana │ + │ :3000 │ + └──────────┘ +``` + +## Быстрый старт + +### 1. Запустить инфраструктуру + +```bash +cd examples/otel_tracing +docker compose up -d +``` + +Дождитесь готовности YDB: + +```bash +docker compose logs ydb -f +# Ждите строку "Database started successfully" +``` + +### 2. Собрать SDK с OTel и тестами + +Из корня репозитория: + +```bash +mkdir -p build && cd build + +cmake .. \ + -DYDB_SDK_TESTS=ON \ + -DYDB_SDK_ENABLE_OTEL_TRACE=ON \ + -DYDB_SDK_ENABLE_OTEL_METRICS=ON + +cmake --build . --target otel_tracing_example -j$(nproc) +``` + +### 3. Запустить демо + +```bash +./examples/otel_tracing/otel_tracing_example \ + --endpoint grpc://localhost:2136 \ + --database /local \ + --otlp http://localhost:4328 \ + --iterations 20 +``` + +### 4. Открыть дашборды + +| Сервис | URL | Описание | +|-----------|------------------------------|---------------------------------| +| Grafana | http://localhost:3000 | Дашборд "YDB QueryService" | +| Jaeger | http://localhost:16686 | Поиск трейсов по сервису | +| Prometheus| http://localhost:9090 | Метрики `ydb_ydb_query_*` | + +**Grafana**: логин `admin` / пароль `admin`. + +### 5. Что смотреть + +#### В Grafana (дашборд "YDB QueryService"): +- **Request Rate by Operation** — RPS по операциям (ExecuteQuery, CreateSession, Commit, Rollback) +- **Error Rate by Operation** — частота ошибок +- **Latency p50/p95/p99** — распределение задержек +- **Error Ratio** — процент ошибок +- **Recent Traces** — таблица трейсов из Jaeger + +#### В Jaeger UI: +- Выберите сервис `ydb-cpp-sdk-demo` +- Каждый спан содержит атрибуты: + - `db.system.name` = `ydb` + - `server.address`, `server.port` + - `network.peer.address`, `network.peer.port` + - `db.query.text` (для ExecuteQuery) + - `db.response.status_code`, `error.type` (при ошибках) + +#### В Prometheus: +- `ydb_ydb_query_requests_total` — счётчик запросов +- `ydb_ydb_query_errors_total` — счётчик ошибок +- `ydb_ydb_query_latency_ms_bucket` — гистограмма задержек + +### 6. Остановить + +```bash +cd examples/otel_tracing +docker compose down -v +``` diff --git a/examples/otel_tracing/docker-compose.yml b/examples/otel_tracing/docker-compose.yml new file mode 100644 index 00000000000..3314c27d421 --- /dev/null +++ b/examples/otel_tracing/docker-compose.yml @@ -0,0 +1,69 @@ +services: + ydb: + image: cr.yandex/yc/yandex-docker-local-ydb:latest + platform: linux/amd64 + ports: + - "2136:2136" + - "8765:8765" + environment: + - GRPC_TLS_PORT=2135 + - GRPC_PORT=2136 + - MON_PORT=8765 + - YDB_DEFAULT_LOG_LEVEL=NOTICE + volumes: + - ydb-data:/ydb_data + healthcheck: + test: /bin/sh -c "/ydb -e grpc://localhost:2136 -d /local scheme ls" + interval: 5s + timeout: 5s + retries: 20 + + jaeger: + image: jaegertracing/all-in-one:1.76.0 + ports: + - "16686:16686" + - "4317:4317" + - "4318:4318" + environment: + - COLLECTOR_OTLP_ENABLED=true + + prometheus: + image: prom/prometheus:v2.53.0 + ports: + - "9090:9090" + volumes: + - ./prometheus/prometheus.yml:/etc/prometheus/prometheus.yml:ro + depends_on: + - otel-collector + + otel-collector: + image: otel/opentelemetry-collector-contrib:latest + ports: + - "4327:4317" + - "4328:4318" + - "8889:8889" + volumes: + - ./otel-collector/config.yml:/etc/otelcol-contrib/config.yaml:ro + depends_on: + - jaeger + + grafana: + image: grafana/grafana:latest + ports: + - "3000:3000" + environment: + - GF_SECURITY_ADMIN_USER=admin + - GF_SECURITY_ADMIN_PASSWORD=admin + - GF_AUTH_ANONYMOUS_ENABLED=true + - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin + volumes: + - ./grafana/provisioning:/etc/grafana/provisioning:ro + - ./grafana/dashboards:/var/lib/grafana/dashboards:ro + - grafana-data:/var/lib/grafana + depends_on: + - jaeger + - prometheus + +volumes: + ydb-data: + grafana-data: diff --git a/examples/otel_tracing/grafana/dashboards/ydb-query-service.json b/examples/otel_tracing/grafana/dashboards/ydb-query-service.json new file mode 100644 index 00000000000..1ce6c11a8f7 --- /dev/null +++ b/examples/otel_tracing/grafana/dashboards/ydb-query-service.json @@ -0,0 +1,129 @@ +{ + "annotations": { "list": [] }, + "editable": true, + "fiscalYearStartMonth": 0, + "graphTooltip": 1, + "links": [], + "panels": [ + { + "title": "Request Rate by Operation", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 0 }, + "datasource": { "type": "prometheus", "uid": "${prometheus_ds}" }, + "fieldConfig": { + "defaults": { + "unit": "ops", + "custom": { "drawStyle": "line", "fillOpacity": 10 } + } + }, + "targets": [ + { + "expr": "rate(ydb_query_requests_total[1m])", + "legendFormat": "{{operation}}" + } + ] + }, + { + "title": "Error Rate by Operation", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 0 }, + "datasource": { "type": "prometheus", "uid": "${prometheus_ds}" }, + "fieldConfig": { + "defaults": { + "unit": "ops", + "custom": { "drawStyle": "line", "fillOpacity": 10 }, + "color": { "mode": "palette-classic" } + } + }, + "targets": [ + { + "expr": "rate(ydb_query_errors_total[1m])", + "legendFormat": "{{operation}}" + } + ] + }, + { + "title": "Latency p50 / p95 / p99 (ms)", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 8 }, + "datasource": { "type": "prometheus", "uid": "${prometheus_ds}" }, + "fieldConfig": { + "defaults": { + "unit": "ms", + "custom": { "drawStyle": "line", "fillOpacity": 5 } + } + }, + "targets": [ + { + "expr": "histogram_quantile(0.50, rate(ydb_query_latency_ms_bucket[1m]))", + "legendFormat": "p50 {{operation}}" + }, + { + "expr": "histogram_quantile(0.95, rate(ydb_query_latency_ms_bucket[1m]))", + "legendFormat": "p95 {{operation}}" + }, + { + "expr": "histogram_quantile(0.99, rate(ydb_query_latency_ms_bucket[1m]))", + "legendFormat": "p99 {{operation}}" + } + ] + }, + { + "title": "Error Ratio (%)", + "type": "stat", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 8 }, + "datasource": { "type": "prometheus", "uid": "${prometheus_ds}" }, + "fieldConfig": { + "defaults": { + "unit": "percentunit", + "thresholds": { + "steps": [ + { "color": "green", "value": null }, + { "color": "yellow", "value": 0.01 }, + { "color": "red", "value": 0.05 } + ] + } + } + }, + "targets": [ + { + "expr": "sum(rate(ydb_query_errors_total[5m])) by (operation) / sum(rate(ydb_query_requests_total[5m])) by (operation)", + "legendFormat": "{{operation}}" + } + ] + }, + { + "title": "Recent Traces", + "type": "table", + "gridPos": { "h": 10, "w": 24, "x": 0, "y": 16 }, + "datasource": { "type": "jaeger", "uid": "${jaeger_ds}" }, + "targets": [ + { + "query": "ydb-cpp-sdk-demo", + "queryType": "search", + "service": "ydb-cpp-sdk-demo" + } + ] + } + ], + "schemaVersion": 39, + "templating": { + "list": [ + { + "name": "prometheus_ds", + "type": "datasource", + "query": "prometheus", + "current": { "text": "Prometheus", "value": "Prometheus" } + }, + { + "name": "jaeger_ds", + "type": "datasource", + "query": "jaeger", + "current": { "text": "Jaeger", "value": "Jaeger" } + } + ] + }, + "time": { "from": "now-30m", "to": "now" }, + "title": "YDB QueryService", + "uid": "ydb-query-service" +} diff --git a/examples/otel_tracing/grafana/provisioning/dashboards/dashboards.yml b/examples/otel_tracing/grafana/provisioning/dashboards/dashboards.yml new file mode 100644 index 00000000000..8336756c137 --- /dev/null +++ b/examples/otel_tracing/grafana/provisioning/dashboards/dashboards.yml @@ -0,0 +1,12 @@ +apiVersion: 1 + +providers: + - name: "YDB" + orgId: 1 + folder: "YDB" + type: file + disableDeletion: false + editable: true + options: + path: /var/lib/grafana/dashboards + foldersFromFilesStructure: false diff --git a/examples/otel_tracing/grafana/provisioning/datasources/datasources.yml b/examples/otel_tracing/grafana/provisioning/datasources/datasources.yml new file mode 100644 index 00000000000..428e06210ab --- /dev/null +++ b/examples/otel_tracing/grafana/provisioning/datasources/datasources.yml @@ -0,0 +1,16 @@ +apiVersion: 1 + +datasources: + - name: Jaeger + type: jaeger + access: proxy + url: http://jaeger:16686 + isDefault: false + editable: true + + - name: Prometheus + type: prometheus + access: proxy + url: http://prometheus:9090 + isDefault: true + editable: true diff --git a/examples/otel_tracing/main.cpp b/examples/otel_tracing/main.cpp new file mode 100644 index 00000000000..f8e53e23669 --- /dev/null +++ b/examples/otel_tracing/main.cpp @@ -0,0 +1,232 @@ +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace nostd = opentelemetry::nostd; +namespace sdktrace = opentelemetry::sdk::trace; +namespace sdkmetrics = opentelemetry::sdk::metrics; +namespace otlp = opentelemetry::exporter::otlp; +namespace resource = opentelemetry::sdk::resource; + +using namespace NYdb; +using namespace NYdb::NQuery; +using namespace NYdb::NStatusHelpers; + +struct TConfig { + std::string Endpoint = "grpc://localhost:2136"; + std::string Database = "/local"; + std::string OtlpEndpoint = "http://localhost:4328"; + int Iterations = 20; +}; + +nostd::shared_ptr InitTracing(const TConfig& cfg) { + otlp::OtlpHttpExporterOptions opts; + opts.url = cfg.OtlpEndpoint + "/v1/traces"; + + auto exporter = otlp::OtlpHttpExporterFactory::Create(opts); + auto processor = sdktrace::SimpleSpanProcessorFactory::Create(std::move(exporter)); + + auto res = resource::Resource::Create({ + {"service.name", "ydb-cpp-sdk-demo"}, + {"service.version", "1.0.0"}, + }); + + std::shared_ptr provider = + std::make_shared(std::move(processor), res); + return nostd::shared_ptr(provider); +} + +nostd::shared_ptr InitMetrics(const TConfig& cfg) { + otlp::OtlpHttpMetricExporterOptions opts; + opts.url = cfg.OtlpEndpoint + "/v1/metrics"; + + auto exporter = otlp::OtlpHttpMetricExporterFactory::Create(opts); + + sdkmetrics::PeriodicExportingMetricReaderOptions readerOpts; + readerOpts.export_interval_millis = std::chrono::milliseconds(5000); + readerOpts.export_timeout_millis = std::chrono::milliseconds(3000); + + auto reader = sdkmetrics::PeriodicExportingMetricReaderFactory::Create(std::move(exporter), readerOpts); + + auto res = resource::Resource::Create({ + {"service.name", "ydb-cpp-sdk-demo"}, + {"service.version", "1.0.0"}, + }); + + auto rawProvider = std::make_shared( + std::unique_ptr(new sdkmetrics::ViewRegistry()), res); + rawProvider->AddMetricReader(std::move(reader)); + + std::shared_ptr provider = rawProvider; + return nostd::shared_ptr(provider); +} + +void RunWorkload(TQueryClient& client, int iterations) { + std::cout << "=== Creating table ===" << std::endl; + + ThrowOnError(client.RetryQuerySync([](TSession session) { + return session.ExecuteQuery(R"( + CREATE TABLE IF NOT EXISTS otel_demo ( + id Uint64, + value Utf8, + PRIMARY KEY (id) + ) + )", TTxControl::NoTx()).GetValueSync(); + })); + + for (int i = 0; i < iterations; ++i) { + std::cout << "--- Iteration " << (i + 1) << "/" << iterations << " ---" << std::endl; + + ThrowOnError(client.RetryQuerySync([i](TSession session) { + auto params = TParamsBuilder() + .AddParam("$id").Uint64(i).Build() + .AddParam("$val").Utf8("item_" + std::to_string(i)).Build() + .Build(); + + return session.ExecuteQuery(R"( + DECLARE $id AS Uint64; + DECLARE $val AS Utf8; + UPSERT INTO otel_demo (id, value) VALUES ($id, $val) + )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + params).GetValueSync(); + })); + + ThrowOnError(client.RetryQuerySync([i](TSession session) { + auto params = TParamsBuilder() + .AddParam("$id").Uint64(i).Build() + .Build(); + + return session.ExecuteQuery(R"( + DECLARE $id AS Uint64; + SELECT id, value FROM otel_demo WHERE id = $id + )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + params).GetValueSync(); + })); + + ThrowOnError(client.RetryQuerySync([](TQueryClient client) -> TStatus { + auto session = client.GetSession().GetValueSync().GetSession(); + auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).GetValueSync(); + if (!beginResult.IsSuccess()) { + return beginResult; + } + auto tx = beginResult.GetTransaction(); + + auto result = session.ExecuteQuery(R"( + SELECT COUNT(*) AS cnt FROM otel_demo + )", TTxControl::Tx(tx)).GetValueSync(); + + if (!result.IsSuccess()) { + return result; + } + + return tx.Commit().GetValueSync(); + })); + + if (i % 5 == 4) { + auto rollbackResult = client.RetryQuerySync([](TQueryClient client) -> TStatus { + auto session = client.GetSession().GetValueSync().GetSession(); + auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).GetValueSync(); + if (!beginResult.IsSuccess()) { + return beginResult; + } + auto tx = beginResult.GetTransaction(); + return tx.Rollback().GetValueSync(); + }); + if (!rollbackResult.IsSuccess()) { + std::cerr << " Rollback status: " << static_cast(rollbackResult.GetStatus()) << std::endl; + } + } + + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + + std::cout << "=== Dropping table ===" << std::endl; + + ThrowOnError(client.RetryQuerySync([](TSession session) { + return session.ExecuteQuery( + "DROP TABLE otel_demo", TTxControl::NoTx()).GetValueSync(); + })); +} + +int main(int argc, char** argv) { + TConfig cfg; + + NLastGetopt::TOpts opts; + opts.AddLongOption('e', "endpoint", "YDB endpoint") + .DefaultValue(cfg.Endpoint).StoreResult(&cfg.Endpoint); + opts.AddLongOption('d', "database", "YDB database") + .DefaultValue(cfg.Database).StoreResult(&cfg.Database); + opts.AddLongOption("otlp", "OTLP HTTP endpoint") + .DefaultValue(cfg.OtlpEndpoint).StoreResult(&cfg.OtlpEndpoint); + opts.AddLongOption('n', "iterations", "Number of iterations") + .DefaultValue(std::to_string(cfg.Iterations)).StoreResult(&cfg.Iterations); + + NLastGetopt::TOptsParseResult parsedOpts(&opts, argc, argv); + + std::cout << "Initializing OpenTelemetry..." << std::endl; + std::cout << " OTLP endpoint: " << cfg.OtlpEndpoint << std::endl; + + auto tracerProvider = InitTracing(cfg); + auto meterProvider = InitMetrics(cfg); + + auto ydbTraceProvider = NMetrics::CreateOtelTraceProvider(tracerProvider); + auto ydbMetricRegistry = NMetrics::CreateOtelMetricRegistry(meterProvider); + + std::cout << "Connecting to YDB at " << cfg.Endpoint << cfg.Database << std::endl; + + auto driverConfig = TDriverConfig() + .SetEndpoint(cfg.Endpoint) + .SetDatabase(cfg.Database) + .SetDiscoveryMode(EDiscoveryMode::Off) + .SetTraceProvider(ydbTraceProvider) + .SetMetricRegistry(ydbMetricRegistry); + + TDriver driver(driverConfig); + TQueryClient client(driver); + + try { + RunWorkload(client, cfg.Iterations); + } catch (const std::exception& e) { + std::cerr << "Error: " << e.what() << std::endl; + } + + std::cout << "Flushing telemetry..." << std::endl; + + driver.Stop(true); + + if (auto* sdkTracerProvider = dynamic_cast(tracerProvider.get())) { + sdkTracerProvider->ForceFlush(); + } + if (auto* sdkMeterProvider = dynamic_cast(meterProvider.get())) { + sdkMeterProvider->ForceFlush(); + } + + std::this_thread::sleep_for(std::chrono::seconds(3)); + + std::cout << "Done. Open Grafana at http://localhost:3000" << std::endl; + std::cout << " Dashboard: YDB QueryService" << std::endl; + std::cout << " Also: Jaeger UI at http://localhost:16686" << std::endl; + + return 0; +} diff --git a/examples/otel_tracing/otel-collector/config.yml b/examples/otel_tracing/otel-collector/config.yml new file mode 100644 index 00000000000..9589c9cd4ee --- /dev/null +++ b/examples/otel_tracing/otel-collector/config.yml @@ -0,0 +1,32 @@ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:4317 + http: + endpoint: 0.0.0.0:4318 + +exporters: + otlp/jaeger: + endpoint: jaeger:4317 + tls: + insecure: true + + prometheus: + endpoint: 0.0.0.0:8889 + +processors: + batch: + timeout: 1s + send_batch_size: 1024 + +service: + pipelines: + traces: + receivers: [otlp] + processors: [batch] + exporters: [otlp/jaeger] + metrics: + receivers: [otlp] + processors: [batch] + exporters: [prometheus] diff --git a/examples/otel_tracing/prometheus/prometheus.yml b/examples/otel_tracing/prometheus/prometheus.yml new file mode 100644 index 00000000000..faeda702aff --- /dev/null +++ b/examples/otel_tracing/prometheus/prometheus.yml @@ -0,0 +1,8 @@ +global: + scrape_interval: 5s + evaluation_interval: 5s + +scrape_configs: + - job_name: "otel-collector" + static_configs: + - targets: ["otel-collector:8889"] diff --git a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h index 054d040f97d..f992c577bf6 100644 --- a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h +++ b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -3,10 +3,7 @@ #include #include - -namespace opentelemetry::metrics { -class MeterProvider; -} +#include namespace NYdb::inline V3::NMetrics { diff --git a/plugins/metrics/otel/src/metrics.cpp b/plugins/metrics/otel/src/metrics.cpp index af07af89bb8..f883fae7946 100644 --- a/plugins/metrics/otel/src/metrics.cpp +++ b/plugins/metrics/otel/src/metrics.cpp @@ -9,74 +9,80 @@ #include #include +#include + namespace NYdb::inline V3::NMetrics { namespace { -using namespace opentelemetry; +namespace otel_metrics = opentelemetry::metrics; +namespace otel_nostd = opentelemetry::nostd; +namespace otel_common = opentelemetry::common; +namespace otel_context = opentelemetry::context; +namespace otel_sdk_metrics = opentelemetry::sdk::metrics; -common::KeyValueIterableView MakeAttributes(const TLabels& labels) { - return common::KeyValueIterableView(labels); +otel_common::KeyValueIterableView MakeAttributes(const TLabels& labels) { + return otel_common::KeyValueIterableView(labels); } class TOtelCounter : public ICounter { public: - TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) + TOtelCounter(otel_nostd::shared_ptr> counter, const TLabels& labels) : Counter_(std::move(counter)) , Labels_(labels) {} void Inc() override { - Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Counter_->Add(1, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } private: - nostd::shared_ptr> Counter_; + otel_nostd::shared_ptr> Counter_; TLabels Labels_; }; class TOtelUpDownCounterGauge : public IGauge { public: - TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) + TOtelUpDownCounterGauge(otel_nostd::shared_ptr> counter, const TLabels& labels) : Counter_(std::move(counter)) , Labels_(labels) {} void Add(double delta) override { - Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Counter_->Add(delta, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); Value_ += delta; } void Set(double value) override { - Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Counter_->Add(value - Value_, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); Value_ = value; } private: - nostd::shared_ptr> Counter_; + otel_nostd::shared_ptr> Counter_; TLabels Labels_; double Value_ = 0; }; class TOtelHistogram : public IHistogram { public: - TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) + TOtelHistogram(otel_nostd::shared_ptr> histogram, const TLabels& labels) : Histogram_(std::move(histogram)) , Labels_(labels) {} void Record(double value) override { - Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Histogram_->Record(value, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } private: - nostd::shared_ptr> Histogram_; + otel_nostd::shared_ptr> Histogram_; TLabels Labels_; }; class TOtelMetricRegistry : public IMetricRegistry { public: - TOtelMetricRegistry(nostd::shared_ptr meterProvider) + TOtelMetricRegistry(otel_nostd::shared_ptr meterProvider) : MeterProvider_(std::move(meterProvider)) , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) {} @@ -116,7 +122,7 @@ class TOtelMetricRegistry : public IMetricRegistry { return; } - auto* sdkProvider = dynamic_cast(MeterProvider_.get()); + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); if (!sdkProvider) { return; } @@ -128,32 +134,32 @@ class TOtelMetricRegistry : public IMetricRegistry { } } - auto selector = std::make_unique( - sdk::metrics::InstrumentType::kHistogram, + auto selector = std::make_unique( + otel_sdk_metrics::InstrumentType::kHistogram, name, unit ); - auto meterSelector = std::make_unique( - "ydb-cpp-sdk", - GetSdkSemver(), - {} + auto meterSelector = std::make_unique( + std::string("ydb-cpp-sdk"), + std::string(GetSdkSemver()), + std::string() ); - auto histogramConfig = std::make_shared(); + auto histogramConfig = std::make_shared(); histogramConfig->boundaries_ = buckets; - auto view = std::make_unique( - {}, - {}, - sdk::metrics::AggregationType::kHistogram, + auto view = std::make_unique( + std::string(), + std::string(), + otel_sdk_metrics::AggregationType::kHistogram, histogramConfig ); sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); } - nostd::shared_ptr MeterProvider_; - nostd::shared_ptr Meter_; + otel_nostd::shared_ptr MeterProvider_; + otel_nostd::shared_ptr Meter_; std::mutex HistogramViewsLock_; std::unordered_set HistogramViews_; }; diff --git a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h index 9bdc12fb25f..68b238d6a41 100644 --- a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h +++ b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -3,10 +3,7 @@ #include #include - -namespace opentelemetry::trace { -class TracerProvider; -} +#include namespace NYdb::inline V3::NMetrics { diff --git a/plugins/trace/otel/src/trace.cpp b/plugins/trace/otel/src/trace.cpp index 7cac3f4c1cb..b904df3ab9d 100644 --- a/plugins/trace/otel/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -8,22 +8,24 @@ namespace NYdb::inline V3::NMetrics { namespace { -using namespace opentelemetry; +namespace otel_trace = opentelemetry::trace; +namespace otel_nostd = opentelemetry::nostd; +namespace otel_common = opentelemetry::common; -trace::SpanKind MapSpanKind(ESpanKind kind) { +otel_trace::SpanKind MapSpanKind(ESpanKind kind) { switch (kind) { - case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; - case ESpanKind::SERVER: return trace::SpanKind::kServer; - case ESpanKind::CLIENT: return trace::SpanKind::kClient; - case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; - case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + case ESpanKind::INTERNAL: return otel_trace::SpanKind::kInternal; + case ESpanKind::SERVER: return otel_trace::SpanKind::kServer; + case ESpanKind::CLIENT: return otel_trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return otel_trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return otel_trace::SpanKind::kConsumer; } - return trace::SpanKind::kInternal; + return otel_trace::SpanKind::kInternal; } class TOtelSpan : public ISpan { public: - TOtelSpan(nostd::shared_ptr span) + TOtelSpan(otel_nostd::shared_ptr span) : Span_(std::move(span)) {} @@ -43,38 +45,38 @@ class TOtelSpan : public ISpan { if (attributes.empty()) { Span_->AddEvent(name); } else { - std::vector> attrs; + std::vector> attrs; attrs.reserve(attributes.size()); for (const auto& [k, v] : attributes) { - attrs.emplace_back(nostd::string_view(k), common::AttributeValue(nostd::string_view(v))); + attrs.emplace_back(otel_nostd::string_view(k), otel_common::AttributeValue(otel_nostd::string_view(v))); } Span_->AddEvent(name, attrs); } } private: - nostd::shared_ptr Span_; + otel_nostd::shared_ptr Span_; }; class TOtelTracer : public ITracer { public: - TOtelTracer(nostd::shared_ptr tracer) + TOtelTracer(otel_nostd::shared_ptr tracer) : Tracer_(std::move(tracer)) {} std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { - trace::StartSpanOptions options; + otel_trace::StartSpanOptions options; options.kind = MapSpanKind(kind); return std::make_shared(Tracer_->StartSpan(name, options)); } private: - nostd::shared_ptr Tracer_; + otel_nostd::shared_ptr Tracer_; }; class TOtelTraceProvider : public ITraceProvider { public: - TOtelTraceProvider(nostd::shared_ptr tracerProvider) + TOtelTraceProvider(otel_nostd::shared_ptr tracerProvider) : TracerProvider_(std::move(tracerProvider)) {} @@ -83,7 +85,7 @@ class TOtelTraceProvider : public ITraceProvider { } private: - nostd::shared_ptr TracerProvider_; + otel_nostd::shared_ptr TracerProvider_; }; } // namespace diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 0fc3f75ec88..2e26c95f8d1 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -104,7 +104,6 @@ class TQueryClient::TImpl: public TClientImplCommon, public CollectParamsSize(params ? ¶ms->GetProtoMap() : nullptr); auto span = std::make_shared(Tracer_, "ExecuteQuery", DbDriverState_->DiscoveryEndpoint); - span->SetQueryText(query); auto metrics = std::make_shared(MetricRegistry_, "ExecuteQuery"); return TExecQueryImpl::ExecuteQuery( diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index df4a5718771..4532b0e6552 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -104,17 +104,6 @@ void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { } } -void TQuerySpan::SetQueryText(const std::string& query) noexcept { - if (!Span_ || query.empty()) { - return; - } - try { - Span_->SetAttribute("db.query.text", query); - } catch (...) { - SafeLogSpanError("failed to set query text"); - } -} - void TQuerySpan::AddEvent(const std::string& name, const std::map& attributes) noexcept { if (!Span_) { return; diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index 75fd0fa830e..096634346b3 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -16,7 +16,6 @@ class TQuerySpan { ~TQuerySpan() noexcept; void SetPeerEndpoint(const std::string& endpoint) noexcept; - void SetQueryText(const std::string& query) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; void End(EStatus status) noexcept; diff --git a/tests/common/fake_trace_provider.h b/tests/common/fake_trace_provider.h new file mode 100644 index 00000000000..79d38337ab5 --- /dev/null +++ b/tests/common/fake_trace_provider.h @@ -0,0 +1,141 @@ +#pragma once + +#include + +#include +#include + +namespace NYdb::NTests { + +struct TFakeEvent { + std::string Name; + std::map Attributes; +}; + +class TFakeSpan : public NMetrics::ISpan { +public: + void End() override { + std::lock_guard lock(Mutex_); + Ended_ = true; + } + + void SetAttribute(const std::string& key, const std::string& value) override { + std::lock_guard lock(Mutex_); + StringAttributes_[key] = value; + } + + void SetAttribute(const std::string& key, int64_t value) override { + std::lock_guard lock(Mutex_); + IntAttributes_[key] = value; + } + + void AddEvent(const std::string& name, const std::map& attributes) override { + std::lock_guard lock(Mutex_); + Events_.push_back({name, attributes}); + } + + bool IsEnded() const { + std::lock_guard lock(Mutex_); + return Ended_; + } + + std::string GetStringAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + auto it = StringAttributes_.find(key); + return it != StringAttributes_.end() ? it->second : ""; + } + + bool HasStringAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + return StringAttributes_.contains(key); + } + + int64_t GetIntAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + auto it = IntAttributes_.find(key); + return it != IntAttributes_.end() ? it->second : 0; + } + + bool HasIntAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + return IntAttributes_.contains(key); + } + + std::vector GetEvents() const { + std::lock_guard lock(Mutex_); + return Events_; + } + +private: + mutable std::mutex Mutex_; + bool Ended_ = false; + std::map StringAttributes_; + std::map IntAttributes_; + std::vector Events_; +}; + +class TFakeTracer : public NMetrics::ITracer { +public: + std::shared_ptr StartSpan(const std::string& name, NMetrics::ESpanKind kind) override { + auto span = std::make_shared(); + std::lock_guard lock(Mutex_); + Spans_.push_back({name, kind, span}); + return span; + } + + struct TSpanRecord { + std::string Name; + NMetrics::ESpanKind Kind; + std::shared_ptr Span; + }; + + std::vector GetSpans() const { + std::lock_guard lock(Mutex_); + return Spans_; + } + + std::shared_ptr GetLastSpan() const { + std::lock_guard lock(Mutex_); + return Spans_.empty() ? nullptr : Spans_.back().Span; + } + + TSpanRecord GetLastSpanRecord() const { + std::lock_guard lock(Mutex_); + return Spans_.back(); + } + + size_t SpanCount() const { + std::lock_guard lock(Mutex_); + return Spans_.size(); + } + +private: + mutable std::mutex Mutex_; + std::vector Spans_; +}; + +class TFakeTraceProvider : public NMetrics::ITraceProvider { +public: + std::shared_ptr GetTracer(const std::string& name) override { + std::lock_guard lock(Mutex_); + auto it = Tracers_.find(name); + if (it != Tracers_.end()) { + return it->second; + } + auto tracer = std::make_shared(); + Tracers_[name] = tracer; + return tracer; + } + + std::shared_ptr GetFakeTracer(const std::string& name) const { + std::lock_guard lock(Mutex_); + auto it = Tracers_.find(name); + return it != Tracers_.end() ? it->second : nullptr; + } + +private: + mutable std::mutex Mutex_; + std::map> Tracers_; +}; + +} // namespace NYdb::NTests diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index de86c3fe274..ee15a8cd2e8 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -115,3 +115,16 @@ add_ydb_test(NAME client-ydb_metrics_ut GTEST LABELS unit ) + +add_ydb_test(NAME client-ydb_query_spans_ut GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + query/query_spans_ut.cpp + LINK_LIBRARIES + yutil + client-ydb_query-impl + client-trace + LABELS + unit +) diff --git a/tests/unit/client/query/query_spans_ut.cpp b/tests/unit/client/query/query_spans_ut.cpp new file mode 100644 index 00000000000..ae553909413 --- /dev/null +++ b/tests/unit/client/query/query_spans_ut.cpp @@ -0,0 +1,242 @@ +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NQuery; +using namespace NYdb::NMetrics; +using namespace NYdb::NTests; + +class QuerySpanTest : public ::testing::Test { +protected: + void SetUp() override { + Tracer = std::make_shared(); + } + + std::shared_ptr Tracer; +}; + +TEST_F(QuerySpanTest, SpanNameFormat) { + TQuerySpan span(Tracer, "ExecuteQuery", "localhost:2135"); + span.End(EStatus::SUCCESS); + + ASSERT_EQ(Tracer->SpanCount(), 1u); + EXPECT_EQ(Tracer->GetLastSpanRecord().Name, "ydb.ExecuteQuery"); +} + +TEST_F(QuerySpanTest, SpanKindIsClient) { + TQuerySpan span(Tracer, "CreateSession", "localhost:2135"); + span.End(EStatus::SUCCESS); + + ASSERT_EQ(Tracer->SpanCount(), 1u); + EXPECT_EQ(Tracer->GetLastSpanRecord().Kind, ESpanKind::CLIENT); +} + +TEST_F(QuerySpanTest, DbSystemAttribute) { + TQuerySpan span(Tracer, "ExecuteQuery", "localhost:2135"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("db.system.name"), "ydb"); +} + +TEST_F(QuerySpanTest, ServerAddressAndPort) { + TQuerySpan span(Tracer, "Commit", "ydb.server:2135"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "ydb.server"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_F(QuerySpanTest, ServerAddressCustomPort) { + TQuerySpan span(Tracer, "Rollback", "myhost:9090"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "myhost"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 9090); +} + +TEST_F(QuerySpanTest, ServerAddressNoPortDefaultsTo2135) { + TQuerySpan span(Tracer, "ExecuteQuery", "myhost"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "myhost"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_F(QuerySpanTest, IPv6EndpointParsing) { + TQuerySpan span(Tracer, "ExecuteQuery", "[::1]:2136"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "::1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2136); +} + +TEST_F(QuerySpanTest, IPv6EndpointNoPort) { + TQuerySpan span(Tracer, "ExecuteQuery", "[fe80::1]"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "fe80::1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_F(QuerySpanTest, PeerEndpointAttributes) { + TQuerySpan span(Tracer, "ExecuteQuery", "discovery.ydb:2135"); + span.SetPeerEndpoint("10.0.0.1:2136"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("network.peer.address"), "10.0.0.1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("network.peer.port"), 2136); +} + +TEST_F(QuerySpanTest, SuccessStatusRecorded) { + TQuerySpan span(Tracer, "Commit", "localhost:2135"); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->HasIntAttribute("db.response.status_code")); + EXPECT_EQ(fakeSpan->GetIntAttribute("db.response.status_code"), static_cast(EStatus::SUCCESS)); + EXPECT_FALSE(fakeSpan->HasStringAttribute("error.type")); +} + +TEST_F(QuerySpanTest, ErrorStatusSetsErrorType) { + TQuerySpan span(Tracer, "Rollback", "localhost:2135"); + span.End(EStatus::UNAVAILABLE); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetIntAttribute("db.response.status_code"), static_cast(EStatus::UNAVAILABLE)); + EXPECT_TRUE(fakeSpan->HasStringAttribute("error.type")); + EXPECT_FALSE(fakeSpan->GetStringAttribute("error.type").empty()); +} + +TEST_F(QuerySpanTest, SpanIsEndedAfterEnd) { + TQuerySpan span(Tracer, "ExecuteQuery", "localhost:2135"); + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + + EXPECT_FALSE(fakeSpan->IsEnded()); + span.End(EStatus::SUCCESS); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_F(QuerySpanTest, NullTracerDoesNotCrash) { + EXPECT_NO_THROW({ + TQuerySpan span(nullptr, "ExecuteQuery", "localhost:2135"); + span.SetPeerEndpoint("10.0.0.1:2136"); + span.AddEvent("retry", {{"attempt", "1"}}); + span.End(EStatus::SUCCESS); + }); +} + +TEST_F(QuerySpanTest, DestructorEndsSpan) { + auto fakeSpan = [&]() -> std::shared_ptr { + TQuerySpan span(Tracer, "CreateSession", "localhost:2135"); + return Tracer->GetLastSpan(); + }(); + + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_F(QuerySpanTest, ExplicitEndThenDestructorDoesNotDoubleEnd) { + auto fakeSpan = [&]() -> std::shared_ptr { + TQuerySpan span(Tracer, "Commit", "localhost:2135"); + span.End(EStatus::SUCCESS); + return Tracer->GetLastSpan(); + }(); + + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_F(QuerySpanTest, AddEventForwarded) { + TQuerySpan span(Tracer, "ExecuteQuery", "localhost:2135"); + span.AddEvent("retry", {{"ydb.attempt", "2"}, {"error.type", "UNAVAILABLE"}}); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + auto events = fakeSpan->GetEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events[0].Name, "retry"); + EXPECT_EQ(events[0].Attributes.at("ydb.attempt"), "2"); + EXPECT_EQ(events[0].Attributes.at("error.type"), "UNAVAILABLE"); +} + +TEST_F(QuerySpanTest, EmptyPeerEndpointIgnored) { + TQuerySpan span(Tracer, "CreateSession", "localhost:2135"); + span.SetPeerEndpoint(""); + span.End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_FALSE(fakeSpan->HasStringAttribute("network.peer.address")); + EXPECT_FALSE(fakeSpan->HasIntAttribute("network.peer.port")); +} + +TEST_F(QuerySpanTest, AllFourOperationNames) { + const std::vector operations = {"CreateSession", "ExecuteQuery", "Commit", "Rollback"}; + + for (const auto& op : operations) { + TQuerySpan span(Tracer, op, "localhost:2135"); + span.End(EStatus::SUCCESS); + } + + auto spans = Tracer->GetSpans(); + ASSERT_EQ(spans.size(), 4u); + EXPECT_EQ(spans[0].Name, "ydb.CreateSession"); + EXPECT_EQ(spans[1].Name, "ydb.ExecuteQuery"); + EXPECT_EQ(spans[2].Name, "ydb.Commit"); + EXPECT_EQ(spans[3].Name, "ydb.Rollback"); + + for (const auto& record : spans) { + EXPECT_EQ(record.Kind, ESpanKind::CLIENT); + } +} + +TEST_F(QuerySpanTest, MultipleErrorStatuses) { + std::vector errorStatuses = { + EStatus::BAD_REQUEST, + EStatus::UNAUTHORIZED, + EStatus::INTERNAL_ERROR, + EStatus::UNAVAILABLE, + EStatus::OVERLOADED, + EStatus::TIMEOUT, + EStatus::NOT_FOUND, + EStatus::CLIENT_INTERNAL_ERROR, + }; + + for (auto status : errorStatuses) { + TQuerySpan span(Tracer, "ExecuteQuery", "localhost:2135"); + span.End(status); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->HasStringAttribute("error.type")) + << "error.type missing for status " << static_cast(status); + EXPECT_EQ(fakeSpan->GetIntAttribute("db.response.status_code"), static_cast(status)); + } +} + +TEST_F(QuerySpanTest, EmptyEndpointDoesNotCrash) { + EXPECT_NO_THROW({ + TQuerySpan span(Tracer, "ExecuteQuery", ""); + span.End(EStatus::SUCCESS); + }); +} From 1fd13a69a763cb01f17fbd52c2a37de31112379c Mon Sep 17 00:00:00 2001 From: maladetska Date: Thu, 26 Mar 2026 02:22:00 +0800 Subject: [PATCH 10/93] make spans nested in traces --- examples/otel_tracing/README.md | 8 +- include/ydb-cpp-sdk/client/trace/trace.h | 6 + plugins/trace/otel/src/trace.cpp | 15 ++ .../impl/observability/operation_span.cpp | 153 ++++++++++++++++++ .../impl/observability/operation_span.h | 32 ++++ src/client/query/client.cpp | 16 +- src/client/query/impl/query_spans.cpp | 12 ++ src/client/query/impl/query_spans.h | 1 + tests/common/fake_trace_provider.h | 27 +++- tests/unit/client/query/query_spans_ut.cpp | 18 +++ 10 files changed, 276 insertions(+), 12 deletions(-) create mode 100644 src/client/impl/observability/operation_span.cpp create mode 100644 src/client/impl/observability/operation_span.h diff --git a/examples/otel_tracing/README.md b/examples/otel_tracing/README.md index 412ea575ad8..075d9b5e5ec 100644 --- a/examples/otel_tracing/README.md +++ b/examples/otel_tracing/README.md @@ -73,7 +73,7 @@ cmake --build . --target otel_tracing_example -j$(nproc) |-----------|------------------------------|---------------------------------| | Grafana | http://localhost:3000 | Дашборд "YDB QueryService" | | Jaeger | http://localhost:16686 | Поиск трейсов по сервису | -| Prometheus| http://localhost:9090 | Метрики `ydb_ydb_query_*` | +| Prometheus| http://localhost:9090 | Метрики `ydb_query_*` | **Grafana**: логин `admin` / пароль `admin`. @@ -96,9 +96,9 @@ cmake --build . --target otel_tracing_example -j$(nproc) - `db.response.status_code`, `error.type` (при ошибках) #### В Prometheus: -- `ydb_ydb_query_requests_total` — счётчик запросов -- `ydb_ydb_query_errors_total` — счётчик ошибок -- `ydb_ydb_query_latency_ms_bucket` — гистограмма задержек +- `ydb_query_requests_total` — счётчик запросов +- `ydb_query_errors_total` — счётчик ошибок +- `ydb_query_latency_ms_bucket` — гистограмма задержек ### 6. Остановить diff --git a/include/ydb-cpp-sdk/client/trace/trace.h b/include/ydb-cpp-sdk/client/trace/trace.h index b86297146a9..054fa258488 100644 --- a/include/ydb-cpp-sdk/client/trace/trace.h +++ b/include/ydb-cpp-sdk/client/trace/trace.h @@ -15,6 +15,11 @@ enum class ESpanKind { CONSUMER }; +class IScope { +public: + virtual ~IScope() = default; +}; + class ISpan { public: virtual ~ISpan() = default; @@ -22,6 +27,7 @@ class ISpan { virtual void SetAttribute(const std::string& key, const std::string& value) = 0; virtual void SetAttribute(const std::string& key, int64_t value) = 0; virtual void AddEvent(const std::string& name, const std::map& attributes = {}) = 0; + virtual std::unique_ptr Activate() = 0; }; class ITracer { diff --git a/plugins/trace/otel/src/trace.cpp b/plugins/trace/otel/src/trace.cpp index b904df3ab9d..b5e6d8ae1ea 100644 --- a/plugins/trace/otel/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include @@ -23,6 +24,16 @@ otel_trace::SpanKind MapSpanKind(ESpanKind kind) { return otel_trace::SpanKind::kInternal; } +class TOtelScope : public IScope { +public: + TOtelScope(otel_nostd::shared_ptr span) + : Scope_(std::move(span)) + {} + +private: + otel_trace::Scope Scope_; +}; + class TOtelSpan : public ISpan { public: TOtelSpan(otel_nostd::shared_ptr span) @@ -54,6 +65,10 @@ class TOtelSpan : public ISpan { } } + std::unique_ptr Activate() override { + return std::make_unique(Span_); + } + private: otel_nostd::shared_ptr Span_; }; diff --git a/src/client/impl/observability/operation_span.cpp b/src/client/impl/observability/operation_span.cpp new file mode 100644 index 00000000000..377608098eb --- /dev/null +++ b/src/client/impl/observability/operation_span.cpp @@ -0,0 +1,153 @@ +#include "operation_span.h" + +#include + +#include + +#include + +namespace NYdb::inline V3::NObservability { + +namespace { + +constexpr int DefaultGrpcPort = 2135; + +void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { + port = DefaultGrpcPort; + + if (endpoint.empty()) { + host = endpoint; + return; + } + + // IPv6 bracket notation: [addr]:port + if (endpoint.front() == '[') { + auto bracketEnd = endpoint.find(']'); + if (bracketEnd != std::string::npos) { + host = endpoint.substr(1, bracketEnd - 1); + if (bracketEnd + 2 < endpoint.size() && endpoint[bracketEnd + 1] == ':') { + try { + port = std::stoi(endpoint.substr(bracketEnd + 2)); + } catch (...) {} + } + return; + } + } + + auto pos = endpoint.rfind(':'); + if (pos != std::string::npos) { + host = endpoint.substr(0, pos); + try { + port = std::stoi(endpoint.substr(pos + 1)); + } catch (...) {} + } else { + host = endpoint; + } +} + +void SafeLogSpanError(TLog& log, const char* message) noexcept { + try { + try { + std::rethrow_exception(std::current_exception()); + } catch (const std::exception& e) { + LOG_LAZY(log, TLOG_ERR, std::string("TOperationSpan: ") + message + ": " + e.what()); + return; + } catch (...) { + } + LOG_LAZY(log, TLOG_ERR, std::string("TOperationSpan: ") + message + ": (unknown)"); + } catch (...) { + } +} + +} // namespace + +TOperationSpan::TOperationSpan(std::shared_ptr tracer, const std::string& operationName, + const std::string& endpoint, const TLog& log) + : Log_(log) +{ + if (!tracer) { + return; + } + + std::string host; + int port; + ParseEndpoint(endpoint, host, port); + + try { + Span_ = tracer->StartSpan("ydb." + operationName, NTrace::ESpanKind::CLIENT); + if (!Span_) { + return; + } + Span_->SetAttribute("db.system.name", "ydb"); + Span_->SetAttribute("server.address", host); + Span_->SetAttribute("server.port", static_cast(port)); + } catch (...) { + SafeLogSpanError(Log_, "failed to initialize span"); + Span_.reset(); + } +} + +TOperationSpan::~TOperationSpan() noexcept { + if (Span_) { + try { + Span_->End(); + } catch (...) { + SafeLogSpanError(Log_, "failed to end span"); + } + } +} + +void TOperationSpan::SetPeerEndpoint(const std::string& endpoint) noexcept { + if (!Span_ || endpoint.empty()) { + return; + } + try { + std::string host; + int port; + ParseEndpoint(endpoint, host, port); + Span_->SetAttribute("network.peer.address", host); + Span_->SetAttribute("network.peer.port", static_cast(port)); + } catch (...) { + SafeLogSpanError(Log_, "failed to set peer endpoint"); + } +} + +void TOperationSpan::AddEvent(const std::string& name, const std::map& attributes) noexcept { + if (!Span_) { + return; + } + try { + Span_->AddEvent(name, attributes); + } catch (...) { + SafeLogSpanError(Log_, "failed to add event"); + } +} + +std::unique_ptr TOperationSpan::Activate() noexcept { + if (!Span_) { + return nullptr; + } + try { + return Span_->Activate(); + } catch (...) { + SafeLogSpanError(Log_, "failed to activate span"); + return nullptr; + } +} + +void TOperationSpan::End(EStatus status) noexcept { + if (Span_) { + try { + Span_->SetAttribute("db.response.status_code", static_cast(status)); + if (status != EStatus::SUCCESS) { + Span_->SetAttribute("error.type", ToString(status)); + } + Span_->End(); + } catch (...) { + SafeLogSpanError(Log_, "failed to finalize span"); + } + Span_.reset(); + } +} + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/operation_span.h b/src/client/impl/observability/operation_span.h new file mode 100644 index 00000000000..5c17d578460 --- /dev/null +++ b/src/client/impl/observability/operation_span.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include + +namespace NYdb::inline V3::NObservability { + +class TOperationSpan { +public: + TOperationSpan(std::shared_ptr tracer, const std::string& operationName, + const std::string& endpoint, const TLog& log); + ~TOperationSpan() noexcept; + + void SetPeerEndpoint(const std::string& endpoint) noexcept; + void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; + std::unique_ptr Activate() noexcept; + + void End(EStatus status) noexcept; + +private: + TLog Log_; + std::shared_ptr Span_; +}; + +} // namespace NYdb::NObservability diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 2e26c95f8d1..8aefee460e2 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -601,6 +601,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public ), NSessionPool::PERIODIC_ACTION_INTERVAL); } + std::shared_ptr CreateRetrySpan(const std::string& operationName) { + return std::make_shared(Tracer_, operationName, DbDriverState_->DiscoveryEndpoint); + } + void CollectRetryStatAsync(EStatus status) { RetryOperationStatCollector_.IncAsyncRetryOperation(status); } @@ -718,13 +722,21 @@ TAsyncStatus TQueryClient::RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRet } TStatus TQueryClient::RetryQuerySync(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings) { + auto parentSpan = Impl_->CreateRetrySpan("RetryQuery"); + auto scope = parentSpan->Activate(); NRetry::Sync::TRetryWithSession ctx(*this, queryFunc, settings); - return ctx.Execute(); + auto status = ctx.Execute(); + parentSpan->End(status.GetStatus()); + return status; } TStatus TQueryClient::RetryQuerySync(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings) { + auto parentSpan = Impl_->CreateRetrySpan("RetryQuery"); + auto scope = parentSpan->Activate(); NRetry::Sync::TRetryWithoutSession ctx(*this, queryFunc, settings); - return ctx.Execute(); + auto status = ctx.Execute(); + parentSpan->End(status.GetStatus()); + return status; } TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, const TTxControl& txControl, diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index 4532b0e6552..72d3e2eec0c 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -115,6 +115,18 @@ void TQuerySpan::AddEvent(const std::string& name, const std::map TQuerySpan::Activate() noexcept { + if (!Span_) { + return nullptr; + } + try { + return Span_->Activate(); + } catch (...) { + SafeLogSpanError("failed to activate span"); + return nullptr; + } +} + void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index 096634346b3..2aeb5d5b79b 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -17,6 +17,7 @@ class TQuerySpan { void SetPeerEndpoint(const std::string& endpoint) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; + std::unique_ptr Activate() noexcept; void End(EStatus status) noexcept; diff --git a/tests/common/fake_trace_provider.h b/tests/common/fake_trace_provider.h index 79d38337ab5..90fd7de8e49 100644 --- a/tests/common/fake_trace_provider.h +++ b/tests/common/fake_trace_provider.h @@ -12,7 +12,10 @@ struct TFakeEvent { std::map Attributes; }; -class TFakeSpan : public NMetrics::ISpan { +class TFakeScope : public NTrace::IScope { +}; + +class TFakeSpan : public NTrace::ISpan { public: void End() override { std::lock_guard lock(Mutex_); @@ -34,11 +37,22 @@ class TFakeSpan : public NMetrics::ISpan { Events_.push_back({name, attributes}); } + std::unique_ptr Activate() override { + std::lock_guard lock(Mutex_); + Activated_ = true; + return std::make_unique(); + } + bool IsEnded() const { std::lock_guard lock(Mutex_); return Ended_; } + bool IsActivated() const { + std::lock_guard lock(Mutex_); + return Activated_; + } + std::string GetStringAttribute(const std::string& key) const { std::lock_guard lock(Mutex_); auto it = StringAttributes_.find(key); @@ -69,14 +83,15 @@ class TFakeSpan : public NMetrics::ISpan { private: mutable std::mutex Mutex_; bool Ended_ = false; + bool Activated_ = false; std::map StringAttributes_; std::map IntAttributes_; std::vector Events_; }; -class TFakeTracer : public NMetrics::ITracer { +class TFakeTracer : public NTrace::ITracer { public: - std::shared_ptr StartSpan(const std::string& name, NMetrics::ESpanKind kind) override { + std::shared_ptr StartSpan(const std::string& name, NTrace::ESpanKind kind) override { auto span = std::make_shared(); std::lock_guard lock(Mutex_); Spans_.push_back({name, kind, span}); @@ -85,7 +100,7 @@ class TFakeTracer : public NMetrics::ITracer { struct TSpanRecord { std::string Name; - NMetrics::ESpanKind Kind; + NTrace::ESpanKind Kind; std::shared_ptr Span; }; @@ -114,9 +129,9 @@ class TFakeTracer : public NMetrics::ITracer { std::vector Spans_; }; -class TFakeTraceProvider : public NMetrics::ITraceProvider { +class TFakeTraceProvider : public NTrace::ITraceProvider { public: - std::shared_ptr GetTracer(const std::string& name) override { + std::shared_ptr GetTracer(const std::string& name) override { std::lock_guard lock(Mutex_); auto it = Tracers_.find(name); if (it != Tracers_.end()) { diff --git a/tests/unit/client/query/query_spans_ut.cpp b/tests/unit/client/query/query_spans_ut.cpp index ae553909413..8618c2a0d1d 100644 --- a/tests/unit/client/query/query_spans_ut.cpp +++ b/tests/unit/client/query/query_spans_ut.cpp @@ -240,3 +240,21 @@ TEST_F(QuerySpanTest, EmptyEndpointDoesNotCrash) { span.End(EStatus::SUCCESS); }); } + +TEST_F(QuerySpanTest, ActivateReturnsScope) { + TQuerySpan span(Tracer, "RetryQuery", "localhost:2135"); + auto scope = span.Activate(); + EXPECT_NE(scope, nullptr); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsActivated()); + + span.End(EStatus::SUCCESS); +} + +TEST_F(QuerySpanTest, ActivateNullTracerReturnsNull) { + TQuerySpan span(nullptr, "RetryQuery", "localhost:2135"); + auto scope = span.Activate(); + EXPECT_EQ(scope, nullptr); +} From 650b039697281aac82d3b344f9e976e70e763a4e Mon Sep 17 00:00:00 2001 From: maladetska Date: Mon, 30 Mar 2026 01:26:32 +0300 Subject: [PATCH 11/93] [draft] --- examples/otel_tracing/CMakeLists.txt | 1 + examples/otel_tracing/README.md | 22 +-- examples/otel_tracing/docker-compose.yml | 1 + .../grafana/dashboards/ydb-query-service.json | 28 +-- examples/otel_tracing/main.cpp | 164 +++++++++++++----- src/client/table/impl/CMakeLists.txt | 1 + src/client/table/impl/table_client.cpp | 47 ++++- src/client/table/impl/table_client.h | 22 ++- src/client/table/impl/table_spans.cpp | 91 ++++++++++ src/client/table/impl/table_spans.h | 24 +++ 10 files changed, 322 insertions(+), 79 deletions(-) create mode 100644 src/client/table/impl/table_spans.cpp create mode 100644 src/client/table/impl/table_spans.h diff --git a/examples/otel_tracing/CMakeLists.txt b/examples/otel_tracing/CMakeLists.txt index e98ad9d3751..b826c66688b 100644 --- a/examples/otel_tracing/CMakeLists.txt +++ b/examples/otel_tracing/CMakeLists.txt @@ -4,6 +4,7 @@ target_link_libraries(otel_tracing_example PUBLIC yutil getopt YDB-CPP-SDK::Query + YDB-CPP-SDK::Table YDB-CPP-SDK::Params YDB-CPP-SDK::Driver YDB-CPP-SDK::OpenTelemetryTrace diff --git a/examples/otel_tracing/README.md b/examples/otel_tracing/README.md index 075d9b5e5ec..0eeb4ddca47 100644 --- a/examples/otel_tracing/README.md +++ b/examples/otel_tracing/README.md @@ -1,6 +1,6 @@ -# YDB C++ SDK — OpenTelemetry Tracing Demo +# YDB C++ SDK — OpenTelemetry Demo -Демонстрация трассировки операций QueryService (CreateSession, ExecuteQuery, Commit, Rollback) +Демонстрация трассировки и метрик операций QueryService и TableService с визуализацией в **Grafana**, **Jaeger** и **Prometheus**. ## Архитектура @@ -73,32 +73,32 @@ cmake --build . --target otel_tracing_example -j$(nproc) |-----------|------------------------------|---------------------------------| | Grafana | http://localhost:3000 | Дашборд "YDB QueryService" | | Jaeger | http://localhost:16686 | Поиск трейсов по сервису | -| Prometheus| http://localhost:9090 | Метрики `ydb_query_*` | +| Prometheus| http://localhost:9090 | Метрики `db_client_operation_*` | **Grafana**: логин `admin` / пароль `admin`. ### 5. Что смотреть #### В Grafana (дашборд "YDB QueryService"): -- **Request Rate by Operation** — RPS по операциям (ExecuteQuery, CreateSession, Commit, Rollback) +- **Request Rate by Operation** — RPS по операциям (ExecuteQuery, ExecuteDataQuery, CreateSession, Commit, Rollback) - **Error Rate by Operation** — частота ошибок -- **Latency p50/p95/p99** — распределение задержек +- **Duration p50/p95/p99** — распределение длительности операций - **Error Ratio** — процент ошибок - **Recent Traces** — таблица трейсов из Jaeger #### В Jaeger UI: - Выберите сервис `ydb-cpp-sdk-demo` -- Каждый спан содержит атрибуты: - - `db.system.name` = `ydb` +- Трейсы от QueryService содержат спаны с атрибутами: + - `db.system.name` = `other_sql` - `server.address`, `server.port` - - `network.peer.address`, `network.peer.port` - `db.query.text` (для ExecuteQuery) - `db.response.status_code`, `error.type` (при ошибках) +- TableService пока генерирует только метрики (без спанов) #### В Prometheus: -- `ydb_query_requests_total` — счётчик запросов -- `ydb_query_errors_total` — счётчик ошибок -- `ydb_query_latency_ms_bucket` — гистограмма задержек +- `db_client_operation_requests_total` — счётчик запросов по операциям +- `db_client_operation_errors_total` — счётчик ошибок по операциям +- `db_client_operation_duration_seconds_bucket` — гистограмма длительности (OTel Semantic Conventions) ### 6. Остановить diff --git a/examples/otel_tracing/docker-compose.yml b/examples/otel_tracing/docker-compose.yml index 3314c27d421..9d01c8fa823 100644 --- a/examples/otel_tracing/docker-compose.yml +++ b/examples/otel_tracing/docker-compose.yml @@ -10,6 +10,7 @@ services: - GRPC_PORT=2136 - MON_PORT=8765 - YDB_DEFAULT_LOG_LEVEL=NOTICE + - YDB_USE_IN_MEMORY_PDISKS=true volumes: - ydb-data:/ydb_data healthcheck: diff --git a/examples/otel_tracing/grafana/dashboards/ydb-query-service.json b/examples/otel_tracing/grafana/dashboards/ydb-query-service.json index 1ce6c11a8f7..12a38d99f15 100644 --- a/examples/otel_tracing/grafana/dashboards/ydb-query-service.json +++ b/examples/otel_tracing/grafana/dashboards/ydb-query-service.json @@ -18,8 +18,8 @@ }, "targets": [ { - "expr": "rate(ydb_query_requests_total[1m])", - "legendFormat": "{{operation}}" + "expr": "rate(db_client_operation_requests_total[1m])", + "legendFormat": "{{db_operation_name}}" } ] }, @@ -37,34 +37,34 @@ }, "targets": [ { - "expr": "rate(ydb_query_errors_total[1m])", - "legendFormat": "{{operation}}" + "expr": "rate(db_client_operation_errors_total[1m])", + "legendFormat": "{{db_operation_name}}" } ] }, { - "title": "Latency p50 / p95 / p99 (ms)", + "title": "Duration p50 / p95 / p99 (s)", "type": "timeseries", "gridPos": { "h": 8, "w": 12, "x": 0, "y": 8 }, "datasource": { "type": "prometheus", "uid": "${prometheus_ds}" }, "fieldConfig": { "defaults": { - "unit": "ms", + "unit": "s", "custom": { "drawStyle": "line", "fillOpacity": 5 } } }, "targets": [ { - "expr": "histogram_quantile(0.50, rate(ydb_query_latency_ms_bucket[1m]))", - "legendFormat": "p50 {{operation}}" + "expr": "histogram_quantile(0.50, rate(db_client_operation_duration_seconds_bucket[1m]))", + "legendFormat": "p50 {{db_operation_name}}" }, { - "expr": "histogram_quantile(0.95, rate(ydb_query_latency_ms_bucket[1m]))", - "legendFormat": "p95 {{operation}}" + "expr": "histogram_quantile(0.95, rate(db_client_operation_duration_seconds_bucket[1m]))", + "legendFormat": "p95 {{db_operation_name}}" }, { - "expr": "histogram_quantile(0.99, rate(ydb_query_latency_ms_bucket[1m]))", - "legendFormat": "p99 {{operation}}" + "expr": "histogram_quantile(0.99, rate(db_client_operation_duration_seconds_bucket[1m]))", + "legendFormat": "p99 {{db_operation_name}}" } ] }, @@ -87,8 +87,8 @@ }, "targets": [ { - "expr": "sum(rate(ydb_query_errors_total[5m])) by (operation) / sum(rate(ydb_query_requests_total[5m])) by (operation)", - "legendFormat": "{{operation}}" + "expr": "sum(rate(db_client_operation_errors_total[5m])) by (db_operation_name) / sum(rate(db_client_operation_requests_total[5m])) by (db_operation_name)", + "legendFormat": "{{db_operation_name}}" } ] }, diff --git a/examples/otel_tracing/main.cpp b/examples/otel_tracing/main.cpp index f8e53e23669..80e31a9f2e8 100644 --- a/examples/otel_tracing/main.cpp +++ b/examples/otel_tracing/main.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include @@ -19,6 +20,9 @@ #include +#include +#include + #include #include #include @@ -30,11 +34,10 @@ namespace otlp = opentelemetry::exporter::otlp; namespace resource = opentelemetry::sdk::resource; using namespace NYdb; -using namespace NYdb::NQuery; using namespace NYdb::NStatusHelpers; struct TConfig { - std::string Endpoint = "grpc://localhost:2136"; + std::string Endpoint = "localhost:2136"; std::string Database = "/local"; std::string OtlpEndpoint = "http://localhost:4328"; int Iterations = 20; @@ -82,37 +85,54 @@ nostd::shared_ptr InitMetrics(const TConf return nostd::shared_ptr(provider); } -void RunWorkload(TQueryClient& client, int iterations) { - std::cout << "=== Creating table ===" << std::endl; +nostd::shared_ptr GetAppTracer() { + return opentelemetry::trace::Provider::GetTracerProvider()->GetTracer("ydb-demo-app", "1.0.0"); +} + +void RunQueryWorkload(NQuery::TQueryClient& client, int iterations) { + std::cout << "\n=== Query Service workload ===" << std::endl; + + auto tracer = GetAppTracer(); - ThrowOnError(client.RetryQuerySync([](TSession session) { - return session.ExecuteQuery(R"( - CREATE TABLE IF NOT EXISTS otel_demo ( - id Uint64, - value Utf8, - PRIMARY KEY (id) - ) - )", TTxControl::NoTx()).GetValueSync(); - })); + { + auto ddlSpan = tracer->StartSpan("QueryService.DDL"); + auto scope = opentelemetry::trace::Scope(ddlSpan); + + ThrowOnError(client.RetryQuerySync([](NQuery::TSession session) { + return session.ExecuteQuery(R"( + CREATE TABLE IF NOT EXISTS otel_demo ( + id Uint64, + value Utf8, + PRIMARY KEY (id) + ) + )", NQuery::TTxControl::NoTx()).GetValueSync(); + })); + + ddlSpan->SetStatus(opentelemetry::trace::StatusCode::kOk); + } for (int i = 0; i < iterations; ++i) { - std::cout << "--- Iteration " << (i + 1) << "/" << iterations << " ---" << std::endl; + auto iterSpan = tracer->StartSpan("QueryService.Iteration"); + auto scope = opentelemetry::trace::Scope(iterSpan); + iterSpan->SetAttribute("iteration", static_cast(i + 1)); + + std::cout << " [Query] Iteration " << (i + 1) << "/" << iterations << std::endl; - ThrowOnError(client.RetryQuerySync([i](TSession session) { + ThrowOnError(client.RetryQuerySync([i](NQuery::TSession session) { auto params = TParamsBuilder() .AddParam("$id").Uint64(i).Build() - .AddParam("$val").Utf8("item_" + std::to_string(i)).Build() + .AddParam("$val").Utf8("query_" + std::to_string(i)).Build() .Build(); return session.ExecuteQuery(R"( DECLARE $id AS Uint64; DECLARE $val AS Utf8; UPSERT INTO otel_demo (id, value) VALUES ($id, $val) - )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + )", NQuery::TTxControl::BeginTx(NQuery::TTxSettings::SerializableRW()).CommitTx(), params).GetValueSync(); })); - ThrowOnError(client.RetryQuerySync([i](TSession session) { + ThrowOnError(client.RetryQuerySync([i](NQuery::TSession session) { auto params = TParamsBuilder() .AddParam("$id").Uint64(i).Build() .Build(); @@ -120,21 +140,86 @@ void RunWorkload(TQueryClient& client, int iterations) { return session.ExecuteQuery(R"( DECLARE $id AS Uint64; SELECT id, value FROM otel_demo WHERE id = $id - )", TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + )", NQuery::TTxControl::BeginTx(NQuery::TTxSettings::SerializableRW()).CommitTx(), params).GetValueSync(); })); - ThrowOnError(client.RetryQuerySync([](TQueryClient client) -> TStatus { - auto session = client.GetSession().GetValueSync().GetSession(); - auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).GetValueSync(); + if (i % 5 == 4) { + ThrowOnError(client.RetryQuerySync([](NQuery::TQueryClient client) -> TStatus { + auto session = client.GetSession().GetValueSync().GetSession(); + auto beginResult = session.BeginTransaction(NQuery::TTxSettings::SerializableRW()).GetValueSync(); + if (!beginResult.IsSuccess()) { + return beginResult; + } + auto tx = beginResult.GetTransaction(); + + auto result = session.ExecuteQuery(R"( + SELECT COUNT(*) AS cnt FROM otel_demo + )", NQuery::TTxControl::Tx(tx)).GetValueSync(); + + if (!result.IsSuccess()) { + return result; + } + + return tx.Commit().GetValueSync(); + })); + } + + iterSpan->SetStatus(opentelemetry::trace::StatusCode::kOk); + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } +} + +void RunTableWorkload(NTable::TTableClient& client, int iterations) { + std::cout << "\n=== Table Service workload ===" << std::endl; + + auto tracer = GetAppTracer(); + + for (int i = 0; i < iterations; ++i) { + int id = 1000 + i; + + auto iterSpan = tracer->StartSpan("TableService.Iteration"); + auto scope = opentelemetry::trace::Scope(iterSpan); + iterSpan->SetAttribute("iteration", static_cast(i + 1)); + + std::cout << " [Table] Iteration " << (i + 1) << "/" << iterations << std::endl; + + ThrowOnError(client.RetryOperationSync([id](NTable::TSession session) { + auto params = session.GetParamsBuilder() + .AddParam("$id").Uint64(id).Build() + .AddParam("$val").Utf8("table_" + std::to_string(id)).Build() + .Build(); + + return session.ExecuteDataQuery(R"( + DECLARE $id AS Uint64; + DECLARE $val AS Utf8; + UPSERT INTO otel_demo (id, value) VALUES ($id, $val) + )", NTable::TTxControl::BeginTx(NTable::TTxSettings::SerializableRW()).CommitTx(), + std::move(params)).GetValueSync(); + })); + + ThrowOnError(client.RetryOperationSync([id](NTable::TSession session) { + auto params = session.GetParamsBuilder() + .AddParam("$id").Uint64(id).Build() + .Build(); + + return session.ExecuteDataQuery(R"( + DECLARE $id AS Uint64; + SELECT id, value FROM otel_demo WHERE id = $id + )", NTable::TTxControl::BeginTx(NTable::TTxSettings::SerializableRW()).CommitTx(), + std::move(params)).GetValueSync(); + })); + + ThrowOnError(client.RetryOperationSync([](NTable::TSession session) -> TStatus { + auto beginResult = session.BeginTransaction(NTable::TTxSettings::SerializableRW()).GetValueSync(); if (!beginResult.IsSuccess()) { return beginResult; } auto tx = beginResult.GetTransaction(); - auto result = session.ExecuteQuery(R"( + auto result = session.ExecuteDataQuery(R"( SELECT COUNT(*) AS cnt FROM otel_demo - )", TTxControl::Tx(tx)).GetValueSync(); + )", NTable::TTxControl::Tx(tx)).GetValueSync(); if (!result.IsSuccess()) { return result; @@ -144,9 +229,8 @@ void RunWorkload(TQueryClient& client, int iterations) { })); if (i % 5 == 4) { - auto rollbackResult = client.RetryQuerySync([](TQueryClient client) -> TStatus { - auto session = client.GetSession().GetValueSync().GetSession(); - auto beginResult = session.BeginTransaction(TTxSettings::SerializableRW()).GetValueSync(); + auto rollbackResult = client.RetryOperationSync([](NTable::TSession session) -> TStatus { + auto beginResult = session.BeginTransaction(NTable::TTxSettings::SerializableRW()).GetValueSync(); if (!beginResult.IsSuccess()) { return beginResult; } @@ -158,15 +242,9 @@ void RunWorkload(TQueryClient& client, int iterations) { } } - std::this_thread::sleep_for(std::chrono::milliseconds(500)); + iterSpan->SetStatus(opentelemetry::trace::StatusCode::kOk); + std::this_thread::sleep_for(std::chrono::milliseconds(200)); } - - std::cout << "=== Dropping table ===" << std::endl; - - ThrowOnError(client.RetryQuerySync([](TSession session) { - return session.ExecuteQuery( - "DROP TABLE otel_demo", TTxControl::NoTx()).GetValueSync(); - })); } int main(int argc, char** argv) { @@ -203,10 +281,18 @@ int main(int argc, char** argv) { .SetMetricRegistry(ydbMetricRegistry); TDriver driver(driverConfig); - TQueryClient client(driver); + NQuery::TQueryClient queryClient(driver); + NTable::TTableClient tableClient(driver); try { - RunWorkload(client, cfg.Iterations); + RunQueryWorkload(queryClient, cfg.Iterations); + RunTableWorkload(tableClient, cfg.Iterations); + + std::cout << "\n=== Cleanup ===" << std::endl; + ThrowOnError(queryClient.RetryQuerySync([](NQuery::TSession session) { + return session.ExecuteQuery( + "DROP TABLE otel_demo", NQuery::TTxControl::NoTx()).GetValueSync(); + })); } catch (const std::exception& e) { std::cerr << "Error: " << e.what() << std::endl; } @@ -225,8 +311,8 @@ int main(int argc, char** argv) { std::this_thread::sleep_for(std::chrono::seconds(3)); std::cout << "Done. Open Grafana at http://localhost:3000" << std::endl; - std::cout << " Dashboard: YDB QueryService" << std::endl; - std::cout << " Also: Jaeger UI at http://localhost:16686" << std::endl; + std::cout << " Jaeger UI at http://localhost:16686" << std::endl; + std::cout << " Prometheus at http://localhost:9090" << std::endl; return 0; } diff --git a/src/client/table/impl/CMakeLists.txt b/src/client/table/impl/CMakeLists.txt index 8ecfe4ead87..8e14331f062 100644 --- a/src/client/table/impl/CMakeLists.txt +++ b/src/client/table/impl/CMakeLists.txt @@ -22,6 +22,7 @@ target_sources(client-ydb_table-impl PRIVATE readers.cpp request_migrator.cpp table_client.cpp + table_spans.cpp transaction.cpp ) diff --git a/src/client/table/impl/table_client.cpp b/src/client/table/impl/table_client.cpp index 4df9e91e24e..92876f3d4fe 100644 --- a/src/client/table/impl/table_client.cpp +++ b/src/client/table/impl/table_client.cpp @@ -22,6 +22,12 @@ TTableClient::TImpl::TImpl(std::shared_ptr&& connections, , Settings_(settings) , SessionPool_(Settings_.SessionPoolSettings_.MaxActiveSessions_) { + MetricRegistry_ = Connections_->GetExternalMetricRegistry(); + + if (auto traceProvider = Connections_->GetTraceProvider()) { + Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-table"); + } + if (!DbDriverState_->StatCollector.IsCollecting()) { return; } @@ -378,8 +384,10 @@ TAsyncCreateSessionResult TTableClient::TImpl::CreateSession(const TCreateSessio auto createSessionPromise = NewPromise(); auto self = shared_from_this(); + auto metrics = std::make_shared(MetricRegistry_, "CreateSession"); + auto span = std::make_shared(Tracer_, "CreateSession", DbDriverState_->DiscoveryEndpoint); - auto createSessionExtractor = [createSessionPromise, self, standalone] + auto createSessionExtractor = [createSessionPromise, self, standalone, metrics, span] (google::protobuf::Any* any, TPlainStatus status) mutable { Ydb::Table::CreateSessionResult result; if (any) { @@ -392,10 +400,11 @@ TAsyncCreateSessionResult TTableClient::TImpl::CreateSession(const TCreateSessio } self->DbDriverState_->StatCollector.IncSessionsOnHost(status.Endpoint); } else { - // We do not use SessionStatusInterception for CreateSession request session.SessionImpl_->MarkBroken(); } TCreateSessionResult val(TStatus(std::move(status)), std::move(session)); + metrics->End(val.GetStatus()); + span->End(val.GetStatus()); createSessionPromise.SetValue(std::move(val)); }; @@ -759,11 +768,21 @@ TAsyncStatus TTableClient::TImpl::ExecuteSchemeQuery(const TSession& session, co request.set_session_id(TStringType{session.GetId()}); request.set_yql_text(TStringType{query}); - return RunSimple( + auto metrics = std::make_shared(MetricRegistry_, "ExecuteSchemeQuery"); + auto span = std::make_shared(Tracer_, "ExecuteSchemeQuery", DbDriverState_->DiscoveryEndpoint); + + auto future = RunSimple( std::move(request), &Ydb::Table::V1::TableService::Stub::AsyncExecuteSchemeQuery, rpcSettings ); + + return future.Apply([metrics, span](NThreading::TFuture f) mutable { + auto status = f.ExtractValue(); + metrics->End(status.GetStatus()); + span->End(status.GetStatus()); + return status; + }); } TAsyncBeginTransactionResult TTableClient::TImpl::BeginTransaction(const TSession& session, const TTxSettings& txSettings, @@ -777,8 +796,9 @@ TAsyncBeginTransactionResult TTableClient::TImpl::BeginTransaction(const TSessio SetTxSettings(txSettings, request.mutable_tx_settings()); auto promise = NewPromise(); + auto span = std::make_shared(Tracer_, "BeginTransaction", DbDriverState_->DiscoveryEndpoint); - auto extractor = [promise, session] + auto extractor = [promise, session, span] (google::protobuf::Any* any, TPlainStatus status) mutable { std::string txId; if (any) { @@ -789,6 +809,7 @@ TAsyncBeginTransactionResult TTableClient::TImpl::BeginTransaction(const TSessio TBeginTransactionResult beginTxResult(TStatus(std::move(status)), TTransaction(session, txId)); + span->End(beginTxResult.GetStatus()); promise.SetValue(std::move(beginTxResult)); }; @@ -816,8 +837,10 @@ TAsyncCommitTransactionResult TTableClient::TImpl::CommitTransaction(const TSess request.set_collect_stats(GetStatsCollectionMode(settings.CollectQueryStats_)); auto promise = NewPromise(); + auto metrics = std::make_shared(MetricRegistry_, "Commit"); + auto span = std::make_shared(Tracer_, "Commit", DbDriverState_->DiscoveryEndpoint); - auto extractor = [promise] + auto extractor = [promise, metrics, span] (google::protobuf::Any* any, TPlainStatus status) mutable { std::optional queryStats; if (any) { @@ -830,6 +853,8 @@ TAsyncCommitTransactionResult TTableClient::TImpl::CommitTransaction(const TSess } TCommitTransactionResult commitTxResult(TStatus(std::move(status)), queryStats); + metrics->End(commitTxResult.GetStatus()); + span->End(commitTxResult.GetStatus()); promise.SetValue(std::move(commitTxResult)); }; @@ -855,11 +880,21 @@ TAsyncStatus TTableClient::TImpl::RollbackTransaction(const TSession& session, c request.set_session_id(TStringType{session.GetId()}); request.set_tx_id(TStringType{txId}); - return RunSimple( + auto metrics = std::make_shared(MetricRegistry_, "Rollback"); + auto span = std::make_shared(Tracer_, "Rollback", DbDriverState_->DiscoveryEndpoint); + + auto future = RunSimple( std::move(request), &Ydb::Table::V1::TableService::Stub::AsyncRollbackTransaction, rpcSettings ); + + return future.Apply([metrics, span](NThreading::TFuture f) mutable { + auto status = f.ExtractValue(); + metrics->End(status.GetStatus()); + span->End(status.GetStatus()); + return status; + }); } TAsyncExplainDataQueryResult TTableClient::TImpl::ExplainDataQuery(const TSession& session, const std::string& query, diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 8fe71287f36..0e3e4734462 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -18,6 +18,9 @@ #include "request_migrator.h" #include "readers.h" +#include "table_metrics.h" +#include "table_spans.h" + #include @@ -237,16 +240,11 @@ class TTableClient::TImpl: public TClientImplCommon, public auto promise = NewPromise(); bool keepInCache = settings.KeepInQueryCache_ && settings.KeepInQueryCache_.value(); - // We don't want to delay call of TSession dtor, so we can't capture it by copy - // otherwise we break session pool and other clients logic. - // Same problem with TDataQuery and TTransaction - // - // The fast solution is: - // - create copy of TSession out of lambda - // - capture pointer - // - call free just before SetValue call + auto metrics = std::make_shared(MetricRegistry_, "ExecuteDataQuery"); + auto span = std::make_shared(Tracer_, "ExecuteDataQuery", DbDriverState_->DiscoveryEndpoint); + auto sessionPtr = new TSession(session); - auto extractor = [promise, sessionPtr, query, fromCache, keepInCache] + auto extractor = [promise, sessionPtr, query, fromCache, keepInCache, metrics, span] (google::protobuf::Any* any, TPlainStatus status) mutable { std::vector res; std::optional tx; @@ -285,6 +283,9 @@ class TTableClient::TImpl: public TClientImplCommon, public TDataQueryResult dataQueryResult(TStatus(std::move(status)), std::move(res), tx, dataQuery, fromCache, queryStats); + metrics->End(dataQueryResult.GetStatus()); + span->End(dataQueryResult.GetStatus()); + delete sessionPtr; tx.reset(); dataQuery.reset(); @@ -326,6 +327,9 @@ class TTableClient::TImpl: public TClientImplCommon, public NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> ParamsSizeHistogram; NSdkStats::TAtomicCounter<::NMonitoring::TRate> SessionRemovedDueBalancing; + std::shared_ptr MetricRegistry_; + std::shared_ptr Tracer_; + private: NSessionPool::TSessionPool SessionPool_; TRequestMigrator RequestMigrator_; diff --git a/src/client/table/impl/table_spans.cpp b/src/client/table/impl/table_spans.cpp new file mode 100644 index 00000000000..2593ef2a0b3 --- /dev/null +++ b/src/client/table/impl/table_spans.cpp @@ -0,0 +1,91 @@ +#include "table_spans.h" + +#include + +namespace NYdb::inline V3::NTable { + +namespace { + +constexpr int DefaultGrpcPort = 2135; + +void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { + port = DefaultGrpcPort; + + if (endpoint.empty()) { + host = endpoint; + return; + } + + if (endpoint.front() == '[') { + auto bracketEnd = endpoint.find(']'); + if (bracketEnd != std::string::npos) { + host = endpoint.substr(1, bracketEnd - 1); + if (bracketEnd + 2 < endpoint.size() && endpoint[bracketEnd + 1] == ':') { + try { + port = std::stoi(endpoint.substr(bracketEnd + 2)); + } catch (...) {} + } + return; + } + } + + auto pos = endpoint.rfind(':'); + if (pos != std::string::npos) { + host = endpoint.substr(0, pos); + try { + port = std::stoi(endpoint.substr(pos + 1)); + } catch (...) {} + } else { + host = endpoint; + } +} + +} // namespace + +TTableSpan::TTableSpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { + if (!tracer) { + return; + } + + std::string host; + int port; + ParseEndpoint(endpoint, host, port); + + try { + Span_ = tracer->StartSpan(operationName, NMetrics::ESpanKind::CLIENT); + if (!Span_) { + return; + } + Span_->SetAttribute("db.system.name", "other_sql"); + Span_->SetAttribute("db.operation.name", operationName); + Span_->SetAttribute("server.address", host); + Span_->SetAttribute("server.port", static_cast(port)); + } catch (...) { + Span_.reset(); + } +} + +TTableSpan::~TTableSpan() noexcept { + if (Span_) { + try { + Span_->End(); + } catch (...) { + } + } +} + +void TTableSpan::End(EStatus status) noexcept { + if (Span_) { + try { + Span_->SetAttribute("db.response.status_code", ToString(status)); + if (status != EStatus::SUCCESS) { + Span_->SetAttribute("error.type", ToString(status)); + } + Span_->End(); + } catch (...) { + } + Span_.reset(); + } +} + +} // namespace NYdb::NTable diff --git a/src/client/table/impl/table_spans.h b/src/client/table/impl/table_spans.h new file mode 100644 index 00000000000..283249b782e --- /dev/null +++ b/src/client/table/impl/table_spans.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace NYdb::inline V3::NTable { + +class TTableSpan { +public: + TTableSpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); + ~TTableSpan() noexcept; + + void End(EStatus status) noexcept; + +private: + std::shared_ptr Span_; +}; + +} // namespace NYdb::NTable From 27fb2a695c8baff37d78d237f5a24e27e4f5981d Mon Sep 17 00:00:00 2001 From: maladetska Date: Fri, 3 Apr 2026 10:43:50 +0300 Subject: [PATCH 12/93] fix test add retry in demo --- examples/otel_tracing/README.md | 100 +++++++++++++++++++++--- examples/otel_tracing/main.cpp | 118 +++++++++++++++++++++++++++++ tests/integration/metrics/main.cpp | 4 +- 3 files changed, 209 insertions(+), 13 deletions(-) diff --git a/examples/otel_tracing/README.md b/examples/otel_tracing/README.md index 0eeb4ddca47..f3def443d96 100644 --- a/examples/otel_tracing/README.md +++ b/examples/otel_tracing/README.md @@ -64,9 +64,69 @@ cmake --build . --target otel_tracing_example -j$(nproc) --endpoint grpc://localhost:2136 \ --database /local \ --otlp http://localhost:4328 \ - --iterations 20 + --iterations 20 \ + --retry-workers 6 \ + --retry-ops 30 ``` +#### Доступные флаги + +| Флаг | По умолчанию | Описание | +|--------------------|---------------------------|--------------------------------------------------------------------------| +| `--endpoint`, `-e` | `localhost:2136` | gRPC-эндпоинт YDB | +| `--database`, `-d` | `/local` | Имя базы | +| `--otlp` | `http://localhost:4328` | OTLP/HTTP endpoint коллектора | +| `--iterations`,`-n`| `20` | Итераций в Query- и Table-нагрузке | +| `--retry-workers` | `6` | Параллельных воркеров в retry-нагрузке (`0` чтобы пропустить) | +| `--retry-ops` | `30` | Операций на каждого retry-воркера | + +#### Демонстрация реальных ретраев + +Третий встроенный сценарий — `RunRetryWorkload` — намеренно провоцирует +**SERIALIZABLE-конфликты**: N параллельных воркеров делают +`SELECT → sleep → UPSERT → COMMIT` на одной и той же «горячей» строке +(`id = 9999`) внутри `RetryQuerySync`. YDB возвращает `ABORTED` +проигравшим транзакциям, и SDK прозрачно ретраит их. + +В трейсах появятся: + +``` +ydb.RunWithRetry (INTERNAL) +├── ydb.Try attempt=0 (INTERNAL, backoff_ms=0) +│ ├── ydb.CreateSession +│ ├── ydb.ExecuteQuery +│ └── ydb.Commit status=ABORTED, error.type=ABORTED, exception event +├── ydb.Try attempt=1 (INTERNAL, backoff_ms=...) +│ └── ... status=ABORTED +└── ydb.Try attempt=N (INTERNAL) + └── ... status=SUCCESS +``` + +Для усиления конфликтов поднимите воркеров и операций: + +```bash +./examples/otel_tracing/otel_tracing_example \ + --retry-workers 12 --retry-ops 80 +``` + +В конце программа печатает счётчик наблюдённых абортов — каждый из них +соответствует одному автоматическому ретраю SDK. + +> **Важно:** для статуса `ABORTED` SDK использует политику +> `RetryImmediately` (см. `src/client/impl/internal/retry/retry.h`), +> поэтому атрибут `ydb.retry.backoff_ms` будет равен `0` — +> это by design. Чтобы увидеть `backoff_ms > 0`, нужны статусы +> `UNAVAILABLE` (FastBackoff, slot 5 ms) или `OVERLOADED` / +> `CLIENT_RESOURCE_EXHAUSTED` (SlowBackoff, slot 1 s). Самый простой способ +> их получить — кратковременно перезапустить YDB во время работы примера: +> +> ```bash +> ./examples/otel_tracing/otel_tracing_example --retry-workers 8 --retry-ops 100 & +> sleep 5 +> docker compose -f examples/otel_tracing/docker-compose.yml restart ydb +> wait +> ``` + ### 4. Открыть дашборды | Сервис | URL | Описание | @@ -87,18 +147,36 @@ cmake --build . --target otel_tracing_example -j$(nproc) - **Recent Traces** — таблица трейсов из Jaeger #### В Jaeger UI: -- Выберите сервис `ydb-cpp-sdk-demo` -- Трейсы от QueryService содержат спаны с атрибутами: - - `db.system.name` = `other_sql` - - `server.address`, `server.port` - - `db.query.text` (для ExecuteQuery) - - `db.response.status_code`, `error.type` (при ошибках) -- TableService пока генерирует только метрики (без спанов) +- Выберите сервис `ydb-cpp-sdk-demo`. +- RPC-спаны (`SpanKind = CLIENT`): + `ydb.CreateSession`, `ydb.ExecuteQuery`, `ydb.ExecuteDataQuery`, + `ydb.BeginTransaction`, `ydb.Commit`, `ydb.Rollback`, + `ydb.ExecuteSchemeQuery`, `ydb.BulkUpsert`. +- Retry-спаны (`SpanKind = INTERNAL`): `ydb.RunWithRetry`, + `ydb.Try` (по одному на каждую попытку, с атрибутами + `ydb.retry.attempt`, `ydb.retry.backoff_ms`). +- Общие атрибуты на всех YDB-спанах: + - `db.system.name = ydb` + - `db.namespace` (имя базы) + - `server.address`, `server.port` (эндпоинт балансера) + - `network.peer.address`, `network.peer.port` (фактический узел кластера) +- На ошибках добавляются: + - `db.response.status_code` — строковый статус YDB (например, `ABORTED`) + - `error.type` — тот же строковый статус + - событие `exception` с `exception.type` и `exception.message` #### В Prometheus: -- `db_client_operation_requests_total` — счётчик запросов по операциям -- `db_client_operation_errors_total` — счётчик ошибок по операциям -- `db_client_operation_duration_seconds_bucket` — гистограмма длительности (OTel Semantic Conventions) +- `db_client_operation_duration_seconds_bucket` — гистограмма длительности + (OTel Semantic Conventions). Лейблы: `db.system.name`, `db.namespace`, + `db.operation.name` (с префиксом `ydb.`), `ydb.client.api` + (`Query` / `Table`). Для ошибок добавляются `db.response.status_code` + и `error.type`. +- `db_client_operation_requests_total` — счётчик начатых операций + (включая каждую попытку ретрая). +- `db_client_operation_errors_total` — счётчик неуспешных попыток. + Полезно сравнивать с `requests_total`: для retry-нагрузки на той же + «горячей» строке коэффициент ошибок будет очень высоким — это и есть + индикатор работы ретраев. ### 6. Остановить diff --git a/examples/otel_tracing/main.cpp b/examples/otel_tracing/main.cpp index 80e31a9f2e8..b06c224c2ca 100644 --- a/examples/otel_tracing/main.cpp +++ b/examples/otel_tracing/main.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -41,6 +42,8 @@ struct TConfig { std::string Database = "/local"; std::string OtlpEndpoint = "http://localhost:4328"; int Iterations = 20; + int RetryWorkers = 6; + int RetryOps = 30; }; nostd::shared_ptr InitTracing(const TConfig& cfg) { @@ -247,6 +250,113 @@ void RunTableWorkload(NTable::TTableClient& client, int iterations) { } } +void RunRetryWorkload(NQuery::TQueryClient& client, int workers, int opsPerWorker) { + std::cout << "\n=== Retry workload (SERIALIZABLE conflicts) ===" + << " workers=" << workers << " ops=" << opsPerWorker << std::endl; + + auto tracer = GetAppTracer(); + + { + auto seedSpan = tracer->StartSpan("RetryWorkload.Seed"); + auto scope = opentelemetry::trace::Scope(seedSpan); + + ThrowOnError(client.RetryQuerySync([](NQuery::TSession session) { + return session.ExecuteQuery(R"( + UPSERT INTO otel_demo (id, value) VALUES (9999u, "seed") + )", NQuery::TTxControl::BeginTx(NQuery::TTxSettings::SerializableRW()).CommitTx()).GetValueSync(); + })); + } + + std::atomic conflicts{0}; + std::atomic successes{0}; + std::vector threads; + threads.reserve(workers); + + for (int w = 0; w < workers; ++w) { + threads.emplace_back([&, w]() { + auto workerTracer = GetAppTracer(); + for (int i = 0; i < opsPerWorker; ++i) { + auto iterSpan = workerTracer->StartSpan("RetryWorkload.Op"); + auto scope = opentelemetry::trace::Scope(iterSpan); + iterSpan->SetAttribute("worker", static_cast(w)); + iterSpan->SetAttribute("op", static_cast(i)); + + auto status = client.RetryQuerySync( + [w, i, &conflicts](NQuery::TQueryClient client) -> TStatus { + auto sessionRes = client.GetSession().GetValueSync(); + if (!sessionRes.IsSuccess()) { + return sessionRes; + } + auto session = sessionRes.GetSession(); + + auto beginRes = session.BeginTransaction( + NQuery::TTxSettings::SerializableRW()).GetValueSync(); + if (!beginRes.IsSuccess()) { + return beginRes; + } + auto tx = beginRes.GetTransaction(); + + auto readRes = session.ExecuteQuery(R"( + SELECT value FROM otel_demo WHERE id = 9999u + )", NQuery::TTxControl::Tx(tx)).GetValueSync(); + if (!readRes.IsSuccess()) { + if (readRes.GetStatus() == EStatus::ABORTED) { + conflicts.fetch_add(1); + } + return readRes; + } + + std::this_thread::sleep_for( + std::chrono::milliseconds(5 + (w * 7 + i * 3) % 20)); + + auto params = TParamsBuilder() + .AddParam("$v").Utf8("w" + std::to_string(w) + + "_i" + std::to_string(i)).Build() + .Build(); + + auto writeRes = session.ExecuteQuery(R"( + DECLARE $v AS Utf8; + UPSERT INTO otel_demo (id, value) VALUES (9999u, $v) + )", NQuery::TTxControl::Tx(tx), params).GetValueSync(); + if (!writeRes.IsSuccess()) { + if (writeRes.GetStatus() == EStatus::ABORTED) { + conflicts.fetch_add(1); + } + return writeRes; + } + + auto commitRes = tx.Commit().GetValueSync(); + if (!commitRes.IsSuccess() + && commitRes.GetStatus() == EStatus::ABORTED) { + conflicts.fetch_add(1); + } + return commitRes; + }); + + if (status.IsSuccess()) { + successes.fetch_add(1); + iterSpan->SetStatus(opentelemetry::trace::StatusCode::kOk); + } else { + iterSpan->SetStatus(opentelemetry::trace::StatusCode::kError, + std::string(ToString(status.GetStatus()))); + std::cerr << " [retry-wl] worker=" << w << " op=" << i + << " final_status=" << static_cast(status.GetStatus()) + << std::endl; + } + } + }); + } + + for (auto& t : threads) { + t.join(); + } + + std::cout << " Retry workload done." + << " successes=" << successes.load() + << " observed_aborts=" << conflicts.load() + << " (each abort triggers one SDK retry attempt)" << std::endl; +} + int main(int argc, char** argv) { TConfig cfg; @@ -259,6 +369,10 @@ int main(int argc, char** argv) { .DefaultValue(cfg.OtlpEndpoint).StoreResult(&cfg.OtlpEndpoint); opts.AddLongOption('n', "iterations", "Number of iterations") .DefaultValue(std::to_string(cfg.Iterations)).StoreResult(&cfg.Iterations); + opts.AddLongOption("retry-workers", "Concurrent workers for retry workload (0 to skip)") + .DefaultValue(std::to_string(cfg.RetryWorkers)).StoreResult(&cfg.RetryWorkers); + opts.AddLongOption("retry-ops", "Operations per retry worker") + .DefaultValue(std::to_string(cfg.RetryOps)).StoreResult(&cfg.RetryOps); NLastGetopt::TOptsParseResult parsedOpts(&opts, argc, argv); @@ -288,6 +402,10 @@ int main(int argc, char** argv) { RunQueryWorkload(queryClient, cfg.Iterations); RunTableWorkload(tableClient, cfg.Iterations); + if (cfg.RetryWorkers > 0 && cfg.RetryOps > 0) { + RunRetryWorkload(queryClient, cfg.RetryWorkers, cfg.RetryOps); + } + std::cout << "\n=== Cleanup ===" << std::endl; ThrowOnError(queryClient.RetryQuerySync([](NQuery::TSession session) { return session.ExecuteQuery( diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index fec3aab583b..bb58495d9bc 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -127,11 +127,11 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto session = client.GetSession().ExtractValueSync(); ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); - auto requests = GetCounter(registry, "db.client.operation.requests", "CreateSession"); + auto requests = GetCounter(registry, "db.client.operation.requests", "GetSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); - auto duration = GetDuration(registry, "CreateSession", EStatus::SUCCESS); + auto duration = GetDuration(registry, "GetSession", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; EXPECT_GE(duration->Count(), 1u); From a1a9c01b6ce92c251b24199348546a06c2032f18 Mon Sep 17 00:00:00 2001 From: Bulat Date: Thu, 12 Feb 2026 13:39:18 +0000 Subject: [PATCH 13/93] [C++ SDK] Supported TCP_NODELAY option for grpc sockets (#32631) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/driver/driver.h | 7 ++ src/client/driver/driver.cpp | 8 ++ .../impl/internal/common/client_pid.cpp | 2 - .../grpc_connections/grpc_connections.cpp | 3 +- .../grpc_connections/grpc_connections.h | 1 + .../impl/internal/grpc_connections/params.h | 1 + src/library/grpc/client/grpc_client_low.cpp | 87 +++++++++++-------- src/library/grpc/client/grpc_client_low.h | 11 +-- .../grpc_client/grpc_client_low_ut.cpp | 2 +- 10 files changed, 76 insertions(+), 48 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e06b8762c3b..f0ff6ba466f 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -7f34dd9500921f4c8501e75e62488659cb276fb4 +f64a82eff2c4bbe405178f8e71acada58dbc3c39 diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index 20fa52d5e60..b5a4df7faf2 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -105,6 +105,13 @@ class TDriverConfig { //! default: true, 30, 5, 10 for linux, and true and OS default for others POSIX TDriverConfig& SetTcpKeepAliveSettings(bool enable, size_t idle, size_t count, size_t interval); + //! Set TCP_NODELAY socket option + //! enable - if true TCP_NODELAY is enabled (default, no Nagle algorithm, low latency, packet fragmentation) + //! - if false TCP_NODELAY is disabled (Nagle algorithm enabled, reduced packet fragmentation) + //! NOTE: This affects network performance. Disable only if you want to reduce packet fragmentation. + //! default: true + TDriverConfig& SetTcpNoDelay(bool enable); + //! Enable or disable drain of client logic (e.g. session pool drain) during dtor call TDriverConfig& SetDrainOnDtors(bool allowed); diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index c0ef98756fe..18ac064e480 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -40,6 +40,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { EDiscoveryMode GetDiscoveryMode() const override { return DiscoveryMode; } size_t GetMaxQueuedRequests() const override { return MaxQueuedRequests; } TTcpKeepAliveSettings GetTcpKeepAliveSettings() const override { return TcpKeepAliveSettings; } + bool GetTcpNoDelay() const override { return TcpNoDelay; } bool GetDrinOnDtors() const override { return DrainOnDtors; } TBalancingPolicy::TImpl GetBalancingSettings() const override { return BalancingSettings; } TDuration GetGRpcKeepAliveTimeout() const override { return GRpcKeepAliveTimeout; } @@ -71,6 +72,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { TCP_KEEPALIVE_COUNT, TCP_KEEPALIVE_INTERVAL }; + bool TcpNoDelay = true; bool DrainOnDtors = true; TBalancingPolicy::TImpl BalancingSettings = TBalancingPolicy::TImpl::UsePreferableLocation(std::nullopt); TDuration GRpcKeepAliveTimeout = TDuration::Seconds(10); @@ -174,6 +176,11 @@ TDriverConfig& TDriverConfig::SetTcpKeepAliveSettings(bool enable, size_t idle, return *this; } +TDriverConfig& TDriverConfig::SetTcpNoDelay(bool enable) { + Impl_->TcpNoDelay = enable; + return *this; +} + TDriverConfig& TDriverConfig::SetGrpcMemoryQuota(uint64_t bytes) { Impl_->MemoryQuota = bytes; return *this; @@ -285,6 +292,7 @@ TDriverConfig TDriver::GetConfig() const { Impl_->TcpKeepAliveSettings_.Count, Impl_->TcpKeepAliveSettings_.Interval ); + config.SetTcpNoDelay(Impl_->TcpNoDelay_); config.SetDrainOnDtors(Impl_->DrainOnDtors_); config.SetBalancingPolicy(std::make_unique(Impl_->BalancingSettings_)); config.SetGRpcKeepAliveTimeout(std::chrono::duration_cast(Impl_->GRpcKeepAliveTimeout_)); diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index 4bc2136a99b..cb5cc8eb0a5 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,8 +3,6 @@ #include -#include - #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index 757d5b777b7..d17cc078195 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -163,9 +163,10 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MaxMessageSize_(params->GetMaxMessageSize()) , QueuedRequests_(0) , TcpKeepAliveSettings_(params->GetTcpKeepAliveSettings()) + , TcpNoDelay_(params->GetTcpNoDelay()) , SocketIdleTimeout_(TDeadline::SafeDurationCast(params->GetSocketIdleTimeout())) #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL - , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout()) + , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout(), TcpNoDelay_) #endif , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 84a25162912..945cb97bcd4 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -713,6 +713,7 @@ class TGRpcConnectionsImpl std::atomic_int64_t QueuedRequests_; const NYdbGrpc::TTcpKeepAliveSettings TcpKeepAliveSettings_; + const bool TcpNoDelay_; const TDeadline::Duration SocketIdleTimeout_; #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL NYdbGrpc::TChannelPool ChannelPool_; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 1e827d3343b..35c0186421d 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -30,6 +30,7 @@ class IConnectionsParams { virtual EDiscoveryMode GetDiscoveryMode() const = 0; virtual size_t GetMaxQueuedRequests() const = 0; virtual NYdbGrpc::TTcpKeepAliveSettings GetTcpKeepAliveSettings() const = 0; + virtual bool GetTcpNoDelay() const = 0; virtual bool GetDrinOnDtors() const = 0; virtual TBalancingPolicy::TImpl GetBalancingSettings() const = 0; virtual TDuration GetGRpcKeepAliveTimeout() const = 0; diff --git a/src/library/grpc/client/grpc_client_low.cpp b/src/library/grpc/client/grpc_client_low.cpp index a8fa3c820cc..a54553e2110 100644 --- a/src/library/grpc/client/grpc_client_low.cpp +++ b/src/library/grpc/client/grpc_client_low.cpp @@ -37,18 +37,20 @@ void EnableGRpcTracing() { } #if !defined(YDB_DISABLE_GRPC_SOCKET_MUTATOR) -class TGRpcKeepAliveSocketMutator : public grpc_socket_mutator { +class TGRpcSocketMutator : public grpc_socket_mutator { public: - TGRpcKeepAliveSocketMutator(int idle, int count, int interval) - : Idle_(idle) + TGRpcSocketMutator(bool isKeepAliveEnabled, int idle, int count, int interval, bool tcpNoDelay) + : IsKeepAliveEnabled_(isKeepAliveEnabled) + , Idle_(idle) , Count_(count) , Interval_(interval) + , TcpNoDelay_(tcpNoDelay) { grpc_socket_mutator_init(this, &VTable); } private: - static TGRpcKeepAliveSocketMutator* Cast(grpc_socket_mutator* mutator) { - return static_cast(mutator); + static TGRpcSocketMutator* Cast(grpc_socket_mutator* mutator) { + return static_cast(mutator); } template @@ -56,24 +58,32 @@ class TGRpcKeepAliveSocketMutator : public grpc_socket_mutator { return setsockopt(fd, level, optname, reinterpret_cast(&value), sizeof(value)) == 0; } bool SetOption(int fd) { - if (!SetOption(fd, SOL_SOCKET, SO_KEEPALIVE, 1)) { - std::cerr << std::format("Failed to set SO_KEEPALIVE option: {}", strerror(errno)) << std::endl; - return false; - } + if (IsKeepAliveEnabled_) { + if (!SetOption(fd, SOL_SOCKET, SO_KEEPALIVE, 1)) { + std::cerr << std::format("Failed to set SO_KEEPALIVE option: {}", strerror(errno)) << std::endl; + return false; + } #ifdef _linux_ - if (Idle_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPIDLE, Idle_)) { - std::cerr << std::format("Failed to set TCP_KEEPIDLE option: {}", strerror(errno)) << std::endl; - return false; - } - if (Count_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPCNT, Count_)) { - std::cerr << std::format("Failed to set TCP_KEEPCNT option: {}", strerror(errno)) << std::endl; - return false; + if (Idle_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPIDLE, Idle_)) { + std::cerr << std::format("Failed to set TCP_KEEPIDLE option: {}", strerror(errno)) << std::endl; + return false; + } + if (Count_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPCNT, Count_)) { + std::cerr << std::format("Failed to set TCP_KEEPCNT option: {}", strerror(errno)) << std::endl; + return false; + } + if (Interval_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPINTVL, Interval_)) { + std::cerr << std::format("Failed to set TCP_KEEPINTVL option: {}", strerror(errno)) << std::endl; + return false; + } +#endif } - if (Interval_ && !SetOption(fd, IPPROTO_TCP, TCP_KEEPINTVL, Interval_)) { - std::cerr << std::format("Failed to set TCP_KEEPINTVL option: {}", strerror(errno)) << std::endl; + + if (!SetOption(fd, IPPROTO_TCP, TCP_NODELAY, static_cast(TcpNoDelay_))) { + std::cerr << std::format("Failed to set TCP_NODELAY option: {}", strerror(errno)) << std::endl; return false; } -#endif + return true; } static bool Mutate(int fd, grpc_socket_mutator* mutator) { @@ -83,8 +93,8 @@ class TGRpcKeepAliveSocketMutator : public grpc_socket_mutator { static int Compare(grpc_socket_mutator* a, grpc_socket_mutator* b) { const auto* selfA = Cast(a); const auto* selfB = Cast(b); - auto tupleA = std::make_tuple(selfA->Idle_, selfA->Count_, selfA->Interval_); - auto tupleB = std::make_tuple(selfB->Idle_, selfB->Count_, selfB->Interval_); + auto tupleA = std::make_tuple(selfA->IsKeepAliveEnabled_, selfA->Idle_, selfA->Count_, selfA->Interval_, selfA->TcpNoDelay_); + auto tupleB = std::make_tuple(selfB->IsKeepAliveEnabled_, selfB->Idle_, selfB->Count_, selfB->Interval_, selfB->TcpNoDelay_); return tupleA < tupleB ? -1 : tupleA > tupleB ? 1 : 0; } static void Destroy(grpc_socket_mutator* mutator) { @@ -96,17 +106,19 @@ class TGRpcKeepAliveSocketMutator : public grpc_socket_mutator { } static grpc_socket_mutator_vtable VTable; + const bool IsKeepAliveEnabled_; const int Idle_; const int Count_; const int Interval_; + const bool TcpNoDelay_; }; -grpc_socket_mutator_vtable TGRpcKeepAliveSocketMutator::VTable = +grpc_socket_mutator_vtable TGRpcSocketMutator::VTable = { - &TGRpcKeepAliveSocketMutator::Mutate, - &TGRpcKeepAliveSocketMutator::Compare, - &TGRpcKeepAliveSocketMutator::Destroy, - &TGRpcKeepAliveSocketMutator::Mutate2 + &TGRpcSocketMutator::Mutate, + &TGRpcSocketMutator::Compare, + &TGRpcSocketMutator::Destroy, + &TGRpcSocketMutator::Mutate2 }; #endif @@ -133,8 +145,9 @@ void TGRpcRequestProcessorCommon::GetInitialMetadata(std::unordered_multimapsecond); LastUsedQueue_.emplace(Pool_.at(channelId).GetLastUseTime(), channelId); @@ -611,17 +624,15 @@ void TGRpcClientLow::ForgetContext(TContextImpl* context) { } } -grpc_socket_mutator* NImpl::CreateGRpcKeepAliveSocketMutator(const TTcpKeepAliveSettings& TcpKeepAliveSettings_) { +grpc_socket_mutator* NImpl::CreateGRpcSocketMutator(const TTcpKeepAliveSettings& TcpKeepAliveSettings_, bool tcpNoDelay) { #if !defined(YDB_DISABLE_GRPC_SOCKET_MUTATOR) - TGRpcKeepAliveSocketMutator* mutator = nullptr; - if (TcpKeepAliveSettings_.Enabled) { - mutator = new TGRpcKeepAliveSocketMutator( - TcpKeepAliveSettings_.Idle, - TcpKeepAliveSettings_.Count, - TcpKeepAliveSettings_.Interval - ); - } - return mutator; + return new TGRpcSocketMutator( + TcpKeepAliveSettings_.Enabled, + TcpKeepAliveSettings_.Idle, + TcpKeepAliveSettings_.Count, + TcpKeepAliveSettings_.Interval, + tcpNoDelay + ); #endif return nullptr; } diff --git a/src/library/grpc/client/grpc_client_low.h b/src/library/grpc/client/grpc_client_low.h index ac865d29b00..fb0ae5b55af 100644 --- a/src/library/grpc/client/grpc_client_low.h +++ b/src/library/grpc/client/grpc_client_low.h @@ -425,10 +425,10 @@ class IStreamRequestReadWriteProcessor : public IStreamRequestReadProcessor cb); @@ -511,6 +511,7 @@ class TChannelPool { std::unordered_map Pool_; std::multimap LastUsedQueue_; [[maybe_unused]] TTcpKeepAliveSettings TcpKeepAliveSettings_; + [[maybe_unused]] bool TcpNoDelay_; TDuration ExpireTime_; TDuration UpdateReUseTime_; void EraseFromQueueByTime(const TInstant& lastUseTime, const std::string& channelId); @@ -1367,8 +1368,8 @@ class TGRpcClientLow } template - std::unique_ptr> CreateGRpcServiceConnection(const TGRpcClientConfig& config, const TTcpKeepAliveSettings& keepAlive) { - auto mutator = NImpl::CreateGRpcKeepAliveSocketMutator(keepAlive); + std::unique_ptr> CreateGRpcServiceConnection(const TGRpcClientConfig& config, const TTcpKeepAliveSettings& keepAlive, bool tcpNoDelay = true) { + auto mutator = NImpl::CreateGRpcSocketMutator(keepAlive, tcpNoDelay); // will be destroyed inside grpc return std::unique_ptr>(new TServiceConnection(CreateChannelInterface(config, mutator), this)); } diff --git a/tests/unit/library/grpc_client/grpc_client_low_ut.cpp b/tests/unit/library/grpc_client/grpc_client_low_ut.cpp index 35a89328cb7..a53ee718334 100644 --- a/tests/unit/library/grpc_client/grpc_client_low_ut.cpp +++ b/tests/unit/library/grpc_client/grpc_client_low_ut.cpp @@ -22,7 +22,7 @@ Y_UNIT_TEST_SUITE(ChannelPoolTests) { 5, // NYdb::TCP_KEEPALIVE_COUNT, unused in UT, but is necessary in constructor 10 // NYdb::TCP_KEEPALIVE_INTERVAL, unused in UT, but is necessary in constructor }; - auto channelPool = TChannelPool(tcpKeepAliveSettings, TDuration::MilliSeconds(250)); + auto channelPool = TChannelPool(tcpKeepAliveSettings, TDuration::MilliSeconds(250), true); std::vector> ChannelInterfacesWeak; { From e35be7e0e1ce78acf0d103483ede7f949e6977a4 Mon Sep 17 00:00:00 2001 From: qyryq Date: Thu, 12 Feb 2026 13:39:24 +0000 Subject: [PATCH 14/93] Enable direct read without consumer (LOGBROKER-9364) (#32846) --- .github/last_commit.txt | 2 +- src/client/topic/impl/direct_reader.cpp | 10 +++++----- src/client/topic/ut/basic_usage_ut.cpp | 6 +----- src/client/topic/ut/direct_read_ut.cpp | 5 +++-- 4 files changed, 10 insertions(+), 13 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index f0ff6ba466f..c26aaaef007 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -f64a82eff2c4bbe405178f8e71acada58dbc3c39 +821f5e81167744aebf39d669f189a7d47c329c7a diff --git a/src/client/topic/impl/direct_reader.cpp b/src/client/topic/impl/direct_reader.cpp index 3c85ec02025..4be57ef89c9 100644 --- a/src/client/topic/impl/direct_reader.cpp +++ b/src/client/topic/impl/direct_reader.cpp @@ -517,23 +517,23 @@ void TDirectReadSession::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t c if (!IsErrorMessage(*ServerMessage)) { if (ServerMessage->server_message_case() != TDirectReadServerMessage::kDirectReadResponse) { - LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX subsession got message = " << ServerMessage->ShortDebugString()); + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "subsession got message = " << ServerMessage->ShortDebugString()); } else { const auto& data = ServerMessage->direct_read_response().partition_data(); const auto partitionSessionId = ServerMessage->direct_read_response().partition_session_id(); auto partitionSessionIt = PartitionSessions.find(partitionSessionId); if (partitionSessionIt == PartitionSessions.end()) { - LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX subsession got message = DirectReadResponse partitionSessionId=" << partitionSessionId << " not found"); + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "subsession got message = DirectReadResponse partitionSessionId=" << partitionSessionId << " not found"); } if (data.batches_size() == 0) { - LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX subsession got message = DirectReadResponse EMPTY"); + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "subsession got message = DirectReadResponse EMPTY"); } else { const auto& firstBatch = data.batches(0); const auto firstOffset = firstBatch.message_data(0).offset(); const auto& lastBatch = data.batches(data.batches_size() - 1); const auto lastOffset = lastBatch.message_data(lastBatch.message_data_size() - 1).offset(); auto partitionId = partitionSessionIt == PartitionSessions.end() ? -1 : partitionSessionIt->second.PartitionId; - LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX subsession got message = DirectReadResponse" + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "subsession got message = DirectReadResponse" << " partitionSessionId = " << partitionSessionId << " partitionId = " << partitionId << " directReadId = " << ServerMessage->direct_read_response().direct_read_id() @@ -788,7 +788,7 @@ void TDirectReadSession::WriteToProcessorImpl(TDirectReadClientMessage&& req) { Y_ABORT_UNLESS(Lock.IsLocked()); if (Processor) { - LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX subsession send message = " << req.ShortDebugString()); + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "subsession send message = " << req.ShortDebugString()); Processor->Write(std::move(req)); } } diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index 478a88927a4..c8589015a81 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -794,10 +794,6 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } Y_UNIT_TEST(ReadWithoutConsumerWithRestarts) { - if (EnableDirectRead) { - // TODO(qyryq) Enable the test when LOGBROKER-9364 is done. - return; - } TTopicSdkTestSetup setup(TEST_CASE_NAME); auto compressor = std::make_shared(); auto decompressor = CreateThreadPoolManagedExecutor(1); @@ -810,7 +806,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .MaxMemoryUsageBytes(1_MB) .DecompressionExecutor(decompressor) .AppendTopics(topic) - // .DirectRead(EnableDirectRead) + .DirectRead(EnableDirectRead) ; TWriteSessionSettings writeSettings; diff --git a/src/client/topic/ut/direct_read_ut.cpp b/src/client/topic/ut/direct_read_ut.cpp index ea7d3227263..29e0e6e6955 100644 --- a/src/client/topic/ut/direct_read_ut.cpp +++ b/src/client/topic/ut/direct_read_ut.cpp @@ -47,7 +47,7 @@ Y_UNIT_TEST_SUITE(DirectReadWithServer) { auto readerSettings = TReadSessionSettings() .ConsumerName(setup.GetConsumerName()) .AppendTopics(setup.GetTopicPath()) - // .DirectRead(true) + .DirectRead(true) ; TIntrusivePtr partitionSession; @@ -129,7 +129,7 @@ Y_UNIT_TEST_SUITE(DirectReadWithServer) { auto readerSettings = TReadSessionSettings() .ConsumerName(setup.GetConsumerName()) .AppendTopics(setup.GetTopicPath()) - // .DirectRead(true) + .DirectRead(true) ; TIntrusivePtr partitionSession; @@ -175,6 +175,7 @@ Y_UNIT_TEST_SUITE(DirectReadWithServer) { reader->Close(); } + } // Y_UNIT_TEST_SUITE_F(DirectReadWithServer) } // namespace NYdb::NTopic::NTests From e455288866180299890de095d0bc179c0d0c6d14 Mon Sep 17 00:00:00 2001 From: qyryq Date: Thu, 12 Feb 2026 13:39:30 +0000 Subject: [PATCH 15/93] Add TSimpleBlockingFederatedWriteSession (#32794) --- .github/last_commit.txt | 2 +- .../client/federated_topic/federated_topic.h | 2 +- .../federated_topic/impl/federated_topic.cpp | 8 +- .../impl/federated_topic_impl.cpp | 23 +- .../impl/federated_write_session.cpp | 86 ++++ .../impl/federated_write_session.h | 41 +- .../federated_topic/ut/fds_mock/fds_mock.h | 43 ++ .../ut/simple_blocking_write_session_ut.cpp | 402 ++++++++++++++++++ .../topic/common/simple_blocking_helpers.h | 48 +++ src/client/topic/impl/write_session.cpp | 30 +- 10 files changed, 642 insertions(+), 43 deletions(-) create mode 100644 src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp create mode 100644 src/client/topic/common/simple_blocking_helpers.h diff --git a/.github/last_commit.txt b/.github/last_commit.txt index c26aaaef007..e9c0845fd70 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -821f5e81167744aebf39d669f189a7d47c329c7a +bfb70b6f3508f8ded0feb4dfb7ff92709975bf02 diff --git a/include/ydb-cpp-sdk/client/federated_topic/federated_topic.h b/include/ydb-cpp-sdk/client/federated_topic/federated_topic.h index 52af62a9aad..1e5b8fffe78 100644 --- a/include/ydb-cpp-sdk/client/federated_topic/federated_topic.h +++ b/include/ydb-cpp-sdk/client/federated_topic/federated_topic.h @@ -517,7 +517,7 @@ class TFederatedTopicClient { std::shared_ptr CreateReadSession(const TFederatedReadSessionSettings& settings); //! Create write session. - // std::shared_ptr CreateSimpleBlockingWriteSession(const TFederatedWriteSessionSettings& settings); + std::shared_ptr CreateSimpleBlockingWriteSession(const TFederatedWriteSessionSettings& settings); std::shared_ptr CreateWriteSession(const TFederatedWriteSessionSettings& settings); struct TClusterInfo { diff --git a/src/client/federated_topic/impl/federated_topic.cpp b/src/client/federated_topic/impl/federated_topic.cpp index 5b0f7f2b19b..33f39ec23a1 100644 --- a/src/client/federated_topic/impl/federated_topic.cpp +++ b/src/client/federated_topic/impl/federated_topic.cpp @@ -73,10 +73,10 @@ std::shared_ptr TFederatedTopicClient::CreateReadSession( return Impl_->CreateReadSession(settings); } -// std::shared_ptr TFederatedTopicClient::CreateSimpleBlockingWriteSession( -// const TFederatedWriteSessionSettings& settings) { -// return Impl_->CreateSimpleBlockingWriteSession(settings); -// } +std::shared_ptr TFederatedTopicClient::CreateSimpleBlockingWriteSession( + const TFederatedWriteSessionSettings& settings) { + return Impl_->CreateSimpleBlockingWriteSession(settings); +} std::shared_ptr TFederatedTopicClient::CreateWriteSession(const TFederatedWriteSessionSettings& settings) { return Impl_->CreateWriteSession(settings); diff --git a/src/client/federated_topic/impl/federated_topic_impl.cpp b/src/client/federated_topic/impl/federated_topic_impl.cpp index 45fc189643c..4da1c013a83 100644 --- a/src/client/federated_topic/impl/federated_topic_impl.cpp +++ b/src/client/federated_topic/impl/federated_topic_impl.cpp @@ -13,14 +13,23 @@ TFederatedTopicClient::TImpl::CreateReadSession(const TFederatedReadSessionSetti return std::move(session); } -// std::shared_ptr -// TFederatedTopicClient::TImpl::CreateSimpleBlockingWriteSession(const TFederatedWriteSessionSettings& settings) { -// InitObserver(); -// auto session = std::make_shared(settings, Connections, ClientSettings, GetObserver()); -// session->Start(); -// return std::move(session); +std::shared_ptr +TFederatedTopicClient::TImpl::CreateSimpleBlockingWriteSession(const TFederatedWriteSessionSettings& settings) { + // Split settings.MaxMemoryUsage_ by two. + // One half goes to subsession. Other half goes to federated session internal buffer. + const ui64 splitSize = (settings.MaxMemoryUsage_ + 1) / 2; + TFederatedWriteSessionSettings splitSettings = settings; + splitSettings.MaxMemoryUsage(splitSize); + InitObserver(); -// } + with_lock(Lock) { + if (!splitSettings.EventHandlers_.HandlersExecutor_) { + splitSettings.EventHandlers_.HandlersExecutor(ClientSettings.DefaultHandlersExecutor_); + } + } + return std::make_shared( + splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs, GetSubsessionHandlersExecutor()); +} std::shared_ptr TFederatedTopicClient::TImpl::CreateWriteSession(const TFederatedWriteSessionSettings& settings) { diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index 8f559f6076a..a6063ea208e 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -1,6 +1,7 @@ #include "federated_write_session.h" #include +#include #include #define INCLUDE_YDB_INTERNAL_H @@ -467,4 +468,89 @@ bool TFederatedWriteSessionImpl::Close(TDuration timeout) { } } +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TSimpleBlockingFederatedWriteSession + +TSimpleBlockingFederatedWriteSession::TSimpleBlockingFederatedWriteSession( + const TFederatedWriteSessionSettings& settings, + std::shared_ptr connections, + const TFederatedTopicClientSettings& clientSettings, + std::shared_ptr observer, + std::shared_ptr>> codecs, + IExecutor::TPtr subsessionHandlersExecutor +) { + TFederatedWriteSessionSettings subSettings = settings; + auto& log = connections->GetLog(); + if (settings.EventHandlers_.AcksHandler_) { + LOG_LAZY(log, TLOG_WARNING, "TSimpleBlockingFederatedWriteSession: Cannot use AcksHandler, resetting."); + subSettings.EventHandlers_.AcksHandler({}); + } + if (settings.EventHandlers_.ReadyToAcceptHandler_) { + LOG_LAZY(log, TLOG_WARNING, "TSimpleBlockingFederatedWriteSession: Cannot use ReadyToAcceptHandler, resetting."); + subSettings.EventHandlers_.ReadyToAcceptHandler({}); + } + if (settings.EventHandlers_.SessionClosedHandler_) { + LOG_LAZY(log, TLOG_WARNING, "TSimpleBlockingFederatedWriteSession: Cannot use SessionClosedHandler, resetting."); + subSettings.EventHandlers_.SessionClosedHandler({}); + } + if (settings.EventHandlers_.CommonHandler_) { + LOG_LAZY(log, TLOG_WARNING, "TSimpleBlockingFederatedWriteSession: Cannot use CommonHandler, resetting."); + subSettings.EventHandlers_.CommonHandler({}); + } + + Writer = std::make_shared( + subSettings, std::move(connections), clientSettings, std::move(observer), std::move(codecs), std::move(subsessionHandlersExecutor)); + Writer->Start(); +} + +uint64_t TSimpleBlockingFederatedWriteSession::GetInitSeqNo() { + return Writer->GetInitSeqNo().GetValueSync(); +} + +bool TSimpleBlockingFederatedWriteSession::Write( + std::string_view data, std::optional seqNo, std::optional createTimestamp, const TDuration& blockTimeout +) { + auto message = NTopic::TWriteMessage(std::move(data)) + .SeqNo(seqNo) + .CreateTimestamp(createTimestamp); + return Write(std::move(message), nullptr, blockTimeout); +} + +bool TSimpleBlockingFederatedWriteSession::Write( + NTopic::TWriteMessage&& message, TTransactionBase* tx, const TDuration& blockTimeout +) { + if (tx || message.GetTxPtr()) { + ythrow yexception() << "transactions are not supported"; + } + auto continuationToken = WaitForToken(blockTimeout); + if (continuationToken.has_value()) { + Writer->Write(std::move(*continuationToken), std::move(message)); + return true; + } + return false; +} + +std::optional TSimpleBlockingFederatedWriteSession::WaitForToken(const TDuration& timeout) { + return NTopic::NDetail::WaitForToken(*Writer, Closed, timeout); +} + +NTopic::TWriterCounters::TPtr TSimpleBlockingFederatedWriteSession::GetCounters() { + ythrow yexception() << "GetCounters is not yet implemented for federated write sessions"; +} + +bool TSimpleBlockingFederatedWriteSession::IsAlive() const { + return !Closed.load(); +} + +bool TSimpleBlockingFederatedWriteSession::Close(TDuration closeTimeout) { + Closed.store(true); + return Writer->Close(closeTimeout); +} + +TSimpleBlockingFederatedWriteSession::~TSimpleBlockingFederatedWriteSession() { + if (!Closed.load()) { + Close(TDuration::Zero()); + } +} + } // namespace NYdb::NFederatedTopic diff --git a/src/client/federated_topic/impl/federated_write_session.h b/src/client/federated_topic/impl/federated_write_session.h index 77376740e90..8e8e3e2dd38 100644 --- a/src/client/federated_topic/impl/federated_write_session.h +++ b/src/client/federated_topic/impl/federated_write_session.h @@ -150,6 +150,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, class TFederatedWriteSession : public NTopic::IWriteSession, public NTopic::TContextOwner { friend class TFederatedTopicClient::TImpl; + friend class TSimpleBlockingFederatedWriteSession; public: @@ -174,13 +175,13 @@ class TFederatedWriteSession : public NTopic::IWriteSession, return TryGetImpl()->GetInitSeqNo(); } void Write(NTopic::TContinuationToken&& continuationToken, NTopic::TWriteMessage&& message, TTransactionBase* tx = nullptr) override { - if (tx) { + if (tx || message.GetTxPtr()) { ythrow yexception() << "transactions are not supported"; } TryGetImpl()->Write(std::move(continuationToken), std::move(message)); } void WriteEncoded(NTopic::TContinuationToken&& continuationToken, NTopic::TWriteMessage&& params, TTransactionBase* tx = nullptr) override { - if (tx) { + if (tx || params.GetTxPtr()) { ythrow yexception() << "transactions are not supported"; } TryGetImpl()->WriteEncoded(std::move(continuationToken), std::move(params)); @@ -206,4 +207,40 @@ class TFederatedWriteSession : public NTopic::IWriteSession, } }; +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TSimpleBlockingFederatedWriteSession + +class TSimpleBlockingFederatedWriteSession : public NTopic::ISimpleBlockingWriteSession { +public: + TSimpleBlockingFederatedWriteSession( + const TFederatedWriteSessionSettings& settings, + std::shared_ptr connections, + const TFederatedTopicClientSettings& clientSettings, + std::shared_ptr observer, + std::shared_ptr>> codecs, + IExecutor::TPtr subsessionHandlersExecutor); + + bool Write(std::string_view data, std::optional seqNo = std::nullopt, std::optional createTimestamp = std::nullopt, + const TDuration& blockTimeout = TDuration::Max()) override; + + bool Write(NTopic::TWriteMessage&& message, + TTransactionBase* tx = nullptr, + const TDuration& blockTimeout = TDuration::Max()) override; + + uint64_t GetInitSeqNo() override; + + bool Close(TDuration closeTimeout = TDuration::Max()) override; + + ~TSimpleBlockingFederatedWriteSession(); + bool IsAlive() const override; + + NTopic::TWriterCounters::TPtr GetCounters() override; + +private: + std::optional WaitForToken(const TDuration& timeout); + + std::shared_ptr Writer; + std::atomic_bool Closed = false; +}; + } // namespace NYdb::NFederatedTopic diff --git a/src/client/federated_topic/ut/fds_mock/fds_mock.h b/src/client/federated_topic/ut/fds_mock/fds_mock.h index 9f9df0ceb2b..dc0e787f1bd 100644 --- a/src/client/federated_topic/ut/fds_mock/fds_mock.h +++ b/src/client/federated_topic/ut/fds_mock/fds_mock.h @@ -57,11 +57,27 @@ class TFederationDiscoveryServiceMock: public Ydb::FederationDiscovery::V1::Fede } while (true); } + void SetAutoRespondSingleDatabase(bool enable) { + with_lock(Lock) { + AutoRespondSingleDatabase = enable; + } + } + virtual grpc::Status ListFederationDatabases(grpc::ServerContext*, const TRequest* request, TResponse* response) override { Y_UNUSED(request); + // Check if auto-respond mode is enabled + with_lock(Lock) { + if (AutoRespondSingleDatabase) { + auto result = ComposeOkResultSingleDatabase(); + Cerr << ">>> Auto-responding with single database" << Endl; + *response = std::move(result.Response); + return result.Status; + } + } + auto p = NThreading::NewPromise(); auto f = p.GetFuture(); @@ -125,6 +141,32 @@ class TFederationDiscoveryServiceMock: public Ydb::FederationDiscovery::V1::Fede return ComposeOkResult(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); } + // Returns a single database response to avoid partition competition between multiple sub-sessions + TGrpcResult ComposeOkResultSingleDatabase() { + Ydb::FederationDiscovery::ListFederationDatabasesResponse okResponse; + + auto op = okResponse.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + okResponse.mutable_operation()->set_ready(true); + okResponse.mutable_operation()->set_id("12345"); + + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("dc1"); + auto c1 = mockResult.add_federation_databases(); + c1->set_name("dc1"); + c1->set_path("/Root"); + c1->set_id("account-dc1"); + c1->set_endpoint("localhost:" + ToString(Port)); + c1->set_location("dc1"); + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c1->set_weight(1000); + + op->mutable_result()->PackFrom(mockResult); + + return {okResponse, grpc::Status::OK}; + } + TGrpcResult ComposeOkResultUnavailableDatabases() { return ComposeOkResult(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_UNAVAILABLE); } @@ -173,6 +215,7 @@ class TFederationDiscoveryServiceMock: public Ydb::FederationDiscovery::V1::Fede ui16 Port; std::deque PendingRequests; TAdaptiveLock Lock; + bool AutoRespondSingleDatabase = false; }; } // namespace NYdb::NFederatedTopic::NTests diff --git a/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp b/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp new file mode 100644 index 00000000000..51260f74b12 --- /dev/null +++ b/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp @@ -0,0 +1,402 @@ +#include +#include + +#include +#include + +#include +#include + +#include + +namespace NYdb::NFederatedTopic::NTests { + +// Test fixture providing common setup for federated topic tests +class TSimpleBlockingWriteSessionTestFixture { +public: + explicit TSimpleBlockingWriteSessionTestFixture(const TString& testName) + : Setup(std::make_shared( + testName, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 1)) + , ThreadPool(CreateThreadPool(2)) + { + Setup->Start(true, true); + + FdsMock.Port = Setup->GetGrpcPort(); + ServicePort = Setup->GetPortManager()->GetPort(4285); + GrpcServer = Setup->StartGrpcService(ServicePort, &FdsMock); + + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << ServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(std::unique_ptr(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG).Release())); + + Driver = std::make_unique(cfg); + TopicClient = std::make_unique(*Driver); + } + + // Creates a write session and responds to the FDS request with available databases + // The session creation triggers FDS discovery, so we start creation async and then respond + std::shared_ptr CreateWriteSessionWithFdsResponse( + const TFederatedWriteSessionSettings& settings + ) { + // Start session creation asynchronously + auto sessionFuture = NThreading::Async([this, settings]() { + return TopicClient->CreateSimpleBlockingWriteSession(settings); + }, *ThreadPool); + + // Wait for FDS request and respond + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultAvailableDatabases()); + + // Wait for session creation to complete + return sessionFuture.GetValueSync(); + } + + std::shared_ptr CreateWriteSessionWithFdsResponse() { + return CreateWriteSessionWithFdsResponse(DefaultWriteSettings()); + } + + // Creates a write session and responds with one unavailable database + std::shared_ptr CreateWriteSessionWithUnavailableDatabaseResponse( + const TFederatedWriteSessionSettings& settings, + int unavailableDbIndex + ) { + auto sessionFuture = NThreading::Async([this, settings]() { + return TopicClient->CreateSimpleBlockingWriteSession(settings); + }, *ThreadPool); + + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultWithUnavailableDatabase(unavailableDbIndex)); + + return sessionFuture.GetValueSync(); + } + + TFederatedWriteSessionSettings DefaultWriteSettings() { + return TFederatedWriteSessionSettings() + .Path(Setup->GetTestTopicPath()) + .MessageGroupId("src_id"); + } + + std::shared_ptr CreateReadSession() { + TFederatedReadSessionSettings settings; + settings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(std::string(Setup->GetTestTopicPath())); + return TopicClient->CreateReadSession(settings); + } + + // Creates a read session and responds to FDS + std::shared_ptr CreateReadSessionWithFdsResponse() { + auto session = CreateReadSession(); + + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultAvailableDatabases()); + + return session; + } + + // Creates a write session with single-database FDS response (avoids partition competition) + std::shared_ptr CreateWriteSessionWithSingleDatabaseFdsResponse( + const TFederatedWriteSessionSettings& settings + ) { + auto sessionFuture = NThreading::Async([this, settings]() { + return TopicClient->CreateSimpleBlockingWriteSession(settings); + }, *ThreadPool); + + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultSingleDatabase()); + + return sessionFuture.GetValueSync(); + } + + std::shared_ptr CreateWriteSessionWithSingleDatabaseFdsResponse() { + return CreateWriteSessionWithSingleDatabaseFdsResponse(DefaultWriteSettings()); + } + + // Creates a read session with single-database FDS response (avoids partition competition) + std::shared_ptr CreateReadSessionWithSingleDatabaseFdsResponse() { + auto session = CreateReadSession(); + + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultSingleDatabase()); + + return session; + } + + // Creates a read session with unavailable database response + std::shared_ptr CreateReadSessionWithUnavailableDatabaseResponse(int unavailableDbIndex) { + auto session = CreateReadSession(); + + auto fdsRequest = FdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(FdsMock.ComposeOkResultWithUnavailableDatabase(unavailableDbIndex)); + + return session; + } + + void ConfirmPartitionStart(std::shared_ptr& session, TDuration timeout = TDuration::Seconds(60)) { + TInstant deadline = TInstant::Now() + timeout; + while (TInstant::Now() < deadline) { + session->WaitEvent().Wait(TDuration::MilliSeconds(100)); + auto event = session->GetEvent(false); + if (!event.has_value()) { + continue; + } + if (auto* startEvent = std::get_if(&*event)) { + startEvent->Confirm(); + return; + } + } + UNIT_FAIL("Timeout waiting for TStartPartitionSessionEvent"); + } + + std::optional ReadOneMessage( + std::shared_ptr& session, + TDuration timeout = TDuration::Seconds(60) + ) { + TInstant deadline = TInstant::Now() + timeout; + while (TInstant::Now() < deadline) { + session->WaitEvent().Wait(TDuration::MilliSeconds(100)); + auto events = session->GetEvents(false, 1); + for (auto& event : events) { + if (auto* dataEvent = std::get_if(&event)) { + auto& messages = dataEvent->GetMessages(); + if (!messages.empty()) { + auto msg = messages[0]; + dataEvent->Commit(); + return msg; + } + } + } + } + return std::nullopt; + } + + TString GetTestTopic() const { return Setup->GetTestTopic(); } + TFederationDiscoveryServiceMock& GetFdsMock() { return FdsMock; } + + // Helper to verify that messages were written correctly by reading them back + void VerifyMessagesWritten(const TVector& expectedMessages) { + // Create read session with single database FDS response + auto readSession = CreateReadSessionWithSingleDatabaseFdsResponse(); + ConfirmPartitionStart(readSession); + + // Enable auto-respond for any subsequent FDS discovery refresh requests + FdsMock.SetAutoRespondSingleDatabase(true); + + // Read and verify all messages + for (size_t i = 0; i < expectedMessages.size(); ++i) { + auto msg = ReadOneMessage(readSession); + UNIT_ASSERT_C(msg.has_value(), "Did not receive message " << i); + UNIT_ASSERT_VALUES_EQUAL_C(msg->GetData(), expectedMessages[i], + "Message " << i << " content mismatch"); + } + + readSession->Close(); + } + + TFederatedTopicClient* GetTopicClient() { return TopicClient.get(); } + IThreadPool* GetThreadPool() { return ThreadPool.Get(); } + +private: + std::shared_ptr Setup; + TFederationDiscoveryServiceMock FdsMock; + ui16 ServicePort; + std::unique_ptr GrpcServer; + std::unique_ptr Driver; + std::unique_ptr TopicClient; + THolder ThreadPool; +}; + +Y_UNIT_TEST_SUITE(SimpleBlockingFederatedWriteSession) { + + Y_UNIT_TEST(BasicWriteAndClose) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + UNIT_ASSERT(session->IsAlive()); + + TString message = "test message"; + UNIT_ASSERT(session->Write(message)); + session->Close(); + UNIT_ASSERT(!session->IsAlive()); + + // Verify the message was written correctly + fixture.VerifyMessagesWritten({message}); + } + + Y_UNIT_TEST(WriteMultipleMessages) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + // Write and verify a single message (multi-message read verification needs more infrastructure) + TString message = "message-0"; + UNIT_ASSERT_C(session->Write(message), "Failed to write message"); + + session->Close(); + + // Verify message was written correctly + fixture.VerifyMessagesWritten({message}); + } + + Y_UNIT_TEST(WriteWithSeqNoAndTimestamp) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + TString message = "message with explicit params"; + UNIT_ASSERT(session->Write(message, 1, TInstant::Now())); + session->Close(); + + // Verify the message was written correctly + fixture.VerifyMessagesWritten({message}); + } + + Y_UNIT_TEST(WriteAndReadBack) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + // Create write session first (which responds to first FDS with single database) + auto writeSession = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(writeSession); + + // Create read session (which triggers second FDS request with single database) + auto readSession = fixture.CreateReadSessionWithSingleDatabaseFdsResponse(); + fixture.ConfirmPartitionStart(readSession); + + // Enable auto-respond for any subsequent FDS discovery refresh requests + fixture.GetFdsMock().SetAutoRespondSingleDatabase(true); + + // Write single message for now to verify flow works + UNIT_ASSERT(writeSession->Write("test-message")); + + auto msg = fixture.ReadOneMessage(readSession); + UNIT_ASSERT_C(msg.has_value(), "Did not receive message"); + UNIT_ASSERT_VALUES_EQUAL(msg->GetData(), "test-message"); + + writeSession->Close(); + readSession->Close(); + } + + Y_UNIT_TEST(CloseEmptySession) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithFdsResponse(); + UNIT_ASSERT(session); + session->Close(); + } + + Y_UNIT_TEST(WriteWithPreferredDatabase) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + // Use single-database FDS to avoid multi-database complications + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + TString message = "message to preferred database"; + UNIT_ASSERT(session->Write(message)); + session->Close(); + + // Verify the message was written correctly + fixture.VerifyMessagesWritten({message}); + } + + Y_UNIT_TEST(WriteWithPreferredDatabaseUnavailableAndFallback) { + // TODO: Enable after fixing test environment to support multiple federated databases/topics + /* + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto settings = fixture.DefaultWriteSettings() + .PreferredDatabase("dc1") + .AllowFallback(true); + + // dc1 is unavailable, but fallback is allowed + auto writeSession = fixture.CreateWriteSessionWithUnavailableDatabaseResponse(settings, 1); + UNIT_ASSERT(writeSession); + + // Create read session to verify where message was written + auto readSession = fixture.CreateReadSessionWithUnavailableDatabaseResponse(1); + fixture.ConfirmPartitionStart(readSession); + + TString message = "message with fallback"; + UNIT_ASSERT(writeSession->Write(message)); + + // Verify message was written to fallback database (not dc1) + auto msg = fixture.ReadOneMessage(readSession); + UNIT_ASSERT(msg.has_value()); + UNIT_ASSERT_VALUES_EQUAL(msg->GetData(), message); + + auto dbName = msg->GetFederatedPartitionSession()->GetDatabaseName(); + UNIT_ASSERT_C(dbName != "dc1", "Message should not be written to unavailable dc1"); + UNIT_ASSERT_C(dbName == "dc2" || dbName == "dc3", "Expected dc2 or dc3, got: " << dbName); + + writeSession->Close(); + readSession->Close(); + */ + } + + Y_UNIT_TEST(WriteLargeMessages) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + // Write single large message for verification + TString largeMessage(10 * 1024, 'x'); // 10KB + UNIT_ASSERT_C(session->Write(largeMessage), "Failed to write large message"); + + session->Close(); + + // Verify the message was written correctly + fixture.VerifyMessagesWritten({largeMessage}); + } + + Y_UNIT_TEST(WriteWithTWriteMessage) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + TString message = "message via TWriteMessage"; + NTopic::TWriteMessage writeMessage(message); + writeMessage.CreateTimestamp(TInstant::Now()); + UNIT_ASSERT(session->Write(std::move(writeMessage))); + + session->Close(); + + // Verify the message was written correctly + fixture.VerifyMessagesWritten({message}); + } + + Y_UNIT_TEST(IsAliveAfterClose) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + auto session = fixture.CreateWriteSessionWithSingleDatabaseFdsResponse(); + UNIT_ASSERT(session); + + // Session should be alive after creation + UNIT_ASSERT(session->IsAlive()); + + // Write a message to verify session is working + UNIT_ASSERT(session->Write("test message")); + + // Session still alive after write + UNIT_ASSERT(session->IsAlive()); + + // Close the session + session->Close(); + + // After close, session should no longer be alive + UNIT_ASSERT(!session->IsAlive()); + + // Write should fail after session is closed + UNIT_ASSERT_C(!session->Write("after close", std::nullopt, std::nullopt, TDuration::MilliSeconds(100)), + "Write should fail after session is closed"); + } + +} + +} // namespace NYdb::NFederatedTopic::NTests diff --git a/src/client/topic/common/simple_blocking_helpers.h b/src/client/topic/common/simple_blocking_helpers.h new file mode 100644 index 00000000000..ddddb94056a --- /dev/null +++ b/src/client/topic/common/simple_blocking_helpers.h @@ -0,0 +1,48 @@ +#pragma once + +#include + +#include +#include + +namespace NYdb::inline V3::NTopic::NDetail { + +// Common helper for blocking write sessions to wait for a continuation token. +// Used by both TSimpleBlockingWriteSession and TSimpleBlockingFederatedWriteSession. +template +std::optional WaitForToken( + TWriter& writer, + std::atomic_bool& closed, + const TDuration& timeout +) { + TInstant startTime = TInstant::Now(); + TDuration remainingTime = timeout; + + std::optional token; + + while (!closed.load() && remainingTime > TDuration::Zero()) { + writer.WaitEvent().Wait(remainingTime); + + for (auto event : writer.GetEvents(false, std::nullopt)) { + if (auto* readyEvent = std::get_if(&event)) { + Y_ABORT_UNLESS(!token.has_value()); + token = std::move(readyEvent->ContinuationToken); + } else if (std::get_if(&event)) { + // discard (maybe log?) + } else if (std::get_if(&event)) { + closed.store(true); + return std::nullopt; + } + } + + if (token.has_value()) { + return token; + } + + remainingTime = timeout - (TInstant::Now() - startTime); + } + + return std::nullopt; +} + +} // namespace NYdb::NTopic::NDetail diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 3b51c48e723..db267c81875 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -1,6 +1,7 @@ #include "write_session.h" #include +#include namespace NYdb::inline V3::NTopic { @@ -135,34 +136,7 @@ bool TSimpleBlockingWriteSession::Write( } std::optional TSimpleBlockingWriteSession::WaitForToken(const TDuration& timeout) { - TInstant startTime = TInstant::Now(); - TDuration remainingTime = timeout; - - std::optional token = std::nullopt; - - while (IsAlive() && remainingTime > TDuration::Zero()) { - Writer->WaitEvent().Wait(remainingTime); - - for (auto event : Writer->GetEvents()) { - if (auto* readyEvent = std::get_if(&event)) { - Y_ABORT_UNLESS(!token.has_value()); - token = std::move(readyEvent->ContinuationToken); - } else if (std::get_if(&event)) { - // discard - } else if (std::get_if(&event)) { - Closed.store(true); - return std::nullopt; - } - } - - if (token.has_value()) { - return token; - } - - remainingTime = timeout - (TInstant::Now() - startTime); - } - - return std::nullopt; + return NDetail::WaitForToken(*Writer, Closed, timeout); } TWriterCounters::TPtr TSimpleBlockingWriteSession::GetCounters() { From 299f85688e3eb28de55b625b8f8700d92b8969eb Mon Sep 17 00:00:00 2001 From: azevaykin <145343289+azevaykin@users.noreply.github.com> Date: Thu, 12 Feb 2026 13:39:37 +0000 Subject: [PATCH 16/93] Access levels in whoami (#32907) --- .github/last_commit.txt | 2 +- .../ydb-cpp-sdk/client/discovery/discovery.h | 12 ++++++++ src/api/protos/ydb_discovery.proto | 12 ++++++++ src/client/discovery/discovery.cpp | 30 +++++++++++++++++++ 4 files changed, 55 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e9c0845fd70..a4e89a9daa0 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -bfb70b6f3508f8ded0feb4dfb7ff92709975bf02 +2d400306362fc343767d227fd8f60c2d640e11a7 diff --git a/include/ydb-cpp-sdk/client/discovery/discovery.h b/include/ydb-cpp-sdk/client/discovery/discovery.h index 91be5831f25..870d8adc5b2 100644 --- a/include/ydb-cpp-sdk/client/discovery/discovery.h +++ b/include/ydb-cpp-sdk/client/discovery/discovery.h @@ -90,9 +90,21 @@ class TWhoAmIResult : public TStatus { TWhoAmIResult(TStatus&& status, const Ydb::Discovery::WhoAmIResult& proto); const std::string& GetUserName() const; const std::vector& GetGroups() const; + bool IsAdministrationAllowed() const; + bool IsMonitoringAllowed() const; + bool IsViewerAllowed() const; + bool IsDatabaseAllowed() const; + bool IsRegisterNodeAllowed() const; + bool IsBootstrapAllowed() const; private: std::string UserName_; std::vector Groups_; + bool IsAdministrationAllowed_ = false; + bool IsMonitoringAllowed_ = false; + bool IsViewerAllowed_ = false; + bool IsDatabaseAllowed_ = false; + bool IsRegisterNodeAllowed_ = false; + bool IsBootstrapAllowed_ = false; }; using TAsyncWhoAmIResult = NThreading::TFuture; diff --git a/src/api/protos/ydb_discovery.proto b/src/api/protos/ydb_discovery.proto index ea6580f85a1..2779becca22 100644 --- a/src/api/protos/ydb_discovery.proto +++ b/src/api/protos/ydb_discovery.proto @@ -63,6 +63,18 @@ message WhoAmIResult { string user = 1; // List of group SIDs (Security IDs) for the user repeated string groups = 2; + // Whether user is allowed to perform administration operations + bool is_administration_allowed = 3; + // Whether user is allowed to perform monitoring operations + bool is_monitoring_allowed = 4; + // Whether user is allowed to view data + bool is_viewer_allowed = 5; + // Whether user is allowed to access database + bool is_database_allowed = 6; + // Whether user is allowed to register dynamic node + bool is_register_node_allowed = 7; + // Whether user is allowed to bootstrap + bool is_bootstrap_allowed = 8; } message WhoAmIResponse { diff --git a/src/client/discovery/discovery.cpp b/src/client/discovery/discovery.cpp index 2da501ead72..21e0d6102b8 100644 --- a/src/client/discovery/discovery.cpp +++ b/src/client/discovery/discovery.cpp @@ -60,6 +60,12 @@ TWhoAmIResult::TWhoAmIResult(TStatus&& status, const Ydb::Discovery::WhoAmIResul for (const auto& group : groups) { Groups_.emplace_back(group); } + IsAdministrationAllowed_ = proto.is_administration_allowed(); + IsMonitoringAllowed_ = proto.is_monitoring_allowed(); + IsViewerAllowed_ = proto.is_viewer_allowed(); + IsDatabaseAllowed_ = proto.is_database_allowed(); + IsRegisterNodeAllowed_ = proto.is_register_node_allowed(); + IsBootstrapAllowed_ = proto.is_bootstrap_allowed(); } const std::string& TWhoAmIResult::GetUserName() const { @@ -70,6 +76,30 @@ const std::vector& TWhoAmIResult::GetGroups() const { return Groups_; } +bool TWhoAmIResult::IsAdministrationAllowed() const { + return IsAdministrationAllowed_; +} + +bool TWhoAmIResult::IsMonitoringAllowed() const { + return IsMonitoringAllowed_; +} + +bool TWhoAmIResult::IsViewerAllowed() const { + return IsViewerAllowed_; +} + +bool TWhoAmIResult::IsDatabaseAllowed() const { + return IsDatabaseAllowed_; +} + +bool TWhoAmIResult::IsRegisterNodeAllowed() const { + return IsRegisterNodeAllowed_; +} + +bool TWhoAmIResult::IsBootstrapAllowed() const { + return IsBootstrapAllowed_; +} + TNodeLocation::TNodeLocation(const Ydb::Discovery::NodeLocation& location) : DataCenterNum(location.has_data_center_num() ? std::make_optional(location.data_center_num()) : std::nullopt) , RoomNum(location.has_room_num() ? std::make_optional(location.room_num()) : std::nullopt) From 1345e7e4709df5aa31c27d6df243d52ec7700b0f Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Thu, 12 Feb 2026 13:39:43 +0000 Subject: [PATCH 17/93] LOGBROKER-10206 KWS (#31966) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/topic/client.h | 8 + .../ydb-cpp-sdk/client/topic/write_session.h | 80 + src/client/topic/impl/topic.cpp | 10 + src/client/topic/impl/topic_impl.cpp | 37 + src/client/topic/impl/topic_impl.h | 2 + src/client/topic/impl/write_session.cpp | 1468 ++++++++++++++++- src/client/topic/impl/write_session.h | 390 +++++ src/client/topic/ut/basic_usage_ut.cpp | 972 +++++++++++ src/client/topic/ut/ut_utils/event_loop.cpp | 84 + src/client/topic/ut/ut_utils/event_loop.h | 37 + tests/integration/topic/basic_usage_it.cpp | 162 ++ 12 files changed, 3226 insertions(+), 26 deletions(-) create mode 100644 src/client/topic/ut/ut_utils/event_loop.cpp create mode 100644 src/client/topic/ut/ut_utils/event_loop.h diff --git a/.github/last_commit.txt b/.github/last_commit.txt index a4e89a9daa0..151499f8102 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -2d400306362fc343767d227fd8f60c2d640e11a7 +b84823f633e98b9b7ec149f4eac112c96de006c9 diff --git a/include/ydb-cpp-sdk/client/topic/client.h b/include/ydb-cpp-sdk/client/topic/client.h index 83d324f7e48..88d50cca684 100644 --- a/include/ydb-cpp-sdk/client/topic/client.h +++ b/include/ydb-cpp-sdk/client/topic/client.h @@ -48,6 +48,14 @@ class TTopicClient { //! Create write session. std::shared_ptr CreateSimpleBlockingWriteSession(const TWriteSessionSettings& settings); + + //! Create simple blocking keyed write session. Experimental feature. DO NOT USE IN PRODUCTION. + std::shared_ptr CreateSimpleBlockingKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + + //! Create keyed write session. Experimental feature. DO NOT USE IN PRODUCTION. + std::shared_ptr CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + + //! Create write session. std::shared_ptr CreateWriteSession(const TWriteSessionSettings& settings); // Commit offset diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index e0c4d4618e2..e2b4a690409 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -144,6 +144,42 @@ struct TWriteSessionSettings : public TRequestSettings { FLUENT_SETTING_DEFAULT(bool, ValidateSeqNo, true); }; +struct TKeyedWriteSessionSettings : public TWriteSessionSettings { + using TSelf = TKeyedWriteSessionSettings; + + enum class EPartitionChooserStrategy { + Bound, + Hash, + }; + + TKeyedWriteSessionSettings() = default; + TKeyedWriteSessionSettings(const TKeyedWriteSessionSettings&) = default; + TKeyedWriteSessionSettings(TKeyedWriteSessionSettings&&) = default; + + TKeyedWriteSessionSettings& operator=(const TKeyedWriteSessionSettings&) = default; + TKeyedWriteSessionSettings& operator=(TKeyedWriteSessionSettings&&) = default; + + //! Session lifetime. + FLUENT_SETTING_DEFAULT(TDuration, SubSessionIdleTimeout, TDuration::Seconds(30)); + + //! Partition chooser strategy. + FLUENT_SETTING_DEFAULT(EPartitionChooserStrategy, PartitionChooserStrategy, EPartitionChooserStrategy::Bound); + + //! Hasher function. + FLUENT_SETTING_DEFAULT(std::function, PartitioningKeyHasher, DefaultPartitioningKeyHasher); + + //! Default partitioning key hasher. + //! Uses MurmurHash. + static std::string DefaultPartitioningKeyHasher(const std::string_view key); + + //! ProducerId prefix to use. + //! ProducerId is generated as ProducerIdPrefix + partition id. + FLUENT_SETTING(std::string, ProducerIdPrefix); + +private: + using TWriteSessionSettings::ProducerId; +}; + //! Contains the message to write and all the options. struct TWriteMessage { using TSelf = TWriteMessage; @@ -276,4 +312,48 @@ class IWriteSession { virtual ~IWriteSession() = default; }; +//! Keyed write session. Experimental SDK. DO NOT USE IN PRODUCTION. +class IKeyedWriteSession { +public: + //! Write single message. + //! continuationToken - a token earlier provided to client with ReadyToAccept event. + virtual void Write(TContinuationToken&& continuationToken, const std::string& key, TWriteMessage&& message, + TTransactionBase* tx = nullptr) = 0; + + //! Future that is set when next event is available. + virtual NThreading::TFuture WaitEvent() = 0; + + //! Wait and return next event. Use WaitEvent() for non-blocking wait. + virtual std::optional GetEvent(bool block = false) = 0; + + //! Get several events in one call. + //! If blocking = false, instantly returns up to maxEventsCount available events. + //! If blocking = true, blocks till maxEventsCount events are available. + //! If maxEventsCount is unset, write session decides the count to return itself. + virtual std::vector GetEvents(bool block = false, std::optional maxEventsCount = std::nullopt) = 0; + + virtual bool Close(TDuration closeTimeout = TDuration::Max()) = 0; + virtual TWriterCounters::TPtr GetCounters() = 0; + virtual ~IKeyedWriteSession() = default; +}; + +//! Simple blocking keyed write session. Experimental SDK. DO NOT USE IN PRODUCTION. +class ISimpleBlockingKeyedWriteSession { +public: + //! Write single message. + //! continuationToken - a token earlier provided to client with ReadyToAccept event. + virtual bool Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx = nullptr, + TDuration blockTimeout = TDuration::Max()) = 0; + + //! Wait for all writes to complete (no more that closeTimeout()), then close. + //! Return true if all writes were completed and acked, false if timeout was reached and some writes were aborted. + virtual bool Close(TDuration closeTimeout = TDuration::Max()) = 0; + + //! Writer counters with different stats (see TWriterConuters). + virtual TWriterCounters::TPtr GetCounters() = 0; + + //! Close() with timeout = 0 and destroy everything instantly. + virtual ~ISimpleBlockingKeyedWriteSession() = default; +}; + } // namespace NYdb::NTopic diff --git a/src/client/topic/impl/topic.cpp b/src/client/topic/impl/topic.cpp index 610a4b921a6..cf8c77e0aac 100644 --- a/src/client/topic/impl/topic.cpp +++ b/src/client/topic/impl/topic.cpp @@ -591,6 +591,16 @@ std::shared_ptr TTopicClient::CreateSimpleBlockingW return Impl_->CreateSimpleWriteSession(settings); } +std::shared_ptr TTopicClient::CreateSimpleBlockingKeyedWriteSession( + const TKeyedWriteSessionSettings& settings) { + return Impl_->CreateSimpleKeyedWriteSession(settings); +} + +std::shared_ptr TTopicClient::CreateKeyedWriteSession( + const TKeyedWriteSessionSettings& settings) { + return Impl_->CreateKeyedWriteSession(settings); +} + std::shared_ptr TTopicClient::CreateWriteSession(const TWriteSessionSettings& settings) { return Impl_->CreateWriteSession(settings); } diff --git a/src/client/topic/impl/topic_impl.cpp b/src/client/topic/impl/topic_impl.cpp index 6fe5e604458..acfe6bde587 100644 --- a/src/client/topic/impl/topic_impl.cpp +++ b/src/client/topic/impl/topic_impl.cpp @@ -61,6 +61,43 @@ std::shared_ptr TTopicClient::TImpl::CreateSimpleWr return std::move(session); } +std::shared_ptr TTopicClient::TImpl::CreateSimpleKeyedWriteSession(const TKeyedWriteSessionSettings& settings) { + auto alteredSettings = settings; + { + std::lock_guard guard(Lock); + if (!settings.CompressionExecutor_) { + alteredSettings.CompressionExecutor(Settings.DefaultCompressionExecutor_); + } + + if (!settings.EventHandlers_.HandlersExecutor_) { + alteredSettings.EventHandlers_.HandlersExecutor(Settings.DefaultHandlersExecutor_); + } + } + + auto session = std::make_shared( + alteredSettings, shared_from_this(), Connections_, DbDriverState_ + ); + return session; +} + +std::shared_ptr TTopicClient::TImpl::CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings) { + auto alteredSettings = settings; + { + std::lock_guard guard(Lock); + if (!settings.CompressionExecutor_) { + alteredSettings.CompressionExecutor(Settings.DefaultCompressionExecutor_); + } + + if (!settings.EventHandlers_.HandlersExecutor_) { + alteredSettings.EventHandlers_.HandlersExecutor(Settings.DefaultHandlersExecutor_); + } + } + + return std::make_shared( + alteredSettings, shared_from_this(), Connections_, DbDriverState_ + ); +} + std::shared_ptr TTopicClient::TImpl::CreateReadSessionConnectionProcessorFactory() { using TService = Ydb::Topic::V1::TopicService; using TRequest = Ydb::Topic::StreamReadMessage::FromClient; diff --git a/src/client/topic/impl/topic_impl.h b/src/client/topic/impl/topic_impl.h index 1983d4734d6..669ae12c39f 100644 --- a/src/client/topic/impl/topic_impl.h +++ b/src/client/topic/impl/topic_impl.h @@ -316,6 +316,8 @@ class TTopicClient::TImpl : public TClientImplCommon { // Runtime API. std::shared_ptr CreateReadSession(const TReadSessionSettings& settings); std::shared_ptr CreateSimpleWriteSession(const TWriteSessionSettings& settings); + std::shared_ptr CreateSimpleKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + std::shared_ptr CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings); std::shared_ptr CreateWriteSession(const TWriteSessionSettings& settings); using IReadSessionConnectionProcessorFactory = diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index db267c81875..06f5b2cc711 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -2,6 +2,12 @@ #include #include +#include + +#include +#include +#include +#include namespace NYdb::inline V3::NTopic { @@ -9,11 +15,12 @@ namespace NYdb::inline V3::NTopic { // TWriteSession TWriteSession::TWriteSession( - const TWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState) - : TContextOwner(settings, std::move(client), std::move(connections), std::move(dbDriverState)) { + const TWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState) + : TContextOwner(settings, std::move(client), std::move(connections), std::move(dbDriverState)) +{ } void TWriteSession::Start(const TDuration& delay) { @@ -37,17 +44,19 @@ NThreading::TFuture TWriteSession::WaitEvent() { } void TWriteSession::WriteEncoded(TContinuationToken&& token, std::string_view data, ECodec codec, ui32 originalSize, - std::optional seqNo, std::optional createTimestamp) { + std::optional seqNo, std::optional createTimestamp) { auto message = TWriteMessage::CompressedMessage(std::move(data), codec, originalSize); - if (seqNo.has_value()) + if (seqNo.has_value()) { message.SeqNo(*seqNo); - if (createTimestamp.has_value()) + } + if (createTimestamp.has_value()) { message.CreateTimestamp(*createTimestamp); + } TryGetImpl()->WriteInternal(std::move(token), std::move(message)); } void TWriteSession::WriteEncoded(TContinuationToken&& token, TWriteMessage&& message, - TTransactionBase* tx) + TTransactionBase* tx) { if (tx) { message.Tx(*tx); @@ -56,17 +65,19 @@ void TWriteSession::WriteEncoded(TContinuationToken&& token, TWriteMessage&& mes } void TWriteSession::Write(TContinuationToken&& token, std::string_view data, std::optional seqNo, - std::optional createTimestamp) { + std::optional createTimestamp) { TWriteMessage message{std::move(data)}; - if (seqNo.has_value()) + if (seqNo.has_value()) { message.SeqNo(*seqNo); - if (createTimestamp.has_value()) + } + if (createTimestamp.has_value()) { message.CreateTimestamp(*createTimestamp); + } TryGetImpl()->WriteInternal(std::move(token), std::move(message)); } void TWriteSession::Write(TContinuationToken&& token, TWriteMessage&& message, - TTransactionBase* tx) { + TTransactionBase* tx) { if (tx) { message.Tx(*tx); } @@ -81,15 +92,1299 @@ TWriteSession::~TWriteSession() { TryGetImpl()->Close(TDuration::Zero()); } +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession + +// TKeyedWriteSessionSettings + +std::string TKeyedWriteSessionSettings::DefaultPartitioningKeyHasher(const std::string_view key) { + std::uint64_t hash = MurmurHash(key.data(), key.size()); + return HexEncode(&hash, sizeof(hash)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TPartitionInfo + +bool TKeyedWriteSession::TPartitionInfo::InRange(const std::string_view key) const { + if (FromBound_ > key) { + return false; + } + if (ToBound_.has_value() && *ToBound_ <= key) { + return false; + } + return true; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TWriteSessionWrapper + +TKeyedWriteSession::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint64_t partition) + : Session(std::move(session)) + , Partition(static_cast(partition)) + , QueueSize(0) +{} + +bool TKeyedWriteSession::TWriteSessionWrapper::IsQueueEmpty() const { + return QueueSize == 0; +} + +bool TKeyedWriteSession::TWriteSessionWrapper::AddToQueue(std::uint64_t delta) { + bool idle = QueueSize == 0; + QueueSize += delta; + return idle; +} + +bool TKeyedWriteSession::TWriteSessionWrapper::RemoveFromQueue(std::uint64_t delta) { + Y_ABORT_UNLESS(QueueSize >= delta, "RemoveFromQueue: underflow"); + QueueSize -= delta; + return QueueSize == 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TIdleSession + +bool TKeyedWriteSession::TIdleSession::Less(const std::shared_ptr& other) const { + if (EmptySince == other->EmptySince) { + return Session->Partition < other->Session->Partition; + } + + return EmptySince < other->EmptySince; +} + +bool TKeyedWriteSession::TIdleSession::Comparator::operator()( + const std::shared_ptr& first, + const std::shared_ptr& second) const { + return first->Less(second); +} + +bool TKeyedWriteSession::TIdleSession::IsExpired() const { + return TInstant::Now() - EmptySince > IdleTimeout; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TSplittedPartitionWorker + +TKeyedWriteSession::TSplittedPartitionWorker::TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId, std::uint64_t partitionIdx) + : Session(session) + , PartitionId(partitionId) + , PartitionIdx(partitionIdx) +{} + +void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { + std::unique_lock lock(Lock); + switch (State) { + case EState::Init: + DescribeTopicFuture = Session->Client->DescribeTopic(Session->Settings.Path_, TDescribeTopicSettings()); + lock.unlock(); + DescribeTopicFuture.Subscribe([this](const NThreading::TFuture&) { + std::lock_guard lock(Lock); + MoveTo(EState::GotDescribe); + }); + lock.lock(); + if (State == EState::Init) { + MoveTo(EState::PendingDescribe); + } + break; + case EState::GotDescribe: + HandleDescribeResult(); + LaunchGetMaxSeqNoFutures(lock); + if (State == EState::GotDescribe) { + MoveTo(EState::PendingMaxSeqNo); + } + break; + case EState::PendingDescribe: + case EState::PendingMaxSeqNo: + case EState::Done: + break; + case EState::GotMaxSeqNo: + Session->MessagesWorker->ScheduleResendMessages(PartitionId, MaxSeqNo); + for (const auto& child : Session->Partitions[PartitionIdx].Children_) { + Session->Partitions[child].Locked(false); + } + MoveTo(EState::Done); + break; + } +} + +void TKeyedWriteSession::TSplittedPartitionWorker::MoveTo(EState state) { + State = state; +} + +void TKeyedWriteSession::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t maxSeqNo) { + MaxSeqNo = std::max(MaxSeqNo, maxSeqNo); +} + +bool TKeyedWriteSession::TSplittedPartitionWorker::IsDone() { + std::lock_guard lock(Lock); + return State == EState::Done; +} + +void TKeyedWriteSession::TSplittedPartitionWorker::HandleDescribeResult() { + std::vector newPartitions; + const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); + for (const auto& partition : partitions) { + if (partition.GetPartitionId() != PartitionId) { + continue; + } + + for (const auto& childPartition : partition.GetChildPartitionIds()) { + newPartitions.push_back(childPartition); + } + break; + } + + std::vector children; + const auto& splittedPartition = Session->Partitions[PartitionIdx]; + Session->PartitionsIndex.erase(splittedPartition.FromBound_); + for (const auto& newPartitionId : newPartitions) { + auto partitionDescribeInfo = std::find_if(partitions.begin(), partitions.end(), [newPartitionId](const auto& partition) { + return partition.GetPartitionId() == newPartitionId; + }); + Y_ABORT_UNLESS(partitionDescribeInfo != partitions.end(), "Partition describe info not found"); + Session->PartitionIdsMapping[newPartitionId] = Session->Partitions.size(); + Session->PartitionsIndex[partitionDescribeInfo->GetFromBound().value_or("")] = Session->Partitions.size(); + Session->Partitions.push_back( + TPartitionInfo() + .PartitionId(newPartitionId) + .FromBound(partitionDescribeInfo->GetFromBound().value_or("")) + .ToBound(partitionDescribeInfo->GetToBound()) + .Locked(true)); + children.push_back(Session->Partitions.size() - 1); + } + Session->Partitions[PartitionIdx].Children(children); +} + +void TKeyedWriteSession::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_lock& lock) { + Y_ABORT_UNLESS(DescribeTopicFuture.IsReady(), "DescribeTopicFuture is not ready yet"); + + std::unordered_map partitionToParent; + const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); + for (const auto& partition : partitions) { + auto parentPartitions = partition.GetParentPartitionIds(); + if (parentPartitions.empty()) { + continue; + } + + // we consider here that each partition has only one parent partition + partitionToParent[partition.GetPartitionId()] = parentPartitions.front(); + } + + std::vector ancestors; + std::uint32_t currentPartition = PartitionId; + while (true) { + ancestors.push_back(currentPartition); + + auto parentPartition = partitionToParent.find(currentPartition); + if (parentPartition == partitionToParent.end()) { + break; + } + currentPartition = parentPartition->second; + } + + NotReadyFutures = ancestors.size(); + for (const auto& ancestor : ancestors) { + auto partitionIdx = Session->PartitionIdsMapping.find(ancestor); + Y_ABORT_UNLESS(partitionIdx != Session->PartitionIdsMapping.end(), "Partition index not found"); + auto wrappedSession = Session->SessionsWorker->GetWriteSession(partitionIdx->second, false); + Y_ABORT_UNLESS(wrappedSession, "Write session not found"); + WriteSessions.push_back(wrappedSession); + + auto future = wrappedSession->Session->GetInitSeqNo(); + lock.unlock(); + future.Subscribe([this, wrappedSession](const NThreading::TFuture& result) { + std::lock_guard lock(Lock); + if (result.HasException()) { + LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, TStringBuilder() << "Failed to get max seq no for partition"); + TSessionClosedEvent sessionClosedEvent(EStatus::INTERNAL_ERROR, {}); + Session->GetSessionClosedEventAndDie(wrappedSession, std::move(sessionClosedEvent)); + MoveTo(EState::Done); + return; + } + + UpdateMaxSeqNo(result.GetValue()); + if (--NotReadyFutures == 0) { + MoveTo(EState::GotMaxSeqNo); + } + }); + lock.lock(); + GetMaxSeqNoFutures.push_back(future); + } +} + +NThreading::TFuture TKeyedWriteSession::TSplittedPartitionWorker::Wait() { + if (DescribeTopicFuture.Initialized() && !DescribeTopicFuture.IsReady()) { + return DescribeTopicFuture.IgnoreResult(); + } + + return NThreading::NWait::WaitAny(GetMaxSeqNoFutures); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TEventsWorkerWrapper + +TKeyedWriteSession::TEventsWorker::TEventsWorker(TKeyedWriteSession* session) + : Session(session) +{ + NotReadyPromise = NThreading::NewPromise(); + NotReadyFuture = NotReadyPromise.GetFuture(); + EventsPromise = NThreading::NewPromise(); + EventsFuture = EventsPromise.GetFuture(); + + // Initialize per-partition futures to a valid, non-ready future to avoid TFutures being uninitialized + // (NThreading::WaitAny throws on uninitialized futures). + Futures.resize(Session->Partitions.size(), NotReadyFuture); + + AddReadyToAcceptEvent(); +} + +void TKeyedWriteSession::TEventsWorker::HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + queueIt->second.push_back(TWriteSessionEvent::TEvent(std::move(event))); +} + +void TKeyedWriteSession::TEventsWorker::HandleReadyToAcceptEvent(std::uint64_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event) { + Session->MessagesWorker->HandleContinuationToken(partition, std::move(event.ContinuationToken)); +} + +void TKeyedWriteSession::TEventsWorker::HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint64_t partition) { + if (event.IsSuccess()) { + return; + } + + if (event.GetStatus() == EStatus::OVERLOADED) { + Session->HandleAutoPartitioning(partition); + return; + } + + if (!CloseEvent.has_value()) { + CloseEvent = std::move(event); + } + Session->NonBlockingClose(); +} + +bool TKeyedWriteSession::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint64_t partition) { + while (true) { + auto event = wrappedSession->Session->GetEvent(false); + if (!event) { + break; + } + + if (auto sessionClosedEvent = std::get_if(&*event); sessionClosedEvent) { + HandleSessionClosedEvent(std::move(*sessionClosedEvent), partition); + return true; + } + + if (auto readyToAcceptEvent = std::get_if(&*event)) { + HandleReadyToAcceptEvent(partition, std::move(*readyToAcceptEvent)); + continue; + } + + if (auto acksEvent = std::get_if(&*event)) { + Session->SessionsWorker->OnReadFromSession(wrappedSession); + HandleAcksEvent(partition, std::move(*acksEvent)); + continue; + } + } + + return false; +} + +void TKeyedWriteSession::TEventsWorker::DoWork() { + std::unique_lock lock(Lock); + + while (!ReadyFutures.empty()) { + auto idx = *ReadyFutures.begin(); + ReadyFutures.erase(idx); + lock.unlock(); + // RunEventLoop without Lock: sub-session's WaitEvent() completion may run the Subscribe + // callback (ReadyFutures.insert) synchronously; that callback takes Lock -> same-thread deadlock. + auto isSessionClosed = RunEventLoop(Session->SessionsWorker->GetWriteSession(idx), idx); + if (!isSessionClosed) { + SubscribeToPartition(idx); + } else { + UnsubscribeFromPartition(idx); + } + lock.lock(); + } + + if (!Session->Done.load()) { + TransferEventsToOutputQueue(); + } +} + +void TKeyedWriteSession::TEventsWorker::SubscribeToPartition(std::uint64_t partition) { + if (auto it = Session->SplittedPartitionWorkers.find(partition); it != Session->SplittedPartitionWorkers.end()) { + Futures[partition] = NotReadyFuture; + return; + } + + auto wrappedSession = Session->SessionsWorker->GetWriteSession(partition); + auto newFuture = wrappedSession->Session->WaitEvent(); + while (partition >= Futures.size()) { + Futures.push_back(NotReadyFuture); + } + newFuture.Subscribe([this, partition](const NThreading::TFuture&) { + std::lock_guard lock(Lock); + this->ReadyFutures.insert(partition); + }); + Futures[partition] = newFuture; +} + +void TKeyedWriteSession::TEventsWorker::HandleNewMessage() { + std::lock_guard lock(Lock); + if (Session->MessagesWorker->IsMemoryUsageOK()) { + AddReadyToAcceptEvent(); + } +} + +void TKeyedWriteSession::TEventsWorker::AddReadyToAcceptEvent() { + EventsOutputQueue.push_back(TWriteSessionEvent::TReadyToAcceptEvent(IssueContinuationToken())); + EventsPromise.TrySetValue(); +} + +bool TKeyedWriteSession::TEventsWorker::AddSessionClosedEvent() { + if (!Session->Closed.load()) { + return false; + } + + if (!CloseEvent.has_value()) { + CloseEvent = TSessionClosedEvent(EStatus::SUCCESS, {}); + } + + if (EventsOutputQueue.empty() && (Session->MessagesWorker->IsQueueEmpty() || Session->Done.load())) { + EventsOutputQueue.push_back(*CloseEvent); + return true; + } + + return false; +} + +void TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { + bool eventsTransferred = false; + bool shouldAddReadyToAcceptEvent = false; + std::unordered_map> acks; + + auto messagesWorker = Session->MessagesWorker; + auto buildOutputAckEvent = [](std::deque& acksQueue, std::optional expectedSeqNo) -> TWriteSessionEvent::TAcksEvent { + TWriteSessionEvent::TAcksEvent ackEvent; + + if (expectedSeqNo.has_value()) { + Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo); + } + + auto ack = std::move(acksQueue.front()); + ackEvent.Acks.push_back(std::move(ack)); + acksQueue.pop_front(); + return ackEvent; + }; + auto finishWithAck = [messagesWorker, &shouldAddReadyToAcceptEvent]() { + bool wasMemoryUsageOk = messagesWorker->IsMemoryUsageOK(); + messagesWorker->HandleAck(); + if (messagesWorker->IsMemoryUsageOK() && !wasMemoryUsageOk) { + shouldAddReadyToAcceptEvent = true; + } + }; + + while (messagesWorker->HasInFlightMessages()) { + const auto& head = messagesWorker->GetFrontInFlightMessage(); + + auto remainingAcks = acks.find(head.Partition); + if (remainingAcks != acks.end() && remainingAcks->second.size() > 0) { + EventsOutputQueue.push_back(buildOutputAckEvent(remainingAcks->second, head.Message.SeqNo_)); + finishWithAck(); + continue; + } + + const auto& eventsQueueIt = PartitionsEventQueues.find(head.Partition); + if (eventsQueueIt == PartitionsEventQueues.end() || eventsQueueIt->second.empty()) { + // No events for this message yet, stop processing (preserve order) + break; + } + + auto event = std::move(eventsQueueIt->second.front()); + auto acksEvent = std::get_if(&event); + Y_ABORT_UNLESS(acksEvent, "Expected AcksEvent only in PartitionsEventQueues"); + + std::deque acksQueue; + std::copy(acksEvent->Acks.begin(), acksEvent->Acks.end(), std::back_inserter(acksQueue)); + EventsOutputQueue.push_back(buildOutputAckEvent(acksQueue, head.Message.SeqNo_)); + acks[head.Partition] = std::move(acksQueue); + eventsQueueIt->second.pop_front(); + eventsTransferred = true; + + finishWithAck(); + } + + // this case handles situation: + // 1st message is written to partition 0 + // 2nd message is written to partition 1 + // 3rd message is written to partition 0 + // 4th message is written to partition 1 + // but AcksEvent for partition 0 looks like: + // [ack1, ack3] + // In this case we can not just forget about ack3, because 3rd message is in-flight + // so we will push 'AcksEvent' back to the queue for partition 0 + for (auto& [partition, acksQueue] : acks) { + if (acksQueue.size() > 0) { + TWriteSessionEvent::TAcksEvent ackEvent; + std::copy(acksQueue.begin(), acksQueue.end(), std::back_inserter(ackEvent.Acks)); + PartitionsEventQueues[partition].push_front(std::move(ackEvent)); + } + } + + if (shouldAddReadyToAcceptEvent) { + AddReadyToAcceptEvent(); + } + + if (eventsTransferred) { + EventsPromise.TrySetValue(); + } +} + +std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueBegin(std::uint64_t partition) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + return queueIt->second.begin(); +} + +std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueEnd(std::uint64_t partition) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + return queueIt->second.end(); +} + +std::optional TKeyedWriteSession::TEventsWorker::GetEventImpl(bool block) { + std::unique_lock lock(Lock); + if (EventsOutputQueue.empty() && block) { + lock.unlock(); + WaitEvent().Wait(); + lock.lock(); + } + + if (!EventsOutputQueue.empty()) { + auto event = std::move(EventsOutputQueue.front()); + EventsOutputQueue.pop_front(); + return event; + } + + return std::nullopt; +} + +std::optional TKeyedWriteSession::TEventsWorker::GetEvent(bool block) { + { + std::unique_lock lock(Lock); + AddSessionClosedEvent(); + } + auto event = GetEventImpl(block); + + return event; +} + +std::vector TKeyedWriteSession::TEventsWorker::GetEvents(bool block, std::optional maxEventsCount) { + if (maxEventsCount.has_value() && maxEventsCount.value() == 0) { + return {}; + } + + { + std::unique_lock lock(Lock); + AddSessionClosedEvent(); + } + + std::vector events; + while (true) { + auto event = GetEventImpl(block); + if (!event) { + break; + } + + events.push_back(std::move(*event)); + if (maxEventsCount.has_value() && events.size() >= maxEventsCount.value()) { + break; + } + } + + return events; +} + +NThreading::TFuture TKeyedWriteSession::TEventsWorker::Wait() { + return NThreading::NWait::WaitAny(Futures); +} + +NThreading::TFuture TKeyedWriteSession::TEventsWorker::WaitEvent() { + std::unique_lock lock(Lock); + + AddSessionClosedEvent(); + if (!EventsOutputQueue.empty()) { + return NThreading::MakeFuture(); + } + + if (EventsFuture.IsReady() && !Session->Closed.load()) { + EventsPromise = NThreading::NewPromise(); + EventsFuture = EventsPromise.GetFuture(); + } + + return EventsFuture; +} + +void TKeyedWriteSession::TEventsWorker::UnsubscribeFromPartition(std::uint64_t partition) { + ReadyFutures.erase(partition); + if (partition < Futures.size()) { + Futures[partition] = NotReadyFuture; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TSessionsWorker + +TKeyedWriteSession::TSessionsWorker::TSessionsWorker(TKeyedWriteSession* session) + : Session(session) +{} + +TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::GetWriteSession(std::uint64_t partition, bool directToPartition) { + auto sessionIter = SessionsIndex.find(partition); + if (sessionIter == SessionsIndex.end()) { + return CreateWriteSession(partition, directToPartition); + } + + if (!directToPartition) { + SessionsIndex.erase(sessionIter); + return CreateWriteSession(partition, directToPartition); + } + + return sessionIter->second; +} + +std::string TKeyedWriteSession::TSessionsWorker::GetProducerId(std::uint64_t partitionId) { + return std::format("{}_{}", Session->Settings.ProducerIdPrefix_, partitionId); +} + +TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::CreateWriteSession(std::uint64_t partition, bool directToPartition) { + auto partitionId = Session->Partitions[partition].PartitionId_; + auto producerId = GetProducerId(partitionId); + TWriteSessionSettings alteredSettings = Session->Settings; + + alteredSettings + .ProducerId(producerId) + .MessageGroupId(producerId) + .MaxMemoryUsage(std::numeric_limits::max()) + .RetryPolicy(Session->RetryPolicy) + .EventHandlers(TWriteSessionSettings::TEventHandlers() + .ReadyToAcceptHandler({}) + .AcksHandler({}) + .SessionClosedHandler({})); + + if (directToPartition) { + alteredSettings.DirectWriteToPartition(true); + alteredSettings.PartitionId(partitionId); + } + auto writeSession = std::make_shared( + Session->Client->CreateWriteSession(alteredSettings), + partition); + + SessionsIndex.emplace(partition, writeSession); + + Session->EventsWorker->SubscribeToPartition(partition); + return writeSession; +} + +void TKeyedWriteSession::TSessionsWorker::DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout, bool mustBeEmpty) { + if (it == SessionsIndex.end() || !it->second) { + return; + } + + auto closeResult = it->second->Session->Close(closeTimeout); + Y_ABORT_UNLESS(!mustBeEmpty || closeResult, "There are still messages in flight"); + const std::uint64_t partition = it->second->Partition; + if (it->second->IdleSession) { + auto itIdle = IdlerSessionsIndex.find(partition); + if (itIdle != IdlerSessionsIndex.end()) { + IdlerSessions.erase(itIdle->second); + IdlerSessionsIndex.erase(itIdle); + } + it->second->IdleSession.reset(); + } + it = SessionsIndex.erase(it); + + Session->EventsWorker->UnsubscribeFromPartition(partition); +} + +void TKeyedWriteSession::TSessionsWorker::OnReadFromSession(WrappedWriteSessionPtr wrappedSession) { + if (wrappedSession->RemoveFromQueue(1)) { + Y_ABORT_UNLESS(!wrappedSession->IdleSession, "IdleSession is already set"); + auto idleSessionPtr = std::make_shared(wrappedSession.get(), TInstant::Now(), Session->Settings.SubSessionIdleTimeout_); + auto [itIdle, inserted] = IdlerSessions.insert(idleSessionPtr); + Y_ABORT_UNLESS(inserted, "Duplicate idle session for partition"); + IdlerSessionsIndex[wrappedSession->Partition] = itIdle; + wrappedSession->IdleSession = idleSessionPtr; + } +} + +void TKeyedWriteSession::TSessionsWorker::OnWriteToSession(WrappedWriteSessionPtr wrappedSession) { + if (wrappedSession->AddToQueue(1) && wrappedSession->IdleSession) { + auto itIdle = IdlerSessionsIndex.find(wrappedSession->Partition); + if (itIdle != IdlerSessionsIndex.end()) { + IdlerSessions.erase(itIdle->second); + IdlerSessionsIndex.erase(itIdle); + } + wrappedSession->IdleSession.reset(); + } +} + +void TKeyedWriteSession::TSessionsWorker::DoWork() { + for (auto it = IdlerSessions.begin(); it != IdlerSessions.end(); ) { + if (!(*it)->IsExpired()) { + break; + } + + auto sessionIter = SessionsIndex.find((*it)->Session->Partition); + ++it; + if (sessionIter != SessionsIndex.end()) { + DestroyWriteSession(sessionIter, TDuration::Zero()); + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TMessagesWorker + +TKeyedWriteSession::TMessagesWorker::TMessagesWorker(TKeyedWriteSession* session) + : Session(session) +{ +} + +void TKeyedWriteSession::TMessagesWorker::RechoosePartitionIfNeeded(TMessageInfo& message) { + const auto& partitionInfo = Session->Partitions[message.Partition]; + if (partitionInfo.Children_.empty()) { + return; + } + + // this case means that partition was split, so we need to rechoose the partition for the message + auto newPartition = Session->PartitionChooser->ChoosePartition(message.Key); + if (newPartition != message.Partition) { + message.Partition = newPartition; + } +} + +void TKeyedWriteSession::TMessagesWorker::DoWork() { + auto sessionsWorker = Session->SessionsWorker; + while (!PendingMessages.empty() && MessagesToResend.empty()) { + auto& head = PendingMessages.front(); + if (Session->Partitions[head.Partition].Locked_ || Session->SplittedPartitionWorkers.contains(head.Partition)) { + break; + } + + RechoosePartitionIfNeeded(head); + auto msgToSave = head; + auto wrappedSession = sessionsWorker->GetWriteSession(msgToSave.Partition); + if (!SendMessage(wrappedSession, std::move(head))) { + break; + } + + PendingMessages.pop_front(); + sessionsWorker->OnWriteToSession(wrappedSession); + PushInFlightMessage(msgToSave.Partition, std::move(msgToSave)); + } + + std::vector partitionsConsumed; + for (auto& [partition, iter] : MessagesToResend) { + auto inFlightMessagesIndexChain = InFlightMessagesIndex.find(partition); + Y_ABORT_UNLESS(inFlightMessagesIndexChain != InFlightMessagesIndex.end(), "InFlightMessagesIndex not found"); + + while (iter != inFlightMessagesIndexChain->second.end()) { + auto msgToSend = **iter; + auto wrappedSession = sessionsWorker->GetWriteSession(msgToSend.Partition); + if (!SendMessage(wrappedSession, std::move(msgToSend))) { + break; + } + + sessionsWorker->OnWriteToSession(wrappedSession); + ++iter; + } + + if (iter == inFlightMessagesIndexChain->second.end()) { + partitionsConsumed.push_back(partition); + } + } + + for (const auto& partition : partitionsConsumed) { + MessagesToResend.erase(partition); + } + + if (Session->MessagesNotEmptyFuture.IsReady()) { + Session->MessagesNotEmptyPromise = NThreading::NewPromise(); + Session->MessagesNotEmptyFuture = Session->MessagesNotEmptyPromise.GetFuture(); + } +} + +bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wrappedSession, TMessageInfo&& message) { + auto continuationToken = GetContinuationToken(message.Partition); + if (!continuationToken) { + return false; + } + + wrappedSession->Session->Write(std::move(*continuationToken), std::move(message.Message), message.Tx); + return true; +} + +void TKeyedWriteSession::TMessagesWorker::PushInFlightMessage(std::uint64_t partition, TMessageInfo&& message) { + InFlightMessages.push_back(std::move(message)); + auto [listIt, _] = InFlightMessagesIndex.try_emplace(partition, std::list::iterator>()); + listIt->second.push_back(std::prev(InFlightMessages.end())); +} + +void TKeyedWriteSession::TMessagesWorker::HandleAck() { + PopInFlightMessage(); +} + +void TKeyedWriteSession::TMessagesWorker::PopInFlightMessage() { + Y_ABORT_UNLESS(!InFlightMessages.empty()); + const std::uint64_t partition = InFlightMessages.front().Partition; + const auto it = InFlightMessages.begin(); + + auto mapIt = InFlightMessagesIndex.find(partition); + if (mapIt != InFlightMessagesIndex.end()) { + auto& list = mapIt->second; + for (auto listIt = list.begin(); listIt != list.end(); ++listIt) { + if (*listIt == it) { + list.erase(listIt); + break; + } + } + if (list.empty()) { + InFlightMessagesIndex.erase(mapIt); + } + } + + Y_ABORT_UNLESS(it->Message.Data.size() <= MemoryUsage, "MemoryUsage is less than the size of the message"); + MemoryUsage -= it->Message.Data.size(); + InFlightMessages.pop_front(); +} + +bool TKeyedWriteSession::TMessagesWorker::IsMemoryUsageOK() const { + return MemoryUsage <= Session->Settings.MaxMemoryUsage_; +} + +void TKeyedWriteSession::TMessagesWorker::AddMessage(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx) { + PendingMessages.push_back(TMessageInfo(key, std::move(message), partition, tx)); + MemoryUsage += message.Data.size(); +} + +std::optional TKeyedWriteSession::TMessagesWorker::GetContinuationToken(std::uint64_t partition) { + auto it = ContinuationTokens.find(partition); + if (it != ContinuationTokens.end() && !it->second.empty()) { + auto token = std::move(it->second.front()); + it->second.pop_front(); + if (it->second.empty()) { + ContinuationTokens.erase(it); + } + return token; + } + + return std::nullopt; +} + +void TKeyedWriteSession::TMessagesWorker::HandleContinuationToken(std::uint64_t partition, TContinuationToken&& continuationToken) { + auto [it, _] = ContinuationTokens.try_emplace(partition, std::deque()); + it->second.push_back(std::move(continuationToken)); +} + +NThreading::TFuture TKeyedWriteSession::TMessagesWorker::Wait() { + return Session->MessagesNotEmptyFuture; +} + +bool TKeyedWriteSession::TMessagesWorker::IsQueueEmpty() const { + return PendingMessages.empty() && InFlightMessages.empty(); +} + +const TKeyedWriteSession::TMessageInfo& TKeyedWriteSession::TMessagesWorker::GetFrontInFlightMessage() const { + Y_ABORT_UNLESS(!InFlightMessages.empty()); + return InFlightMessages.front(); +} + +bool TKeyedWriteSession::TMessagesWorker::HasInFlightMessages() const { + return !InFlightMessages.empty(); +} + +void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint64_t partition, std::uint64_t afterSeqNo) { + auto it = InFlightMessagesIndex.find(partition); + if (it == InFlightMessagesIndex.end()) { + return; + } + + auto& list = it->second; + auto resendIt = list.begin(); + auto ackQueueIt = Session->EventsWorker->AckQueueBegin(partition); + size_t ackIdx = 0; + auto ackQueueEnd = Session->EventsWorker->AckQueueEnd(partition); + std::vector acksToSend; + + while (resendIt != list.end()) { + if (!(*resendIt)->Message.SeqNo_.has_value() || (*resendIt)->Message.SeqNo_.value() > afterSeqNo) { + break; + } + + auto seqNo = (*resendIt)->Message.SeqNo_.value(); + if (ackQueueIt == ackQueueEnd) { + // this case can happen if the message was sent, but session was closed before the ack was received + TWriteSessionEvent::TWriteAck ack; + ack.SeqNo = seqNo; + acksToSend.push_back(std::move(ack)); + } else { + auto acksEvent = std::get_if(&*ackQueueIt); + if (ackIdx == acksEvent->Acks.size()) { + ++ackQueueIt; + ackIdx = 0; + continue; + } + + if (acksEvent->Acks[ackIdx].SeqNo > seqNo) { + // this case can happen if the message was sent, but session was closed before the ack was received + TWriteSessionEvent::TWriteAck ack; + ack.SeqNo = seqNo; + acksEvent->Acks.insert(acksEvent->Acks.begin() + ackIdx, std::move(ack)); + } + ++ackIdx; + } + ++resendIt; + } + + if (!acksToSend.empty()) { + LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, TStringBuilder() << "Sending acks to partition " << partition << ": " << acksToSend.size()); + TWriteSessionEvent::TAcksEvent event; + event.Acks = std::move(acksToSend); + Session->EventsWorker->HandleAcksEvent(partition, std::move(event)); + } + + for (auto iter = resendIt; iter != list.end(); ++iter) { + auto newPartition = Session->PartitionChooser->ChoosePartition((*iter)->Key); + (*iter)->Partition = newPartition; + + auto [listIt, _] = InFlightMessagesIndex.try_emplace(newPartition, std::list::iterator>()); + listIt->second.push_back(*iter); + } + + if (resendIt != list.end()) { + for (const auto& child : Session->Partitions[partition].Children_) { + if (auto childList = InFlightMessagesIndex.find(child); childList != InFlightMessagesIndex.end()) { + MessagesToResend.try_emplace(child, childList->second.begin()); + } + } + + list.erase(resendIt, list.end()); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TKeyedWriteSessionRetryPolicy + +TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::TKeyedWriteSessionRetryPolicy(TKeyedWriteSession* session) + : Session(session) +{} + +typename TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::IRetryState::TPtr TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::CreateRetryState() const { + struct TRetryState : public IRetryState { + TRetryState(TKeyedWriteSession* session) + : Session(session) + {} + ~TRetryState() = default; + TMaybe GetNextRetryDelay(EStatus status) override { + if (status == EStatus::OVERLOADED) { + return Nothing(); + } + + if (!UserRetryState) { + auto policy = Session->Settings.RetryPolicy_ ? Session->Settings.RetryPolicy_ : NYdb::NTopic::IRetryPolicy::GetDefaultPolicy(); + UserRetryState = policy->CreateRetryState(); + } + + return UserRetryState->GetNextRetryDelay(status); + } + + private: + TKeyedWriteSession* Session; + IRetryState::TPtr UserRetryState; + }; + + return std::make_unique(Session); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession + +TKeyedWriteSession::TKeyedWriteSession( + const TKeyedWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState) + : Connections(connections), + Client(client), + DbDriverState(dbDriverState), + Settings(settings) +{ + if (settings.ProducerIdPrefix_.empty()) { + ythrow TContractViolation("ProducerIdPrefix is required for KeyedWriteSession"); + } + + if (!settings.ProducerId_.empty()) { + ythrow TContractViolation("ProducerId should be empty for KeyedWriteSession, use ProducerIdPrefix instead"); + } + + if (!settings.MessageGroupId_.empty()) { + ythrow TContractViolation("MessageGroupId should be empty for KeyedWriteSession"); + } + + TDescribeTopicSettings describeTopicSettings; + auto topicConfig = client->DescribeTopic(settings.Path_, describeTopicSettings).GetValueSync(); + const auto& partitions = topicConfig.GetTopicDescription().GetPartitions(); + auto partitionChooserStrategy = settings.PartitionChooserStrategy_; + + auto strategy = topicConfig.GetTopicDescription().GetPartitioningSettings().GetAutoPartitioningSettings().GetStrategy(); + auto autoPartitioningEnabled = (strategy != EAutoPartitioningStrategy::Disabled && + strategy != EAutoPartitioningStrategy::Unspecified); + + for (const auto& partition : partitions) { + if (!partition.GetActive()) { + continue; + } + + auto partitionId = partition.GetPartitionId(); + PartitionIdsMapping[partitionId] = Partitions.size(); + Partitions.push_back( + TPartitionInfo() + .PartitionId(partitionId) + .FromBound(partition.GetFromBound().value_or("")) + .ToBound(partition.GetToBound())); + } + + Y_ABORT_UNLESS(!Partitions.empty(), "No active partitions found"); + + switch (partitionChooserStrategy) { + case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound: + PartitioningKeyHasher = settings.PartitioningKeyHasher_; + PartitionChooser = std::make_unique(this); + for (size_t i = 0; i < Partitions.size(); ++i) { + if (i > 0 && Partitions[i].FromBound_.empty() && !Partitions[i].ToBound_.has_value()) { + ythrow TContractViolation("Unbounded partition is not supported for Bound partition chooser strategy"); + } + + PartitionsIndex[Partitions[i].FromBound_] = i; + } + break; + case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash: + if (autoPartitioningEnabled) { + ythrow TContractViolation("Hash partition chooser strategy is not supported for topic with auto partitioning"); + } + + PartitionChooser = std::make_unique(this); + break; + } + + ClosePromise = NThreading::NewPromise(); + CloseFuture = ClosePromise.GetFuture(); + ShutdownPromise = NThreading::NewPromise(); + ShutdownFuture = ShutdownPromise.GetFuture(); + MessagesNotEmptyPromise = NThreading::NewPromise(); + MessagesNotEmptyFuture = MessagesNotEmptyPromise.GetFuture(); + + SessionsWorker = std::make_shared(this); + MessagesWorker = std::make_shared(this); + EventsWorker = std::make_shared(this); + RetryPolicy = std::make_shared(this); + + // Start handlers executor for user callbacks (Acks/ReadyToAccept/SessionClosed/Common). + Settings.EventHandlers_.HandlersExecutor_->Start(); + + RunUserEventLoop(); + NextFuture = Next(false); + NextFuture.Subscribe([this](const NThreading::TFuture&) { + RunMainWorker(); + }); +} + +const std::vector& TKeyedWriteSession::GetPartitions() const { + return Partitions; +} + +void TKeyedWriteSession::Write(TContinuationToken&&, const std::string& key, TWriteMessage&& message, TTransactionBase* tx) { + { + std::lock_guard lock(GlobalLock); + if (Closed.load()) { + return; + } + + if ((message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithoutSeqNo) + || (!message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithSeqNo)) { + ythrow TContractViolation("Can not mix messages with and without seqNo"); + } + + if (SeqNoStrategy == ESeqNoStrategy::NotInitialized) { + SeqNoStrategy = message.SeqNo_.has_value() ? ESeqNoStrategy::WithSeqNo : ESeqNoStrategy::WithoutSeqNo; + } + + auto partition = PartitionChooser->ChoosePartition(key); + MessagesWorker->AddMessage(key, std::move(message), partition, tx); + EventsWorker->HandleNewMessage(); + } + + MessagesNotEmptyPromise.TrySetValue(); +} + +bool TKeyedWriteSession::Close(TDuration closeTimeout) { + if (Closed.exchange(true)) { + std::lock_guard lock(GlobalLock); + return MessagesWorker->IsQueueEmpty(); + } + + SetCloseDeadline(closeTimeout); + + ClosePromise.TrySetValue(); + ShutdownFuture.Wait(CloseDeadline); + RunUserEventLoop(); + Done.store(true); + + // No need to lock here, because we are waiting for the shutdown future and it will block until the main worker is done + return MessagesWorker->IsQueueEmpty(); +} + +void TKeyedWriteSession::NonBlockingClose() { + Closed.store(true); + Done.store(true); +} + +void TKeyedWriteSession::SetCloseDeadline(const TDuration& closeTimeout) { + std::lock_guard lock(GlobalLock); + CloseDeadline = TInstant::Now() + closeTimeout; +} + +TKeyedWriteSession::~TKeyedWriteSession() { + Close(TDuration::Zero()); + Settings.EventHandlers_.HandlersExecutor_->Stop(); + ShutdownFuture.Wait(); +} + +NThreading::TFuture TKeyedWriteSession::WaitEvent() { + return EventsWorker->WaitEvent(); +} + +std::optional TKeyedWriteSession::GetEvent(bool block) { + return EventsWorker->GetEvent(block); +} + +std::vector TKeyedWriteSession::GetEvents(bool block, std::optional maxEventsCount) { + return EventsWorker->GetEvents(block, maxEventsCount); +} + +TDuration TKeyedWriteSession::GetCloseTimeout() { + std::lock_guard lock(GlobalLock); + auto now = TInstant::Now(); + if (CloseDeadline <= now) { + return TDuration::Zero(); + } + return CloseDeadline - now; +} + +void TKeyedWriteSession::RunSplittedPartitionWorkers() { + if (SplittedPartitionWorkers.empty()) { + return; + } + + std::vector toRemove; + for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { + if (splittedPartitionWorker->IsDone()) { + toRemove.push_back(partition); + continue; + } + + splittedPartitionWorker->DoWork(); + } + + for (const auto& partition : toRemove) { + SplittedPartitionWorkers.erase(partition); + } +} + +NThreading::TFuture TKeyedWriteSession::Next(bool isClosed) { + std::vector> futures{ + EventsWorker->Wait(), + MessagesWorker->Wait() + }; + + for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { + futures.push_back(splittedPartitionWorker->Wait()); + } + + if (!isClosed) { + futures.push_back(CloseFuture); + } + + return NThreading::NWait::WaitAny(futures); +} + +void TKeyedWriteSession::RunUserEventLoop() { + if (!Settings.EventHandlers_.AcksHandler_ && + !Settings.EventHandlers_.ReadyToAcceptHandler_ && + !Settings.EventHandlers_.SessionClosedHandler_) { + return; + } + + auto handlersExecutor = Settings.EventHandlers_.HandlersExecutor_; + if (!handlersExecutor) { + return; + } + + while (true) { + auto event = GetEvent(false); + if (!event) { + break; + } + + if (auto* readyToAcceptEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.ReadyToAcceptHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*readyToAcceptEvent)]() mutable { + Settings.EventHandlers_.ReadyToAcceptHandler_(ev); + }); + } else { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + continue; + } + + if (auto* acksEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.AcksHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*acksEvent)]() mutable { + Settings.EventHandlers_.AcksHandler_(ev); + }); + } else { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + continue; + } + + if (auto* sessionClosedEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.SessionClosedHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*sessionClosedEvent)]() mutable { + Settings.EventHandlers_.SessionClosedHandler_(ev); + }); + } else if (Settings.EventHandlers_.CommonHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + break; + } + } +} + +void TKeyedWriteSession::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent) { + std::optional receivedSessionClosedEvent; + while (true) { + auto event = wrappedSession->Session->GetEvent(false); + if (!event) { + break; + } + + if (auto* closedEvent = std::get_if(&*event)) { + receivedSessionClosedEvent = std::move(*closedEvent); + break; + } + } + + if (!receivedSessionClosedEvent) { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, TStringBuilder() << "Failed to get session closed event"); + EventsWorker->HandleSessionClosedEvent(std::move(*sessionClosedEvent), wrappedSession->Partition); + } else { + EventsWorker->HandleSessionClosedEvent(std::move(*receivedSessionClosedEvent), wrappedSession->Partition); + } +} + +void TKeyedWriteSession::RunMainWorker() { + RunSplittedPartitionWorkers(); + { + std::unique_lock lock(GlobalLock); + EventsWorker->DoWork(); + if (!Done.load()) { + SessionsWorker->DoWork(); + MessagesWorker->DoWork(); + } + } + RunUserEventLoop(); + + auto isClosed = Closed.load(); + auto closeTimeout = GetCloseTimeout(); + if (isClosed && (Done.load() || MessagesWorker->IsQueueEmpty() || closeTimeout == TDuration::Zero())) { + ShutdownPromise.TrySetValue(); + EventsWorker->EventsPromise.TrySetValue(); + ClosePromise.TrySetValue(); + return; + } + + NextFuture = Next(isClosed); + NextFuture.Subscribe([this](const NThreading::TFuture&) { + RunMainWorker(); + }); +} + +TInstant TKeyedWriteSession::GetCloseDeadline() { + std::lock_guard lock(GlobalLock); + return CloseDeadline; +} + +void TKeyedWriteSession::HandleAutoPartitioning(std::uint64_t partition) { + auto splittedPartitionWorker = std::make_shared(this, Partitions[partition].PartitionId_, partition); + SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); +} + +std::string TKeyedWriteSession::GetProducerId(std::uint64_t partition) { + return std::format("{}_{}", Settings.ProducerIdPrefix_, partition); +} + +TWriterCounters::TPtr TKeyedWriteSession::GetCounters() { + // what should we return here? + return nullptr; +} + +TKeyedWriteSession::TBoundPartitionChooser::TBoundPartitionChooser(TKeyedWriteSession* session) + : Session(session) +{} + +std::uint32_t TKeyedWriteSession::TBoundPartitionChooser::ChoosePartition(const std::string_view key) { + auto hashedKey = Session->PartitioningKeyHasher(key); + + auto lowerBound = Session->PartitionsIndex.lower_bound(hashedKey); + if (lowerBound != Session->PartitionsIndex.end() && lowerBound->first == hashedKey) { + return lowerBound->second; + } + + Y_ABORT_IF(lowerBound == Session->PartitionsIndex.begin(), "Lower bound is the first element"); + return std::prev(lowerBound)->second; +} + +TKeyedWriteSession::THashPartitionChooser::THashPartitionChooser(TKeyedWriteSession* session) + : Session(session) +{ +} + +std::uint32_t TKeyedWriteSession::THashPartitionChooser::ChoosePartition(const std::string_view key) { + std::uint64_t hash = MurmurHash(key.data(), key.size()); + return hash % Session->Partitions.size(); +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TSimpleBlockingWriteSession TSimpleBlockingWriteSession::TSimpleBlockingWriteSession( - const TWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState -) { + const TWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState) { auto subSettings = settings; if (settings.EventHandlers_.AcksHandler_) { LOG_LAZY(dbDriverState->Log, TLOG_WARNING, "TSimpleBlockingWriteSession: Cannot use AcksHandler, resetting."); @@ -116,17 +1411,15 @@ uint64_t TSimpleBlockingWriteSession::GetInitSeqNo() { } bool TSimpleBlockingWriteSession::Write( - std::string_view data, std::optional seqNo, std::optional createTimestamp, const TDuration& blockTimeout -) { + std::string_view data, std::optional seqNo, std::optional createTimestamp, const TDuration& blockTimeout) { auto message = TWriteMessage(std::move(data)) - .SeqNo(seqNo) - .CreateTimestamp(createTimestamp); + .SeqNo(seqNo) + .CreateTimestamp(createTimestamp); return Write(std::move(message), nullptr, blockTimeout); } bool TSimpleBlockingWriteSession::Write( - TWriteMessage&& message, TTransactionBase* tx, const TDuration& blockTimeout -) { + TWriteMessage&& message, TTransactionBase* tx, const TDuration& blockTimeout) { auto continuationToken = WaitForToken(blockTimeout); if (continuationToken.has_value()) { Writer->Write(std::move(*continuationToken), std::move(message), tx); @@ -152,4 +1445,129 @@ bool TSimpleBlockingWriteSession::Close(TDuration closeTimeout) { return Writer->Close(std::move(closeTimeout)); } -} // namespace NYdb::NTopic +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TSimpleBlockingKeyedWriteSession + +TSimpleBlockingKeyedWriteSession::TSimpleBlockingKeyedWriteSession( + const TKeyedWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState) + : Writer(std::make_shared(settings, client, connections, dbDriverState)) +{ + ClosePromise = NThreading::NewPromise(); + CloseFuture = ClosePromise.GetFuture(); +} + +void TSimpleBlockingKeyedWriteSession::RunEventLoop() { + while (true) { + auto event = Writer->GetEvent(false); + if (!event) { + break; + } + + if (auto readyToAcceptEvent = std::get_if(&*event)) { + ContinuationTokensQueue.push(std::move(readyToAcceptEvent->ContinuationToken)); + continue; + } + if (std::get_if(&*event)) { + Closed.store(true); + return; + } + if (auto acksEvent = std::get_if(&*event)) { + HandleAcksEvent(std::move(*acksEvent)); + } + } +} + +void TSimpleBlockingKeyedWriteSession::HandleAcksEvent(const TWriteSessionEvent::TAcksEvent& acksEvent) { + for (auto ack : acksEvent.Acks) { + AckedSeqNos.insert(ack.SeqNo); + } +} + +template +bool TSimpleBlockingKeyedWriteSession::Wait(const TDuration& timeout, F&& stopFunc) { + std::unique_lock lock(Lock); + + auto deadline = TInstant::Now() + timeout; + while (true) { + if (TInstant::Now() > deadline) { + return false; + } + + RunEventLoop(); + + if (stopFunc()) { + return true; + } + + if (Closed.load()) { + return false; + } + + std::vector> futures; + futures.push_back(CloseFuture); + futures.push_back(Writer->WaitEvent()); + lock.unlock(); + NThreading::NWait::WaitAny(futures).Wait(deadline); + lock.lock(); + } +} + +std::optional TSimpleBlockingKeyedWriteSession::GetContinuationToken(TDuration timeout) { + std::optional token; + + Wait(timeout, [&]() { + if (!ContinuationTokensQueue.empty()) { + token = std::move(ContinuationTokensQueue.front()); + ContinuationTokensQueue.pop(); + return true; + } + return false; + }); + + return token; +} + +bool TSimpleBlockingKeyedWriteSession::WaitForAck(std::optional seqNo, TDuration timeout) { + return Wait(timeout, [&]() { + if (!seqNo.has_value()) { + if (AckedSeqNos.empty()) { + return false; + } + + AckedSeqNos.erase(AckedSeqNos.begin()); + return true; + } + + if (AckedSeqNos.contains(*seqNo)) { + AckedSeqNos.erase(*seqNo); + return true; + } + return false; + }); +} + +bool TSimpleBlockingKeyedWriteSession::Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx, TDuration blockTimeout) { + auto continuationToken = GetContinuationToken(blockTimeout); + if (!continuationToken) { + return false; + } + + auto seqNo = message.SeqNo_; + Writer->Write(std::move(*continuationToken), std::move(key), std::move(message), tx); + return WaitForAck(seqNo, blockTimeout); +} + +bool TSimpleBlockingKeyedWriteSession::Close(TDuration closeTimeout) { + Closed.store(true); + ClosePromise.TrySetValue(); + return Writer->Close(closeTimeout); +} + +TWriterCounters::TPtr TSimpleBlockingKeyedWriteSession::GetCounters() { + return nullptr; +} + +} // namespace NYdb::inline V3::NTopic diff --git a/src/client/topic/impl/write_session.h b/src/client/topic/impl/write_session.h index 3112fd1361f..d28df6435b8 100644 --- a/src/client/topic/impl/write_session.h +++ b/src/client/topic/impl/write_session.h @@ -1,12 +1,19 @@ #pragma once +#include +#include #include #include #include +#include + #include #include +#include +#include +#include namespace NYdb::inline V3::NTopic { @@ -55,6 +62,347 @@ class TWriteSession : public IWriteSession, void Start(const TDuration& delay); }; +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession + +class TKeyedWriteSession : public IKeyedWriteSession, + public TContinuationTokenIssuer, + public std::enable_shared_from_this { +private: + using WriteSessionPtr = std::shared_ptr; + + struct TPartitionInfo { + using TSelf = TPartitionInfo; + + bool InRange(const std::string_view key) const; + + FLUENT_SETTING(std::string, FromBound); + FLUENT_SETTING(std::optional, ToBound); + FLUENT_SETTING(std::uint32_t, PartitionId); + FLUENT_SETTING(std::vector, Children); + FLUENT_SETTING_DEFAULT(bool, Locked, false); + }; + + struct TMessageInfo { + TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx) + : Key(key) + , Message(std::move(message)) + , Partition(partition) + , Tx(tx) + {} + + std::string Key; + TWriteMessage Message; + std::uint32_t Partition; + TTransactionBase* Tx; + bool MovedToNewPartition = false; + }; + + struct TIdleSession; + + struct TWriteSessionWrapper { + WriteSessionPtr Session; + const std::uint32_t Partition; + std::uint64_t QueueSize = 0; + std::shared_ptr IdleSession = nullptr; + + TWriteSessionWrapper(WriteSessionPtr session, std::uint64_t partition); + + bool IsQueueEmpty() const; + bool AddToQueue(std::uint64_t delta); + bool RemoveFromQueue(std::uint64_t delta); + }; + + using WrappedWriteSessionPtr = std::shared_ptr; + + struct TIdleSession { + TIdleSession(TWriteSessionWrapper* session, TInstant emptySince, TDuration idleTimeout) + : Session(session) + , EmptySince(emptySince) + , IdleTimeout(idleTimeout) + {} + + const TWriteSessionWrapper* Session; + const TInstant EmptySince; + const TDuration IdleTimeout; + + bool Less(const std::shared_ptr& other) const; + bool IsExpired() const; + + struct Comparator { + bool operator()(const std::shared_ptr& first, const std::shared_ptr& second) const; + }; + }; + + using IdleSessionPtr = std::shared_ptr; + + enum class ESeqNoStrategy { + NotInitialized, + WithoutSeqNo, + WithSeqNo, + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Custom retry policy + + struct TKeyedWriteSessionRetryPolicy : public ::IRetryPolicy { + using TSelf = TKeyedWriteSessionRetryPolicy; + using TPtr = std::shared_ptr; + + TKeyedWriteSessionRetryPolicy(TKeyedWriteSession* session); + ~TKeyedWriteSessionRetryPolicy() = default; + typename IRetryState::TPtr CreateRetryState() const override; + + private: + TKeyedWriteSession* Session; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Workers + + struct TEventsWorker; + + struct TSessionsWorker { + TSessionsWorker(TKeyedWriteSession* session); + WrappedWriteSessionPtr GetWriteSession(std::uint64_t partition, bool directToPartition = true); + void OnReadFromSession(WrappedWriteSessionPtr wrappedSession); + void OnWriteToSession(WrappedWriteSessionPtr wrappedSession); + void DoWork(); + + private: + void AddIdleSession(WrappedWriteSessionPtr wrappedSession, TInstant emptySince, TDuration idleTimeout); + void RemoveIdleSession(std::uint64_t partition); + WrappedWriteSessionPtr CreateWriteSession(std::uint64_t partition, bool directToPartition = true); + + using TSessionsIndexIterator = std::unordered_map::iterator; + void DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout, bool mustBeEmpty = true); + + std::string GetProducerId(std::uint64_t partitionId); + + TKeyedWriteSession* Session; + std::set IdlerSessions; + using IdlerSessionsIterator = std::set::iterator; + std::unordered_map IdlerSessionsIndex; + std::unordered_map SessionsIndex; + }; + + struct TMessagesWorker { + TMessagesWorker(TKeyedWriteSession* session); + + void DoWork(); + + void AddMessage(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx); + void ScheduleResendMessages(std::uint64_t partition, std::uint64_t afterSeqNo); + void HandleAck(); + void HandleContinuationToken(std::uint64_t partition, TContinuationToken&& continuationToken); + bool IsMemoryUsageOK() const; + NThreading::TFuture Wait(); + bool IsQueueEmpty() const; + bool HasInFlightMessages() const; + const TMessageInfo& GetFrontInFlightMessage() const; + + private: + void PushInFlightMessage(std::uint64_t partition, TMessageInfo&& message); + void PopInFlightMessage(); + bool SendMessage(WrappedWriteSessionPtr wrappedSession, TMessageInfo&& message); + std::optional GetContinuationToken(std::uint64_t partition); + void RechoosePartitionIfNeeded(TMessageInfo& message); + + TKeyedWriteSession* Session; + + std::list PendingMessages; + std::list InFlightMessages; + std::unordered_map::iterator>> InFlightMessagesIndex; + + using InFlightMessagesIndexIter = std::list::iterator>::iterator; + std::unordered_map MessagesToResend; + std::unordered_map> ContinuationTokens; + + std::uint64_t MemoryUsage = 0; + }; + + struct TSplittedPartitionWorker { + private: + enum class EState { + Init, + PendingDescribe, + GotDescribe, + PendingMaxSeqNo, + GotMaxSeqNo, + Done, + }; + + void MoveTo(EState state); + void UpdateMaxSeqNo(uint64_t maxSeqNo); + void LaunchGetMaxSeqNoFutures(std::unique_lock& lock); + void HandleDescribeResult(); + + public: + TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId, std::uint64_t partitionIdx); + void DoWork(); + NThreading::TFuture Wait(); + bool IsDone(); + + private: + TKeyedWriteSession* Session; + NThreading::TFuture DescribeTopicFuture; + EState State = EState::Init; + std::uint32_t PartitionId; + std::uint64_t PartitionIdx; + std::uint64_t MaxSeqNo = 0; + std::vector WriteSessions; + std::vector> GetMaxSeqNoFutures; + std::mutex Lock; + std::uint64_t NotReadyFutures = 0; + }; + + struct TEventsWorker { + TEventsWorker(TKeyedWriteSession* session); + + void DoWork(); + NThreading::TFuture Wait(); + NThreading::TFuture WaitEvent(); + void UnsubscribeFromPartition(std::uint64_t partition); + void SubscribeToPartition(std::uint64_t partition); + void HandleNewMessage(); + void HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event); + std::optional GetEvent(bool block); + std::vector GetEvents(bool block, std::optional maxEventsCount = std::nullopt); + std::list::iterator AckQueueBegin(std::uint64_t partition); + std::list::iterator AckQueueEnd(std::uint64_t partition); + + private: + void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint64_t partition); + void HandleReadyToAcceptEvent(std::uint64_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); + bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint64_t partition); + void TransferEventsToOutputQueue(); + void AddReadyToAcceptEvent(); + bool AddSessionClosedEvent(); + std::optional GetEventImpl(bool block); + + TKeyedWriteSession* Session; + + std::vector> Futures; + std::unordered_set ReadyFutures; + std::unordered_map> PartitionsEventQueues; + std::list EventsOutputQueue; + std::mutex Lock; + + NThreading::TFuture NotReadyFuture; + NThreading::TPromise NotReadyPromise; + NThreading::TPromise EventsPromise; + NThreading::TFuture EventsFuture; + + std::optional CloseEvent; + + friend class TKeyedWriteSession; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Partition chooser + + struct IPartitionChooser { + virtual std::uint32_t ChoosePartition(const std::string_view key) = 0; + virtual ~IPartitionChooser() = default; + }; + + struct TBoundPartitionChooser : IPartitionChooser { + TBoundPartitionChooser(TKeyedWriteSession* session); + std::uint32_t ChoosePartition(const std::string_view key) override; + private: + TKeyedWriteSession* Session; + }; + + struct THashPartitionChooser : IPartitionChooser { + THashPartitionChooser(TKeyedWriteSession* session); + std::uint32_t ChoosePartition(const std::string_view key) override; + private: + TKeyedWriteSession* Session; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + + void RunMainWorker(); + + void NonBlockingClose(); + + void SetCloseDeadline(const TDuration& closeTimeout); + + TDuration GetCloseTimeout(); + + std::string GetProducerId(std::uint64_t partition); + + void HandleAutoPartitioning(std::uint64_t partition); + + void RunSplittedPartitionWorkers(); + + NThreading::TFuture Next(bool isClosed); + + void RunUserEventLoop(); + + TInstant GetCloseDeadline(); + + void GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent = std::nullopt); + +public: + TKeyedWriteSession(const TKeyedWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState); + + void Write(TContinuationToken&& continuationToken, const std::string& key, TWriteMessage&& message, + TTransactionBase* tx = nullptr) override; + + NThreading::TFuture WaitEvent() override; + + std::optional GetEvent(bool block = false) override; + + std::vector GetEvents(bool block = false, std::optional maxEventsCount = std::nullopt) override; + + bool Close(TDuration closeTimeout = TDuration::Max()) override; + + TWriterCounters::TPtr GetCounters() override; + + const std::vector& GetPartitions() const; + + ~TKeyedWriteSession(); + +private: + std::shared_ptr Connections; + std::shared_ptr Client; + TDbDriverStatePtr DbDriverState; + + std::vector Partitions; + std::unordered_map PartitionIdsMapping; + std::map PartitionsIndex; + + TKeyedWriteSessionSettings Settings; + ESeqNoStrategy SeqNoStrategy = ESeqNoStrategy::NotInitialized; + + NThreading::TPromise ClosePromise; + NThreading::TFuture CloseFuture; + NThreading::TFuture NextFuture; + NThreading::TPromise ShutdownPromise; + NThreading::TFuture ShutdownFuture; + NThreading::TPromise MessagesNotEmptyPromise; + NThreading::TFuture MessagesNotEmptyFuture; + + std::mutex GlobalLock; + std::atomic_bool Closed = false; + std::atomic_bool Done = false; + TInstant CloseDeadline = TInstant::Now(); + + std::unique_ptr PartitionChooser; + + std::function PartitioningKeyHasher; + + std::shared_ptr EventsWorker; + std::shared_ptr SessionsWorker; + std::unordered_map> SplittedPartitionWorkers; + std::shared_ptr MessagesWorker; + std::shared_ptr RetryPolicy; +}; + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TSimpleBlockingWriteSession @@ -89,5 +437,47 @@ class TSimpleBlockingWriteSession : public ISimpleBlockingWriteSession { std::atomic_bool Closed = false; }; +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TSimpleBlockingKeyedWriteSession + +class TSimpleBlockingKeyedWriteSession : public ISimpleBlockingKeyedWriteSession { +private: + std::optional GetContinuationToken(TDuration timeout); + + void HandleAcksEvent(const TWriteSessionEvent::TAcksEvent& acksEvent); + + bool WaitForAck(std::optional seqNo, TDuration timeout); + + template + bool Wait(const TDuration& timeout, F&& stopFunc); + + void RunEventLoop(); + +public: + TSimpleBlockingKeyedWriteSession( + const TKeyedWriteSessionSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState); + + + bool Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx = nullptr, + TDuration blockTimeout = TDuration::Max()) override; + + bool Close(TDuration closeTimeout = TDuration::Max()) override; + + TWriterCounters::TPtr GetCounters() override; + +protected: + std::shared_ptr Writer; + std::unordered_set AckedSeqNos; + std::queue ContinuationTokensQueue; + + NThreading::TPromise ClosePromise; + NThreading::TFuture CloseFuture; + + std::mutex Lock; + std::atomic_bool Closed = false; +}; } // namespace NYdb::NTopic diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index c8589015a81..b0f04b02117 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -1,5 +1,6 @@ #include "ut_utils/topic_sdk_test_setup.h" +#include #include #include @@ -12,9 +13,11 @@ #include #include #include +#include #include #include +#include #include #include @@ -25,8 +28,13 @@ #include #include +#include #include +#include +#include +#include +#include using namespace std::chrono_literals; @@ -129,6 +137,33 @@ void WriteBinaryProducerIdWithDirectTabletWrite(TTopicSdkTestSetup& setup, UNIT_ASSERT_VALUES_EQUAL(result->Record.GetPartitionResponse().CmdWriteResultSize(), 1); } +static std::string FindKeyForBucket(size_t bucket, size_t bucketsCount) { + for (size_t i = 0; i < 1'000'000; ++i) { + std::string key = "key-" + ToString(i); + if (MurmurHash(key.data(), key.size()) % bucketsCount == bucket) { + return key; + } + } + UNIT_FAIL("Failed to find a key for bucket"); + return {}; +} + +void CreateTopicWithAutoPartitioning(TTopicClient& client) { + TCreateTopicSettings createSettings; + createSettings + .BeginConfigurePartitioningSettings() + .MinActivePartitions(2) + .MaxActivePartitions(100) + .BeginConfigureAutoPartitioningSettings() + .UpUtilizationPercent(2) + .DownUtilizationPercent(1) + .StabilizationWindow(TDuration::Seconds(2)) + .Strategy(EAutoPartitioningStrategy::ScaleUp) + .EndConfigureAutoPartitioningSettings() + .EndConfigurePartitioningSettings(); + client.CreateTopic(TEST_TOPIC, createSettings).Wait(); +} + void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSessionSettings writeSettings, const std::string& message, std::uint32_t count, TTopicSdkTestSetup& setup, std::shared_ptr decompressor) { auto client = setup.MakeClient(); auto session = client.CreateSimpleBlockingWriteSession(writeSettings); @@ -259,6 +294,16 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_VALUES_EQUAL(gotMetaProducerId, expectedEncoded); UNIT_ASSERT_VALUES_EQUAL(Base64Decode(expectedEncoded), binaryProducerId); } + + Y_UNIT_TEST(CreateTopicWithManyPartitions) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + const TString name = "test-topic-" + ToString(TInstant::Now().Seconds()); + setup.CreateTopic(name, TEST_CONSUMER, 100); + + auto describe = setup.MakeClient().DescribeTopic(name).GetValueSync(); + UNIT_ASSERT_C(describe.IsSuccess(), describe.GetIssues().ToOneLineString()); + UNIT_ASSERT_VALUES_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 100); + } Y_UNIT_TEST(CreateTopicWithStreamingConsumer) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; @@ -889,6 +934,933 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } + Y_UNIT_TEST(KeyedWriteSession_UserEventHandlers) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + std::atomic readyCount{0}; + std::atomic acksCount{0}; + std::atomic closedCount{0}; + std::atomic commonCount{0}; + + std::mutex tokensMutex; + std::condition_variable tokensCv; + std::deque readyTokens; + + writeSettings.EventHandlers_.HandlersExecutor(std::make_shared()); + + writeSettings.EventHandlers_.ReadyToAcceptHandler( + [&](TWriteSessionEvent::TReadyToAcceptEvent& ev) { + readyCount.fetch_add(1); + { + std::lock_guard lock(tokensMutex); + readyTokens.emplace_back(std::move(ev.ContinuationToken)); + } + tokensCv.notify_one(); + }); + + writeSettings.EventHandlers_.AcksHandler( + [&](TWriteSessionEvent::TAcksEvent& ev) { + Y_UNUSED(ev); + acksCount.fetch_add(1); + }); + + writeSettings.EventHandlers_.SessionClosedHandler( + [&](const TSessionClosedEvent& ev) { + Y_UNUSED(ev); + closedCount.fetch_add(1); + }); + + writeSettings.EventHandlers_.CommonHandler( + [&](TWriteSessionEvent::TEvent& ev) { + Y_UNUSED(ev); + commonCount.fetch_add(1); + }); + + auto getReadyToken = [&]() -> std::optional { + std::unique_lock lock(tokensMutex); + tokensCv.wait_for(lock, std::chrono::seconds(30), [&]() { return !readyTokens.empty(); }); + if (readyTokens.empty()) { + return std::nullopt; + } + auto token = std::move(readyTokens.front()); + readyTokens.pop_front(); + return token; + }; + + auto session = client.CreateKeyedWriteSession(writeSettings); + + const ui64 messages = 5; + for (ui64 i = 0; i < messages; ++i) { + auto token = getReadyToken(); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + TWriteMessage msg("payload"); + msg.SeqNo(i + 1); + session->Write(std::move(*token), "key-" + ToString(i), std::move(msg)); + } + + UNIT_ASSERT_C(session->Close(TDuration::Seconds(30)), "Failed to close keyed write session"); + + UNIT_ASSERT_C(readyCount.load() > 0, "ReadyToAcceptHandler was not called"); + UNIT_ASSERT_C(acksCount.load() == messages, "AcksHandler does not work properly"); + UNIT_ASSERT_C(closedCount.load() > 0, "SessionClosedHandler was not called"); + UNIT_ASSERT_C(commonCount.load() == 0, "CommonHandler should not be called when type-specific handlers are set"); + } + + Y_UNIT_TEST(KeyedWriteSession_ProducerIdPrefixRequired) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 1); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + UNIT_ASSERT_EXCEPTION(setup.MakeClient().CreateKeyedWriteSession(writeSettings), TContractViolation); + } + + Y_UNIT_TEST(KeyedWriteSession_SessionClosedDueToUserError) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); + auto publicClient = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + auto session = publicClient.CreateKeyedWriteSession(writeSettings); + TKeyedWriteSessionEventLoop eventLoop(session); + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + + TWriteMessage msg("msg0"); + msg.SeqNo(0); + session->Write(std::move(*token), "key", std::move(msg)); + + auto readyToAcceptEvent = session->GetEvent(false); + UNIT_ASSERT_C(std::holds_alternative(*readyToAcceptEvent), "ReadyToAcceptEvent is not received"); + + UNIT_ASSERT_C(session->WaitEvent().Wait(TDuration::Seconds(1000)), "Timed out waiting for event"); + auto event = session->GetEvent(false); + UNIT_ASSERT_C(event, "Event is not received"); + auto sessionClosedEvent = std::get_if(&*event); + UNIT_ASSERT_C(sessionClosedEvent, "SessionClosedEvent is not received"); + UNIT_ASSERT_C(sessionClosedEvent->GetStatus() == EStatus::BAD_REQUEST, "Status is not BAD_REQUEST"); + UNIT_ASSERT(!session->Close(TDuration::Seconds(10))); + } + + Y_UNIT_TEST(KeyedWriteSession_NoAutoPartitioning_HashPartitionChooser) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); + + // Capture partition ids in the same order as DescribeTopic returns them + // (the keyed session uses the same DescribeTopic ordering to map hash bucket -> partition id). + auto publicClient = setup.MakeClient(); + auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); + auto before = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + UNIT_ASSERT_C(before.IsSuccess(), before.GetIssues().ToOneLineString()); + const auto& beforePartitions = before.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_VALUES_EQUAL(beforePartitions.size(), 2); + const ui64 partitionId0 = beforePartitions[0].GetPartitionId(); + const ui64 partitionId1 = beforePartitions[1].GetPartitionId(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + auto session = publicClient.CreateKeyedWriteSession(writeSettings); + + const std::string key0 = FindKeyForBucket(0, 2); + const std::string key1 = FindKeyForBucket(1, 2); + + const ui64 count0 = 7; + const ui64 count1 = 11; + + TKeyedWriteSessionEventLoop eventLoop(session); + + auto seqNo = 1; + for (ui64 i = 0; i < count0; ++i) { + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + + TWriteMessage msg("msg0"); + msg.SeqNo(seqNo++); + session->Write(std::move(*token), key0, std::move(msg)); + } + for (ui64 i = 0; i < count1; ++i) { + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + TWriteMessage msg("msg1"); + msg.SeqNo(seqNo++); + session->Write(std::move(*token), key1, std::move(msg)); + } + + UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + + auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); + const auto& afterPartitions = after.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_VALUES_EQUAL(afterPartitions.size(), 2); + + std::unordered_map endOffsets; + for (const auto& p : afterPartitions) { + auto stats = p.GetPartitionStats(); + UNIT_ASSERT(stats.has_value()); + endOffsets[p.GetPartitionId()] = stats->GetEndOffset(); + } + + auto it0 = endOffsets.find(partitionId0); + auto it1 = endOffsets.find(partitionId1); + UNIT_ASSERT(it0 != endOffsets.end()); + UNIT_ASSERT(it1 != endOffsets.end()); + + const ui64 endOffset0 = it0->second; + const ui64 endOffset1 = it1->second; + + // Partition ordering in DescribeTopic is not a part of public API contract, so allow swapping. + UNIT_ASSERT_VALUES_EQUAL(endOffset0 + endOffset1, count0 + count1); + UNIT_ASSERT_C( + (endOffset0 == count0 && endOffset1 == count1) || (endOffset0 == count1 && endOffset1 == count0), + TStringBuilder() << "Unexpected end offsets distribution: " + << "partitionId0=" << partitionId0 << " endOffset0=" << endOffset0 << ", " + << "partitionId1=" << partitionId1 << " endOffset1=" << endOffset1 << ", " + << "expected (" << count0 << "," << count1 << ") in any order" + ); + } + + Y_UNIT_TEST(KeyedWriteSession_NoAutoPartitioning_BoundPartitionChooser) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopicWithAutoscale(TEST_TOPIC, TEST_CONSUMER, 5, 10); + + auto publicClient = setup.MakeClient(); + auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); + auto before = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + + UNIT_ASSERT_C(before.IsSuccess(), before.GetIssues().ToOneLineString()); + const auto& beforePartitions = before.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_VALUES_EQUAL(beforePartitions.size(), 5); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { + return std::string{key}; + }); + + auto session = publicClient.CreateKeyedWriteSession(writeSettings); + auto keyedSession = std::dynamic_pointer_cast(session); + const auto& partitions = keyedSession->GetPartitions(); + + TKeyedWriteSessionEventLoop eventLoop(session); + + std::unordered_map keysCount; + for (const auto& p : partitions) { + keysCount[p.PartitionId_] = 0; + } + + for (size_t i = 0; i < 100; ++i) { + auto key = CreateGuidAsString(); + for (const auto& p : partitions) { + if (p.InRange(key)) { + keysCount[p.PartitionId_]++; + break; + } + } + + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + TWriteMessage msg("msg"); + msg.SeqNo(i + 1); + session->Write(std::move(*token), key, std::move(msg)); + } + + UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + + auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); + const auto& afterPartitions = after.GetTopicDescription().GetPartitions(); + + std::unordered_map endOffsets; + for (const auto& p : afterPartitions) { + auto stats = p.GetPartitionStats(); + UNIT_ASSERT(stats.has_value()); + endOffsets[p.GetPartitionId()] = stats->GetEndOffset(); + } + + for (const auto& p : partitions) { + auto sb = TStringBuilder() << "partitionId=" << p.PartitionId_ << " endOffset=" << endOffsets[p.PartitionId_] << " keysCount=" << keysCount[p.PartitionId_]; + UNIT_ASSERT_VALUES_EQUAL_C(endOffsets[p.PartitionId_], keysCount[p.PartitionId_], sb.c_str()); + } + } + + Y_UNIT_TEST(KeyedWriteSession_EventLoop_Acks) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateKeyedWriteSession(writeSettings); + TKeyedWriteSessionEventLoop eventLoop(session); + + const ui64 count = 3000; + for (ui64 i = 1; i <= count; ++i) { + auto key = CreateGuidAsString(); + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + auto msg = TWriteMessage("data"); + msg.SeqNo(i); + session->Write(std::move(*token), key, std::move(msg)); + } + + UNIT_ASSERT(eventLoop.WaitForAcks(count, TDuration::Seconds(60))); + eventLoop.CheckAcksOrder(); + UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + } + + Y_UNIT_TEST(KeyedWriteSession_MultiThreadedWrite_Acks) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateKeyedWriteSession(writeSettings); + + constexpr ui64 threadsCount = 4; + constexpr ui64 perThread = 25; + constexpr ui64 total = threadsCount * perThread; + + std::atomic nextSeqNo{1}; + std::vector threads; + threads.reserve(threadsCount); + + TKeyedWriteSessionEventLoop eventLoop(session); + + for (ui64 t = 0; t < threadsCount; ++t) { + threads.emplace_back([&, t]() { + auto key = TStringBuilder() << "key-" << t; + for (ui64 i = 0; i < perThread; ++i) { + std::cout << "thread " << t << " writing message " << i << std::endl; + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + const ui64 seqNo = nextSeqNo.fetch_add(1); + auto msg = TWriteMessage("data"); + msg.SeqNo(seqNo); + session->Write(std::move(*token), key, std::move(msg)); + } + }); + } + + UNIT_ASSERT(eventLoop.WaitForAcks(total, TDuration::Seconds(60))); + UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + } + + Y_UNIT_TEST(KeyedWriteSession_IdleSessionsTimeout) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(5)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateKeyedWriteSession(writeSettings); + + TKeyedWriteSessionEventLoop eventLoop(session); + constexpr ui64 messages = 100; + ui64 seqNo = 1; + + for (ui64 i = 0; i < messages; ++i) { + auto key = CreateGuidAsString(); + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + auto msg = TWriteMessage("data"); + msg.SeqNo(seqNo++); + session->Write(std::move(*token), key, std::move(msg)); + } + + UNIT_ASSERT(eventLoop.WaitForAcks(messages, TDuration::Seconds(60))); + eventLoop.CheckAcksOrder(); + + Sleep(TDuration::Seconds(6)); + + for (ui64 i = 0; i < messages; ++i) { + auto key = CreateGuidAsString(); + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + auto msg = TWriteMessage("data"); + msg.SeqNo(seqNo++); + session->Write(std::move(*token), key, std::move(msg)); + } + + UNIT_ASSERT(eventLoop.WaitForAcks(messages * 2, TDuration::Seconds(60))); + eventLoop.CheckAcksOrder(); + } + + Y_UNIT_TEST(KeyedWriteSession_BoundPartitionChooser_SplitPartition_MultiThreadedAcksOrder) { + NKikimr::NPQ::NTest::TTopicSdkTestSetup setup = NKikimr::NPQ::NTest::CreateSetup(); + setup.CreateTopicWithAutoscale(TEST_TOPIC, TEST_CONSUMER, 1, 100); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { + return std::string{key}; + }); + + auto session = client.CreateKeyedWriteSession(writeSettings); + + constexpr ui64 messages = 1000; + TKeyedWriteSessionEventLoop eventLoop(session); + + std::jthread writer([&]() { + for (ui64 i = 1; i <= messages; ++i) { + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + auto key = CreateGuidAsString(); + auto msg = TWriteMessage("data"); + msg.SeqNo(i); + session->Write(std::move(*token), key, std::move(msg)); + } + }); + + std::jthread splitter([&]() { + Sleep(TDuration::Seconds(1)); + ui64 txId = 1006; + NKikimr::NPQ::NTest::SplitPartition(setup, ++txId, 0, "a"); + }); + + UNIT_ASSERT(eventLoop.WaitForAcks(messages, TDuration::Seconds(60))); + eventLoop.CheckAcksOrder(); + UNIT_ASSERT(session->Close(TDuration::Seconds(30))); + } + + Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_BasicWrite) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 5); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); + + const std::string key1 = "key1"; + const std::string key2 = "key2"; + + // Write several messages with different keys + size_t seqNo = 1; + for (int i = 0; i < 5; ++i) { + TWriteMessage msg("message1-" + ToString(i)); + msg.SeqNo(seqNo++); + bool res = session->Write(key1, std::move(msg)); + UNIT_ASSERT(res); + } + + for (int i = 0; i < 5; ++i) { + TWriteMessage msg("message2-" + ToString(i)); + msg.SeqNo(seqNo++); + bool res = session->Write(key2, std::move(msg)); + UNIT_ASSERT(res); + } + + UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + } + + Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_NoSeqNo) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); + + const ui64 messages = 10; + for (ui64 i = 0; i < messages; ++i) { + TWriteMessage msg("payload-" + ToString(i)); + bool res = session->Write("key-" + ToString(i % 3), std::move(msg)); + UNIT_ASSERT(res); + } + + bool closeRes = session->Close(TDuration::Seconds(30)); + UNIT_ASSERT(closeRes); + } + + Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_ManyMessages) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + + auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); + + ui64 seqNo = 1; + + for (ui64 i = 0; i < 1000; ++i) { + auto key = CreateGuidAsString(); + TWriteMessage msg("payload-" + ToString(seqNo)); + msg.SeqNo(seqNo++); + bool res = session->Write(key, std::move(msg)); + UNIT_ASSERT(res); + } + + bool closeRes = session->Close(TDuration::Seconds(60)); + UNIT_ASSERT(closeRes); + } + + Y_UNIT_TEST(KeyedWriteSession_CloseTimeout) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); + + auto client = setup.MakeClient(); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + auto session = client.CreateKeyedWriteSession(writeSettings); + + TKeyedWriteSessionEventLoop eventLoop(session); + + for (int i = 0; i < 1000; ++i) { + auto token = eventLoop.GetContinuationToken(TDuration::Seconds(10)); + UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); + TWriteMessage msg("message-" + ToString(i)); + msg.SeqNo(i + 1); + session->Write(std::move(*token), "key1", std::move(msg)); + } + + // Test Close timeout + const TDuration closeTimeout = TDuration::Seconds(2); + const TInstant startTime = TInstant::Now(); + session->Close(closeTimeout); + const TDuration actualDuration = TInstant::Now() - startTime; + + // Verify that Close didn't block longer than timeout (with some tolerance) + const TDuration maxExpectedDuration = closeTimeout + TDuration::MilliSeconds(100) + closeTimeout / 10; + UNIT_ASSERT_C( + actualDuration <= maxExpectedDuration + maxExpectedDuration / 10, + TStringBuilder() << "Close() took " << actualDuration << " but timeout was " << closeTimeout + ); + + int attempts = 0; + constexpr int maxAttempts = 1100; + for (attempts = 0; attempts < maxAttempts; ++attempts) { + auto event = session->GetEvent(false); + if (!event) { + break; + } + + auto sessionClosedEvent = std::get_if(&*event); + if (!sessionClosedEvent) { + continue; + } + + UNIT_ASSERT(sessionClosedEvent->IsSuccess()); + break; + } + + UNIT_ASSERT(attempts < maxAttempts); + } + + Y_UNIT_TEST(AutoPartitioning_KeyedWriteSession) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + + std::queue readyTokens1; + std::queue readyTokens2; + std::optional sessionClosedEvent; + std::unordered_set ackedSeqNos; + bool closed = false; + + auto createMessage = [](std::string_view payload, ui64 seqNo) -> TWriteMessage { + TWriteMessage msg(payload); + msg.SeqNo(seqNo); + return msg; + }; + + TCreateTopicSettings createSettings; + createSettings + .BeginConfigurePartitioningSettings() + .MinActivePartitions(2) + .MaxActivePartitions(100) + .BeginConfigureAutoPartitioningSettings() + .UpUtilizationPercent(2) + .DownUtilizationPercent(1) + .StabilizationWindow(TDuration::Seconds(2)) + .Strategy(EAutoPartitioningStrategy::ScaleUp) + .EndConfigureAutoPartitioningSettings() + .EndConfigurePartitioningSettings(); + client.CreateTopic(TEST_TOPIC, createSettings).Wait(); + + auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 2); + + TKeyedWriteSessionSettings writeSettings1; + writeSettings1 + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings1.ProducerIdPrefix("autopartitioning_keyed_1"); + writeSettings1.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); + + TKeyedWriteSessionSettings writeSettings2 = writeSettings1; + writeSettings2.ProducerIdPrefix("autopartitioning_keyed_2"); + + auto session1 = client.CreateKeyedWriteSession(writeSettings1); + auto session2 = client.CreateKeyedWriteSession(writeSettings2); + auto msgData = TString(1_MB, 'a'); + + std::vector keys; + for (const auto& partition : describe.GetTopicDescription().GetPartitions()) { + keys.push_back(partition.GetFromBound().value_or("")); + } + + auto getQueue = [&](const std::shared_ptr& s) -> std::queue& { + if (s == session1) { + return readyTokens1; + } + if (s == session2) { + return readyTokens2; + } + Y_ABORT("Unknown session pointer in AutoPartitioning_KeyedWriteSession"); + }; + + auto eventLoop = [&](std::shared_ptr s) { + while (true) { + auto event = s->GetEvent(false); + if (!event) { + break; + } + if (auto* ready = std::get_if(&*event)) { + getQueue(s).push(std::move(ready->ContinuationToken)); + continue; + } + if (auto* closedEv = std::get_if(&*event)) { + sessionClosedEvent = std::move(*closedEv); + closed = true; + break; + } + if (auto* acks = std::get_if(&*event)) { + for (const auto& ack : acks->Acks) { + UNIT_ASSERT_C( + ackedSeqNos.insert(ack.SeqNo).second, + "Duplicate ack for seqNo " << ack.SeqNo); + } + } + } + }; + + auto getReadyToken = [&](std::shared_ptr s) -> std::optional { + auto& q = getQueue(s); + while (q.empty() && !closed) { + s->WaitEvent().Wait(TDuration::Seconds(5)); + eventLoop(s); + } + if (q.empty()) { + return std::nullopt; + } + auto t = std::move(q.front()); + q.pop(); + return t; + }; + + auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { + auto token = getReadyToken(s); + UNIT_ASSERT(token); + auto key = keys[seqNo % keys.size()]; + if (key.empty()) { + key = "lalala"; + } + s->Write(std::move(*token), key, createMessage(payload, seqNo)); + }; + + { + writeMessage(session1, msgData, 1); + writeMessage(session1, msgData, 2); + Sleep(TDuration::Seconds(5)); + auto d = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + UNIT_ASSERT_EQUAL(d.GetTopicDescription().GetPartitions().size(), 2); + } + + { + writeMessage(session1, msgData, 3); + writeMessage(session1, msgData, 4); + writeMessage(session1, msgData, 5); + writeMessage(session1, msgData, 6); + writeMessage(session1, msgData, 7); + writeMessage(session2, msgData, 8); + writeMessage(session1, msgData, 9); + writeMessage(session1, msgData, 10); + writeMessage(session2, msgData, 11); + writeMessage(session1, msgData, 12); + Sleep(TDuration::Seconds(30)); + for (int i = 0; i < 50 && ackedSeqNos.size() < 12 && !closed; ++i) { + eventLoop(session1); + eventLoop(session2); + if (ackedSeqNos.size() < 12) { + Sleep(TDuration::MilliSeconds(200)); + } + } + UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), 12, + "Expected exactly 12 distinct acks, each seqNo exactly once; got " << ackedSeqNos.size()); + } + + auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + auto partitionsCount = describeResult.GetTopicDescription().GetPartitions().size(); + UNIT_ASSERT_C(partitionsCount >= 4, + TStringBuilder() << "Partitions count: " << partitionsCount << ", expected at least 4"); + + writeMessage(session1, msgData, 13); + writeMessage(session1, msgData, 14); + Sleep(TDuration::Seconds(20)); + for (int i = 0; i < 50 && ackedSeqNos.size() < 14 && !closed; ++i) { + eventLoop(session1); + eventLoop(session2); + if (ackedSeqNos.size() < 14) { + Sleep(TDuration::MilliSeconds(200)); + } + } + + UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), 14, + "Expected exactly 14 distinct acks, each seqNo exactly once; got " << ackedSeqNos.size()); + auto sessionPartitions = dynamic_cast(session1.get())->GetPartitions(); + UNIT_ASSERT_EQUAL_C(sessionPartitions.size(), partitionsCount, + "Expected exactly" << partitionsCount << " partitions, actual: " << sessionPartitions.size()); + + UNIT_ASSERT(session1->Close(TDuration::Seconds(30))); + UNIT_ASSERT(session2->Close(TDuration::Seconds(30))); + } + + Y_UNIT_TEST(AutoPartitioning_KeyedWriteSession_SmallMessages) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + + std::queue readyTokens1; + std::queue readyTokens2; + std::optional sessionClosedEvent; + std::unordered_set ackedSeqNos; + bool closed = false; + + auto createMessage = [](std::string_view payload, ui64 seqNo) -> TWriteMessage { + TWriteMessage msg(payload); + msg.SeqNo(seqNo); + return msg; + }; + + TCreateTopicSettings createSettings; + createSettings + .BeginConfigurePartitioningSettings() + .MinActivePartitions(2) + .MaxActivePartitions(100) + .BeginConfigureAutoPartitioningSettings() + .UpUtilizationPercent(2) + .DownUtilizationPercent(1) + .StabilizationWindow(TDuration::Seconds(2)) + .Strategy(EAutoPartitioningStrategy::ScaleUp) + .EndConfigureAutoPartitioningSettings() + .EndConfigurePartitioningSettings(); + client.CreateTopic(TEST_TOPIC, createSettings).Wait(); + + auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 2); + + TKeyedWriteSessionSettings writeSettings1; + writeSettings1 + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings1.ProducerIdPrefix("autopartitioning_keyed_small_1"); + writeSettings1.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); + + TKeyedWriteSessionSettings writeSettings2 = writeSettings1; + writeSettings2.ProducerIdPrefix("autopartitioning_keyed_small_2"); + + auto session1 = client.CreateKeyedWriteSession(writeSettings1); + auto session2 = client.CreateKeyedWriteSession(writeSettings2); + const size_t msgSize = 256_KB; + auto msgData = TString(msgSize, 'a'); + const ui64 totalMessages = 44; + + std::vector keys; + for (const auto& partition : describe.GetTopicDescription().GetPartitions()) { + keys.push_back(partition.GetFromBound().value_or("")); + } + + auto getQueue = [&](const std::shared_ptr& s) -> std::queue& { + if (s == session1) return readyTokens1; + if (s == session2) return readyTokens2; + Y_ABORT("Unknown session pointer in AutoPartitioning_KeyedWriteSession_SmallMessages"); + }; + + auto eventLoop = [&](std::shared_ptr s) { + while (true) { + auto event = s->GetEvent(false); + if (!event) break; + if (auto* ready = std::get_if(&*event)) { + getQueue(s).push(std::move(ready->ContinuationToken)); + continue; + } + if (auto* closedEv = std::get_if(&*event)) { + sessionClosedEvent = std::move(*closedEv); + closed = true; + break; + } + if (auto* acks = std::get_if(&*event)) { + for (const auto& ack : acks->Acks) { + UNIT_ASSERT_C(ackedSeqNos.insert(ack.SeqNo).second, + "Duplicate ack for seqNo " << ack.SeqNo); + } + } + } + }; + + auto getReadyToken = [&](std::shared_ptr s) -> std::optional { + auto& q = getQueue(s); + while (q.empty() && !closed) { + s->WaitEvent().Wait(TDuration::Seconds(5)); + eventLoop(s); + } + if (q.empty()) return std::nullopt; + auto t = std::move(q.front()); + q.pop(); + return t; + }; + + auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { + auto token = getReadyToken(s); + UNIT_ASSERT(token); + auto key = keys[seqNo % keys.size()]; + if (key.empty()) key = "a"; + s->Write(std::move(*token), key, createMessage(payload, seqNo)); + }; + + { + writeMessage(session1, msgData, 1); + writeMessage(session1, msgData, 2); + Sleep(TDuration::Seconds(5)); + auto d = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + UNIT_ASSERT_EQUAL(d.GetTopicDescription().GetPartitions().size(), 2); + } + + { + for (ui64 seq = 3; seq <= totalMessages - 2; ++seq) { + auto s = (seq % 4 == 0) ? session2 : session1; + writeMessage(s, msgData, seq); + } + Sleep(TDuration::Seconds(30)); + for (int i = 0; i < 80 && ackedSeqNos.size() < totalMessages - 2 && !closed; ++i) { + eventLoop(session1); + eventLoop(session2); + if (ackedSeqNos.size() < totalMessages - 2) Sleep(TDuration::MilliSeconds(200)); + } + UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), totalMessages - 2, + "Expected " << totalMessages - 2 << " acks; got " << ackedSeqNos.size()); + } + + auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + auto partitionsCount = describeResult.GetTopicDescription().GetPartitions().size(); + UNIT_ASSERT_C(partitionsCount >= 3, + TStringBuilder() << "Partitions count: " << partitionsCount << ", expected at least 3 (auto-partitioning)"); + + writeMessage(session1, msgData, totalMessages - 1); + writeMessage(session1, msgData, totalMessages); + Sleep(TDuration::Seconds(20)); + for (int i = 0; i < 80 && ackedSeqNos.size() < totalMessages && !closed; ++i) { + eventLoop(session1); + eventLoop(session2); + if (ackedSeqNos.size() < totalMessages) Sleep(TDuration::MilliSeconds(200)); + } + + UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), totalMessages, + "Expected " << totalMessages << " acks; got " << ackedSeqNos.size()); + auto sessionPartitions = dynamic_cast(session1.get())->GetPartitions(); + UNIT_ASSERT_EQUAL_C(sessionPartitions.size(), partitionsCount, + "Session partitions " << sessionPartitions.size() << " != topic partitions " << partitionsCount); + + UNIT_ASSERT(session1->Close(TDuration::Seconds(30))); + UNIT_ASSERT(session2->Close(TDuration::Seconds(30))); + } + } // Y_UNIT_TEST_SUITE(BasicUsage) } // namespace diff --git a/src/client/topic/ut/ut_utils/event_loop.cpp b/src/client/topic/ut/ut_utils/event_loop.cpp new file mode 100644 index 00000000000..f6130422adc --- /dev/null +++ b/src/client/topic/ut/ut_utils/event_loop.cpp @@ -0,0 +1,84 @@ +#include "event_loop.h" +#include + +namespace NYdb::inline V3::NTopic::NTests { + +TKeyedWriteSessionEventLoop::TKeyedWriteSessionEventLoop(std::shared_ptr session) + : Session_(std::move(session)) +{} + +void TKeyedWriteSessionEventLoop::Run() { + while (true) { + auto event = Session_->GetEvent(false); + if (!event) { + break; + } + if (auto* ready = std::get_if(&*event)) { + std::lock_guard lk(Lock_); + ReadyTokens_.push(std::move(ready->ContinuationToken)); + continue; + } + if (std::get_if(&*event)) { + break; + } + if (auto* acks = std::get_if(&*event)) { + std::lock_guard lk(Lock_); + for (const auto& ack : acks->Acks) { + auto [it, inserted] = AckedSeqNos_.insert(ack.SeqNo); + UNIT_ASSERT_C(inserted, TStringBuilder() << "Ack already received: " << ack.SeqNo); + AckOrder_.push_back(ack.SeqNo); + } + } + } +} + +std::optional TKeyedWriteSessionEventLoop::GetContinuationToken(TDuration timeout) { + const TInstant deadline = TInstant::Now() + timeout; + while (TInstant::Now() < deadline) { + { + std::lock_guard lock(Lock_); + if (!ReadyTokens_.empty()) { + auto token = std::move(ReadyTokens_.front()); + ReadyTokens_.pop(); + return token; + } + } + Session_->WaitEvent().Wait(deadline); + Run(); + } + + std::lock_guard lock(Lock_); + if (!ReadyTokens_.empty()) { + auto token = std::move(ReadyTokens_.front()); + ReadyTokens_.pop(); + return token; + } + return std::nullopt; +} + +bool TKeyedWriteSessionEventLoop::WaitForAcks(size_t count, TDuration timeout) { + const TInstant deadline = TInstant::Now() + timeout; + while (TInstant::Now() < deadline) { + { + std::lock_guard lock(Lock_); + if (AckedSeqNos_.size() >= count) { + return true; + } + } + Session_->WaitEvent().Wait(deadline); + Run(); + } + return false; +} + +void TKeyedWriteSessionEventLoop::CheckAcksOrder() { + std::lock_guard lock(Lock_); + size_t expectedAck = 1; + UNIT_ASSERT_C(AckedSeqNos_.size() == AckOrder_.size(), TStringBuilder() << "Unexpected number of acks: got " << AckOrder_.size() << ", expected " << AckedSeqNos_.size()); + for (const auto& ack : AckOrder_) { + UNIT_ASSERT_VALUES_EQUAL_C(ack, expectedAck, TStringBuilder() << "Unexpected ack order: got " << ack << ", expected " << expectedAck); + expectedAck++; + } +} + +} // namespace NYdb::inline V3::NTopic::NTests diff --git a/src/client/topic/ut/ut_utils/event_loop.h b/src/client/topic/ut/ut_utils/event_loop.h new file mode 100644 index 00000000000..a1fa6ca60f6 --- /dev/null +++ b/src/client/topic/ut/ut_utils/event_loop.h @@ -0,0 +1,37 @@ +#pragma once + +#include + +#include + +#include +#include +#include +#include + +namespace NYdb::inline V3::NTopic::NTests { + +//! Helper for keyed write session tests: runs event loop and provides continuation tokens. +class TKeyedWriteSessionEventLoop { +public: + explicit TKeyedWriteSessionEventLoop(std::shared_ptr session); + + //! Block until a continuation token is available or timeout. Calls Run() while waiting. + //! Returns nullopt on timeout or if session was closed before a token appeared. + std::optional GetContinuationToken(TDuration timeout); + + bool WaitForAcks(size_t count, TDuration timeout); + void CheckAcksOrder(); + +private: + //! Process all currently available events. Returns true if SessionClosed was seen. + void Run(); + + std::shared_ptr Session_; + std::queue ReadyTokens_; + std::unordered_set AckedSeqNos_; + std::vector AckOrder_; + std::mutex Lock_; +}; + +} // namespace NYdb::inline V3::NTopic::NTests diff --git a/tests/integration/topic/basic_usage_it.cpp b/tests/integration/topic/basic_usage_it.cpp index 585c59bcb8b..d4e05d0ce47 100644 --- a/tests/integration/topic/basic_usage_it.cpp +++ b/tests/integration/topic/basic_usage_it.cpp @@ -15,6 +15,7 @@ #include #include +#include namespace NYdb::inline V3::NPersQueue::NTests { @@ -39,6 +40,83 @@ std::uint64_t TSimpleWriteSessionTestAdapter::GetAcquiredMessagesCount() const { return 0; } +class TKeyedWriteSessionTestAdapter { +public: + TKeyedWriteSessionTestAdapter(NTopic::IKeyedWriteSession* session); + + void WaitForAcks(size_t count, TDuration timeout); + std::optional GetContinuationToken(TDuration timeout); + size_t GetAckedSeqNosCount() const; + bool ValidateAcksOrder() const; + +private: + void RunEventLoop(TDuration timeout, size_t stopOnAcksCount, bool stopOnContinuationToken = false); + + NTopic::IKeyedWriteSession* Session; + std::queue tokens; + std::vector ackedSeqNos; +}; + +TKeyedWriteSessionTestAdapter::TKeyedWriteSessionTestAdapter(NTopic::IKeyedWriteSession* session) + : Session(session) +{} + +void TKeyedWriteSessionTestAdapter::WaitForAcks(size_t count, TDuration timeout) { + RunEventLoop(timeout, count, false); +} + +bool TKeyedWriteSessionTestAdapter::ValidateAcksOrder() const { + size_t expectedSeqNo = 1; + for (const auto& seqNo : ackedSeqNos) { + if (seqNo != expectedSeqNo) { + return false; + } + expectedSeqNo++; + } + return true; +} + +size_t TKeyedWriteSessionTestAdapter::GetAckedSeqNosCount() const { + return ackedSeqNos.size(); +} + +std::optional TKeyedWriteSessionTestAdapter::GetContinuationToken(TDuration timeout) { + RunEventLoop(timeout, 0, true); + if (tokens.empty()) { + return std::nullopt; + } + auto token = std::move(tokens.front()); + tokens.pop(); + return token; +} + +void TKeyedWriteSessionTestAdapter::RunEventLoop(TDuration timeout, size_t stopOnAcksCount, bool stopOnContinuationToken) { + auto deadline = TInstant::Now() + timeout; + while (TInstant::Now() < deadline) { + Session->WaitEvent().Wait(deadline); + auto event = Session->GetEvent(false); + if (!event) { + continue; + } + if (auto ev = std::get_if(&*event)) { + tokens.push(std::move(ev->ContinuationToken)); + if (stopOnContinuationToken) { + return; + } + continue; + } + if (auto ev = std::get_if(&*event)) { + for (const auto& ack : ev->Acks) { + ackedSeqNos.push_back(ack.SeqNo); + } + + if (ackedSeqNos.size() >= stopOnAcksCount) { + return; + } + } + } +} + } namespace NYdb::inline V3::NTopic::NTests { @@ -854,6 +932,90 @@ TEST_F(BasicUsage, TEST_NAME(TWriteSession_WriteEncoded_Broken)) { } } +TEST_F(BasicUsage, TEST_NAME(TKeyedWriteSessionBasicWrite_NoAutoPartitioning)) { + // Basic write test for keyed write session. + // Write 10 messages with different keys and check that they are written to different partitions. + // Check that the order of messages is preserved. + constexpr auto TOPIC_NAME = "test-topic-2"; + constexpr auto CONSUMER_NAME = "test-consumer-2"; + + CreateTopic(TOPIC_NAME, CONSUMER_NAME, 5, 5); + + auto driver = MakeDriver(); + TTopicClient client(driver); + + auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); + + TKeyedWriteSessionSettings writeSettings; + writeSettings + .Path(GetTopicPath(TOPIC_NAME)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { + return std::string{key}; + }); + + auto session = client.CreateKeyedWriteSession(writeSettings); + auto keyedSession = std::dynamic_pointer_cast(session); + + NPersQueue::NTests::TKeyedWriteSessionTestAdapter testAdapter(session.get()); + for (size_t i = 0; i < 100; ++i) { + auto key = CreateGuidAsString(); + auto token = testAdapter.GetContinuationToken(TDuration::Seconds(30)); + ASSERT_TRUE(token.has_value()) << "Timed out waiting for ReadyToAcceptEvent"; + TWriteMessage msg("msg"); + msg.SeqNo(i + 1); + session->Write(std::move(*token), key, std::move(msg)); + } + + testAdapter.WaitForAcks(100, TDuration::Seconds(30)); + ASSERT_TRUE(session->Close(TDuration::Seconds(10))); + ASSERT_EQ(testAdapter.GetAckedSeqNosCount(), 100ull); + ASSERT_TRUE(testAdapter.ValidateAcksOrder()); + + auto after = client.DescribeTopic(GetTopicPath(TOPIC_NAME), describeTopicSettings).GetValueSync(); + ASSERT_TRUE(after.IsSuccess()) << after.GetIssues().ToOneLineString(); + + auto readSettings = TReadSessionSettings() + .ConsumerName(GetConsumerName(CONSUMER_NAME)) + .AppendTopics(GetTopicPath(TOPIC_NAME)) + ; + std::shared_ptr readSession = client.CreateReadSession(readSettings); + std::uint32_t readMessageCount = 0; + while (readMessageCount < 100) { + std::cerr << "Get event on client\n"; + auto event = *readSession->GetEvent(true); + std::visit(TOverloaded { + [&](TReadSessionEvent::TDataReceivedEvent& event) { + readMessageCount += event.GetMessages().size(); + }, + [&](TReadSessionEvent::TCommitOffsetAcknowledgementEvent&) { + FAIL(); + }, + [&](TReadSessionEvent::TStartPartitionSessionEvent& event) { + event.Confirm(); + }, + [&](TReadSessionEvent::TStopPartitionSessionEvent& event) { + event.Confirm(); + }, + [&](TReadSessionEvent::TEndPartitionSessionEvent& event) { + event.Confirm(); + }, + [&](TReadSessionEvent::TPartitionSessionStatusEvent&) { + FAIL() << "Test does not support lock sessions yet"; + }, + [&](TReadSessionEvent::TPartitionSessionClosedEvent&) { + FAIL() << "Test does not support lock sessions yet"; + }, + [&](TSessionClosedEvent&) { + FAIL() << "Session closed"; + } + }, event); + } +} + namespace { enum class EExpectedTestResult { SUCCESS, From 0bed519923921ef4d42e283f626f4564f9258ec1 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Thu, 12 Feb 2026 13:39:50 +0000 Subject: [PATCH 18/93] LOGBROKER-9686 Add limits to in flight bytes per partition (#33024) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/topic/read_session.h | 3 +++ src/api/protos/ydb_topic.proto | 2 ++ src/client/topic/impl/read_session_impl.ipp | 4 ++++ 4 files changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 151499f8102..802fa67486e 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -b84823f633e98b9b7ec149f4eac112c96de006c9 +37b40e8e2dfa697f068d6561a7edad00aaf3b6e7 diff --git a/include/ydb-cpp-sdk/client/topic/read_session.h b/include/ydb-cpp-sdk/client/topic/read_session.h index 1a7d4be36cc..f79ba7e8f9b 100644 --- a/include/ydb-cpp-sdk/client/topic/read_session.h +++ b/include/ydb-cpp-sdk/client/topic/read_session.h @@ -198,6 +198,9 @@ struct TReadSessionSettings: public TRequestSettings { //! Log. FLUENT_SETTING_OPTIONAL(TLog, Log); + + //! InFlightMemoryController. + FLUENT_SETTING_OPTIONAL(std::uint64_t, PartitionMaxInFlightBytes); }; struct TReadSessionGetEventSettings : public TCommonClientSettingsBase { diff --git a/src/api/protos/ydb_topic.proto b/src/api/protos/ydb_topic.proto index 40f2c31ea4f..f3fdf7c7a8d 100644 --- a/src/api/protos/ydb_topic.proto +++ b/src/api/protos/ydb_topic.proto @@ -327,6 +327,8 @@ message StreamReadMessage { bool direct_read = 4; // Indicates that the SDK supports auto partitioning. bool auto_partitioning_support = 5; + // Max in flight bytes per partition + uint64 partition_max_in_flight_bytes = 6; message TopicReadSettings { // Topic path. diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index f9a4a67d873..e49e76b4f03 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -530,6 +530,10 @@ inline void TSingleClusterReadSessionImpl::InitImpl(TDeferredActionsset_path(TStringType{topic.Path_}); From 694ec90aa1c0c2f5623db2e39a793872f5d85b7c Mon Sep 17 00:00:00 2001 From: flown4qqqq Date: Thu, 12 Feb 2026 13:39:56 +0000 Subject: [PATCH 19/93] Fix typos 'recieve' -> 'receive' (#33517) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_persqueue_v1.proto | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 802fa67486e..23c1c1dd77f 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -37b40e8e2dfa697f068d6561a7edad00aaf3b6e7 +62e09582c29317f5c5cfb51d8ef88c6c77dc8a16 diff --git a/src/api/protos/ydb_persqueue_v1.proto b/src/api/protos/ydb_persqueue_v1.proto index 04655c3852d..3eac02b2d89 100644 --- a/src/api/protos/ydb_persqueue_v1.proto +++ b/src/api/protos/ydb_persqueue_v1.proto @@ -385,7 +385,7 @@ message StreamingReadClientMessage { int64 partition_session_id = 1; } - // Signal for server that client is not ready to recieve more data from this partition. + // Signal for server that client is not ready to receive more data from this partition. message PauseReadRequest { repeated int64 partition_session_ids = 1; } @@ -490,7 +490,7 @@ message StreamingReadServerMessage { } // Command to create and start a partition session. - // Client must react on this signal by sending StartRead when ready recieve data from this partition. + // Client must react on this signal by sending StartRead when ready receive data from this partition. message StartPartitionSessionRequest { // Partition partition stream description. PartitionSession partition_session = 1; From f5a630150dc13cc7cb0eadc62da964b69906fe87 Mon Sep 17 00:00:00 2001 From: Andrei Rykov Date: Thu, 12 Feb 2026 13:40:03 +0000 Subject: [PATCH 20/93] [EXT-1921] Proto annotation-based masking (#33259) --- .github/last_commit.txt | 2 +- src/api/client/yc_private/iam/iam_token.proto | 3 ++- src/api/client/yc_private/iam/iam_token_service.proto | 7 ++++--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 23c1c1dd77f..ca082cb2b06 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -62e09582c29317f5c5cfb51d8ef88c6c77dc8a16 +dbd7178531a80a5d7acad45d423b2b1d975e3a5b diff --git a/src/api/client/yc_private/iam/iam_token.proto b/src/api/client/yc_private/iam/iam_token.proto index 3d7f41f6eb8..47b89f82792 100644 --- a/src/api/client/yc_private/iam/iam_token.proto +++ b/src/api/client/yc_private/iam/iam_token.proto @@ -3,8 +3,9 @@ syntax = "proto3"; package yandex.cloud.priv.iam.v1; import "google/protobuf/timestamp.proto"; +import "src/api/client/yc_private/accessservice/sensitive.proto"; message IamToken { - string iam_token = 1; + string iam_token = 1 [(sensitive) = true]; google.protobuf.Timestamp expires_at = 2; } diff --git a/src/api/client/yc_private/iam/iam_token_service.proto b/src/api/client/yc_private/iam/iam_token_service.proto index d900dff629f..b94bd9da098 100644 --- a/src/api/client/yc_private/iam/iam_token_service.proto +++ b/src/api/client/yc_private/iam/iam_token_service.proto @@ -4,6 +4,7 @@ package yandex.cloud.priv.iam.v1; import "google/api/annotations.proto"; import "google/protobuf/timestamp.proto"; +import "src/api/client/yc_private/accessservice/sensitive.proto"; import "src/api/client/yc_private/iam/iam_token_service_subject.proto"; import "src/api/client/yc_private/iam/yandex_passport_cookie.proto"; import "src/api/client/yc_private/iam/oauth_request.proto"; @@ -35,8 +36,8 @@ service IamTokenService { message CreateIamTokenRequest { oneof identity { - string yandex_passport_oauth_token = 1; - string jwt = 2; + string yandex_passport_oauth_token = 1 [(sensitive) = true]; + string jwt = 2 [(sensitive) = true]; string iam_cookie = 3; YandexPassportCookies yandex_passport_cookies = 4; } @@ -67,7 +68,7 @@ message CreateIamTokenForComputeInstanceRequest { } message CreateIamTokenResponse { - string iam_token = 1; + string iam_token = 1 [(sensitive) = true]; google.protobuf.Timestamp issued_at = 4; google.protobuf.Timestamp expires_at = 2; ts.Subject subject = 3; From 044be497b5147ddcc0ffec49c079f162b5425e71 Mon Sep 17 00:00:00 2001 From: Ilia Shakhov Date: Thu, 12 Feb 2026 13:40:09 +0000 Subject: [PATCH 21/93] Add smart mode to CMS (#33297) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_maintenance.proto | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index ca082cb2b06..defb79aa31e 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -dbd7178531a80a5d7acad45d423b2b1d975e3a5b +58a6a4cd0aa8fb88009fa7ae59317730185aec79 diff --git a/src/api/protos/draft/ydb_maintenance.proto b/src/api/protos/draft/ydb_maintenance.proto index 2d3af6457fd..17f96b9847b 100644 --- a/src/api/protos/draft/ydb_maintenance.proto +++ b/src/api/protos/draft/ydb_maintenance.proto @@ -80,6 +80,12 @@ enum AvailabilityMode { // Ignore any storage group & state storage checks. // Using this mode might cause data unavailability. AVAILABILITY_MODE_FORCE = 3; + + // In this mode: + // - attempts to apply AVAILABILITY_MODE_STRONG; + // - if strong constraints cannot be satisfied, falls back to AVAILABILITY_MODE_WEAK; + // - never escalates to AVAILABILITY_MODE_FORCE. + AVAILABILITY_MODE_SMART = 4; } message MaintenanceTaskOptions { From 88aa044c034c2dc98547a1941029d15e3449009a Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Thu, 12 Feb 2026 13:40:15 +0000 Subject: [PATCH 22/93] LOGBROKER-10274 Restart if got unknown ack from server (#33934) --- .github/last_commit.txt | 2 +- .../impl/write_session_impl.cpp | 25 +++++++++++-------- .../impl/write_session_impl.h | 2 +- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index defb79aa31e..125cba34c45 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -58a6a4cd0aa8fb88009fa7ae59317730185aec79 +607da489f3430e282553917988a4fb7ed2c8e194 diff --git a/src/client/persqueue_public/impl/write_session_impl.cpp b/src/client/persqueue_public/impl/write_session_impl.cpp index e090309c60b..44915db0586 100644 --- a/src/client/persqueue_public/impl/write_session_impl.cpp +++ b/src/client/persqueue_public/impl/write_session_impl.cpp @@ -661,7 +661,7 @@ void TWriteSessionImpl::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t co } else { processResult = ProcessServerMessageImpl(); needSetValue = !InitSeqNoSetDone && processResult.InitSeqNo.has_value() && (InitSeqNoSetDone = true); - if (errorStatus.Ok() && processResult.Ok) { + if (errorStatus.Ok() && processResult.Ok && !processResult.HandleResult.DoRestart) { doRead = true; } } @@ -672,10 +672,8 @@ void TWriteSessionImpl::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t co { std::lock_guard guard(Lock); - if (!errorStatus.Ok()) { - if (processResult.Ok) { // Otherwise, OnError was already called - processResult.HandleResult = RestartImpl(errorStatus); - } + if ((!errorStatus.Ok() && processResult.Ok) || processResult.HandleResult.DoRestart) { // Otherwise, OnError was already called + processResult.HandleResult = RestartImpl(errorStatus); } if (processResult.HandleResult.DoStop) { CloseImpl(std::move(errorStatus)); @@ -785,9 +783,13 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess writeStat, }); - if (CleanupOnAcknowledged(GetIdImpl(sequenceNumber))) { + if (CleanupOnAcknowledged(GetIdImpl(sequenceNumber), result)) { result.Events.emplace_back(TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); } + + if (result.HandleResult.DoRestart) { + return result; + } } //EventsQueue->PushEvent(std::move(acksEvent)); result.Events.emplace_back(std::move(acksEvent)); @@ -803,16 +805,17 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess return result; } -bool TWriteSessionImpl::CleanupOnAcknowledged(ui64 id) { +bool TWriteSessionImpl::CleanupOnAcknowledged(ui64 id, TProcessSrvMessageResult& processResult) { bool result = false; LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "Write session: acknoledged message " << id); UpdateTimedCountersImpl(); if (SentOriginalMessages.empty() || SentOriginalMessages.front().Id != id){ - std::cerr << "State before restart was:\n" << StateStr << "\n\n"; + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "State before restart was: " << StateStr); DumpState(); - std::cerr << "State on ack with id " << id << " is:\n"; - std::cerr << StateStr << "\n\n"; - Y_ABORT("got unknown ack"); + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "State on ack with id " << id << " is:\n" << StateStr); + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Write session: got unknown ack " << id); + processResult.HandleResult.DoRestart = true; + return false; } const auto& sentFront = SentOriginalMessages.front(); diff --git a/src/client/persqueue_public/impl/write_session_impl.h b/src/client/persqueue_public/impl/write_session_impl.h index 483570c944c..86be4fc8d3a 100644 --- a/src/client/persqueue_public/impl/write_session_impl.h +++ b/src/client/persqueue_public/impl/write_session_impl.h @@ -356,7 +356,7 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, //std::string GetDebugIdentity() const; Ydb::PersQueue::V1::StreamingWriteClientMessage GetInitClientMessage(); - bool CleanupOnAcknowledged(ui64 id); + bool CleanupOnAcknowledged(ui64 id, TProcessSrvMessageResult& processResult); bool IsReadyToSendNextImpl(); void DumpState(); ui64 GetNextIdImpl(const std::optional& seqNo); From 2d5d349ed83e0830d3486f22b2dd7d4b43190959 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Thu, 12 Feb 2026 13:40:21 +0000 Subject: [PATCH 23/93] LOGBROKER-10206 Add keyed write session workload & KWS fixes (#33427) --- .github/last_commit.txt | 2 +- .../ydb-cpp-sdk/client/topic/write_session.h | 3 + src/client/topic/impl/write_session.cpp | 898 ++++++++++++------ src/client/topic/impl/write_session.h | 198 ++-- src/client/topic/ut/basic_usage_ut.cpp | 48 +- src/client/topic/ut/ut_utils/event_loop.cpp | 11 +- 6 files changed, 777 insertions(+), 383 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 125cba34c45..1604fb18241 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -607da489f3430e282553917988a4fb7ed2c8e194 +3b235ed1f2fc3977cfc6f99a74123c0097ef9795 diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index e2b4a690409..b3bba9904d8 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -176,6 +176,9 @@ struct TKeyedWriteSessionSettings : public TWriteSessionSettings { //! ProducerId is generated as ProducerIdPrefix + partition id. FLUENT_SETTING(std::string, ProducerIdPrefix); + //! SessionID to use. + FLUENT_SETTING_DEFAULT(std::string, SessionId, ""); + private: using TWriteSessionSettings::ProducerId; }; diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 06f5b2cc711..d9070eb7e5c 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -98,8 +98,17 @@ TWriteSession::~TWriteSession() { // TKeyedWriteSessionSettings std::string TKeyedWriteSessionSettings::DefaultPartitioningKeyHasher(const std::string_view key) { - std::uint64_t hash = MurmurHash(key.data(), key.size()); - return HexEncode(&hash, sizeof(hash)); + const std::uint64_t lo = MurmurHash(key.data(), key.size(), std::uint64_t{0}); + const std::uint64_t hi = MurmurHash(key.data(), key.size(), std::uint64_t{0x9E3779B97F4A7C15ull}); // fixed seed + + const std::uint64_t hiBe = y_absl::gntohll(hi); + const std::uint64_t loBe = y_absl::gntohll(lo); + + std::string out; + out.resize(16); + memcpy(out.data() + 0, &hiBe, 8); + memcpy(out.data() + 8, &loBe, 8); + return out; // 16 bytes } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -115,12 +124,43 @@ bool TKeyedWriteSession::TPartitionInfo::InRange(const std::string_view key) con return true; } +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::TMessageInfo + +TKeyedWriteSession::TMessageInfo::TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx) + : Key(key) + , Data(message.Data) + , Codec(message.Codec) + , OriginalSize(message.OriginalSize) + , SeqNo(message.SeqNo_) + , CreateTimestamp(message.CreateTimestamp_) + , TxInMessage(message.Tx_) + , Tx(tx) + , Partition(partition) +{ + for (const auto& [key, value] : message.MessageMeta_) { + MessageMeta.Fields.emplace_back(key, value); + } +} + +TWriteMessage TKeyedWriteSession::TMessageInfo::BuildMessage() const { + TWriteMessage message(Data); + message.Codec = Codec; + message.OriginalSize = OriginalSize; + message.SeqNo(SeqNo); + message.CreateTimestamp(CreateTimestamp); + for (const auto& [key, value] : MessageMeta.Fields) { + message.MessageMeta_.emplace_back(key, value); + } + message.Tx(TxInMessage); + return message; +} //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TKeyedWriteSession::TWriteSessionWrapper -TKeyedWriteSession::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint64_t partition) +TKeyedWriteSession::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition) : Session(std::move(session)) - , Partition(static_cast(partition)) + , Partition(partition) , QueueSize(0) {} @@ -164,21 +204,49 @@ bool TKeyedWriteSession::TIdleSession::IsExpired() const { //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TKeyedWriteSession::TSplittedPartitionWorker -TKeyedWriteSession::TSplittedPartitionWorker::TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId, std::uint64_t partitionIdx) +TKeyedWriteSession::TSplittedPartitionWorker::TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId) : Session(session) , PartitionId(partitionId) - , PartitionIdx(partitionIdx) -{} +{ + LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "Creating splitted partition worker for partition " << PartitionId); +} + +std::string TKeyedWriteSession::TSplittedPartitionWorker::GetStateName() const { + switch (State) { + case EState::Init: + return "Init"; + case EState::PendingDescribe: + return "PendingDescribe"; + case EState::GotDescribe: + return "GotDescribe"; + case EState::PendingMaxSeqNo: + return "PendingMaxSeqNo"; + case EState::Done: + return "Done"; + case EState::GotMaxSeqNo: + return "GotMaxSeqNo"; + } +} void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { std::unique_lock lock(Lock); + std::weak_ptr session = Session->shared_from_this(); switch (State) { case EState::Init: DescribeTopicFuture = Session->Client->DescribeTopic(Session->Settings.Path_, TDescribeTopicSettings()); lock.unlock(); - DescribeTopicFuture.Subscribe([this](const NThreading::TFuture&) { - std::lock_guard lock(Lock); - MoveTo(EState::GotDescribe); + DescribeTopicFuture.Subscribe([this, session](const NThreading::TFuture&) { + auto sessionPtr = session.lock(); + if (!sessionPtr) { + return; + } + + { + std::lock_guard lock(Lock); + MoveTo(EState::GotDescribe); + } + + sessionPtr->RunMainWorker(); }); lock.lock(); if (State == EState::Init) { @@ -187,6 +255,10 @@ void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { break; case EState::GotDescribe: HandleDescribeResult(); + if (State != EState::GotDescribe) { + break; + } + LaunchGetMaxSeqNoFutures(lock); if (State == EState::GotDescribe) { MoveTo(EState::PendingMaxSeqNo); @@ -197,10 +269,12 @@ void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { case EState::Done: break; case EState::GotMaxSeqNo: + Session->MessagesWorker->RebuildPendingMessagesIndex(PartitionId); Session->MessagesWorker->ScheduleResendMessages(PartitionId, MaxSeqNo); - for (const auto& child : Session->Partitions[PartitionIdx].Children_) { + for (const auto& child : Session->Partitions[PartitionId].Children_) { Session->Partitions[child].Locked(false); } + Session->Partitions[PartitionId].Locked_ = false; MoveTo(EState::Done); break; } @@ -208,6 +282,7 @@ void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { void TKeyedWriteSession::TSplittedPartitionWorker::MoveTo(EState state) { State = state; + LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "Moving splitted partition worker for partition " << PartitionId << " to state " << GetStateName()); } void TKeyedWriteSession::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t maxSeqNo) { @@ -219,45 +294,59 @@ bool TKeyedWriteSession::TSplittedPartitionWorker::IsDone() { return State == EState::Done; } +bool TKeyedWriteSession::TSplittedPartitionWorker::IsInit() { + std::lock_guard lock(Lock); + return State == EState::Init; +} + void TKeyedWriteSession::TSplittedPartitionWorker::HandleDescribeResult() { - std::vector newPartitions; + std::vector newPartitionsIds; const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); for (const auto& partition : partitions) { if (partition.GetPartitionId() != PartitionId) { continue; } - for (const auto& childPartition : partition.GetChildPartitionIds()) { - newPartitions.push_back(childPartition); + LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Found partition " << partition.GetPartitionId() << " for partition " << PartitionId << " children: " << partition.GetChildPartitionIds().size()); + for (const auto& childPartitionId : partition.GetChildPartitionIds()) { + newPartitionsIds.push_back(childPartitionId); } break; } + if (newPartitionsIds.empty()) { + // describe response is incomplete, we need to resend describe request + MoveTo(EState::Init); + Y_ABORT_UNLESS(++Retries < 20, "Too many retries for partition %lu", PartitionId); + LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Describe response is incomplete, we need to resend describe request for partition " << PartitionId); + return; + } + std::vector children; - const auto& splittedPartition = Session->Partitions[PartitionIdx]; + const auto& splittedPartition = Session->Partitions[PartitionId]; Session->PartitionsIndex.erase(splittedPartition.FromBound_); - for (const auto& newPartitionId : newPartitions) { + + for (const auto& newPartitionId : newPartitionsIds) { auto partitionDescribeInfo = std::find_if(partitions.begin(), partitions.end(), [newPartitionId](const auto& partition) { return partition.GetPartitionId() == newPartitionId; }); Y_ABORT_UNLESS(partitionDescribeInfo != partitions.end(), "Partition describe info not found"); - Session->PartitionIdsMapping[newPartitionId] = Session->Partitions.size(); - Session->PartitionsIndex[partitionDescribeInfo->GetFromBound().value_or("")] = Session->Partitions.size(); - Session->Partitions.push_back( - TPartitionInfo() + Session->PartitionsIndex[partitionDescribeInfo->GetFromBound().value_or("")] = newPartitionId; + Session->Partitions[newPartitionId] = TPartitionInfo() .PartitionId(newPartitionId) .FromBound(partitionDescribeInfo->GetFromBound().value_or("")) .ToBound(partitionDescribeInfo->GetToBound()) - .Locked(true)); - children.push_back(Session->Partitions.size() - 1); + .Locked(true); + children.push_back(newPartitionId); } - Session->Partitions[PartitionIdx].Children(children); + + Session->Partitions[PartitionId].Children(children); } void TKeyedWriteSession::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_lock& lock) { Y_ABORT_UNLESS(DescribeTopicFuture.IsReady(), "DescribeTopicFuture is not ready yet"); - std::unordered_map partitionToParent; + std::unordered_map partitionIdToParentId; const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); for (const auto& partition : partitions) { auto parentPartitions = partition.GetParentPartitionIds(); @@ -266,57 +355,70 @@ void TKeyedWriteSession::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std: } // we consider here that each partition has only one parent partition - partitionToParent[partition.GetPartitionId()] = parentPartitions.front(); + partitionIdToParentId[partition.GetPartitionId()] = parentPartitions.front(); } std::vector ancestors; - std::uint32_t currentPartition = PartitionId; + std::uint32_t currentPartitionId = PartitionId; while (true) { - ancestors.push_back(currentPartition); + ancestors.push_back(currentPartitionId); - auto parentPartition = partitionToParent.find(currentPartition); - if (parentPartition == partitionToParent.end()) { + auto parentPartitionId = partitionIdToParentId.find(currentPartitionId); + if (parentPartitionId == partitionIdToParentId.end()) { break; } - currentPartition = parentPartition->second; + currentPartitionId = parentPartitionId->second; } NotReadyFutures = ancestors.size(); for (const auto& ancestor : ancestors) { - auto partitionIdx = Session->PartitionIdsMapping.find(ancestor); - Y_ABORT_UNLESS(partitionIdx != Session->PartitionIdsMapping.end(), "Partition index not found"); - auto wrappedSession = Session->SessionsWorker->GetWriteSession(partitionIdx->second, false); + auto wrappedSession = Session->SessionsWorker->GetWriteSession(ancestor, false); Y_ABORT_UNLESS(wrappedSession, "Write session not found"); WriteSessions.push_back(wrappedSession); auto future = wrappedSession->Session->GetInitSeqNo(); + std::weak_ptr session = Session->shared_from_this(); lock.unlock(); - future.Subscribe([this, wrappedSession](const NThreading::TFuture& result) { - std::lock_guard lock(Lock); - if (result.HasException()) { - LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, TStringBuilder() << "Failed to get max seq no for partition"); - TSessionClosedEvent sessionClosedEvent(EStatus::INTERNAL_ERROR, {}); - Session->GetSessionClosedEventAndDie(wrappedSession, std::move(sessionClosedEvent)); - MoveTo(EState::Done); + future.Subscribe([this, session, wrappedSession, ancestor](const NThreading::TFuture& result) { + auto sessionPtr = session.lock(); + if (!sessionPtr) { return; } - UpdateMaxSeqNo(result.GetValue()); - if (--NotReadyFutures == 0) { - MoveTo(EState::GotMaxSeqNo); + if (IsDone()) { + return; + } + + bool gotMaxSeqNo = false; + { + std::lock_guard lock(Lock); + if (result.HasException()) { + LOG_LAZY(sessionPtr->DbDriverState->Log, TLOG_ERR, sessionPtr->LogPrefix() << "Failed to get max seq no for partition " << ancestor << " for splitted partition " << PartitionId); + TSessionClosedEvent sessionClosedEvent(EStatus::INTERNAL_ERROR, {}); + sessionPtr->GetSessionClosedEventAndDie(wrappedSession, std::move(sessionClosedEvent)); + MoveTo(EState::Done); + return; + } + + UpdateMaxSeqNo(result.GetValue()); + if (--NotReadyFutures == 0) { + MoveTo(EState::GotMaxSeqNo); + gotMaxSeqNo = true; + } + } + + if (gotMaxSeqNo) { + sessionPtr->RunMainWorker(); } }); lock.lock(); GetMaxSeqNoFutures.push_back(future); } -} - -NThreading::TFuture TKeyedWriteSession::TSplittedPartitionWorker::Wait() { - if (DescribeTopicFuture.Initialized() && !DescribeTopicFuture.IsReady()) { - return DescribeTopicFuture.IgnoreResult(); + + if (ancestors.empty()) { + LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "No ancestors found for partition " << PartitionId); + MoveTo(EState::Init); } - - return NThreading::NWait::WaitAny(GetMaxSeqNoFutures); } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -325,15 +427,9 @@ NThreading::TFuture TKeyedWriteSession::TSplittedPartitionWorker::Wait() { TKeyedWriteSession::TEventsWorker::TEventsWorker(TKeyedWriteSession* session) : Session(session) { - NotReadyPromise = NThreading::NewPromise(); - NotReadyFuture = NotReadyPromise.GetFuture(); EventsPromise = NThreading::NewPromise(); EventsFuture = EventsPromise.GetFuture(); - // Initialize per-partition futures to a valid, non-ready future to avoid TFutures being uninitialized - // (NThreading::WaitAny throws on uninitialized futures). - Futures.resize(Session->Partitions.size(), NotReadyFuture); - AddReadyToAcceptEvent(); } @@ -342,15 +438,17 @@ void TKeyedWriteSession::TEventsWorker::HandleAcksEvent(std::uint64_t partition, queueIt->second.push_back(TWriteSessionEvent::TEvent(std::move(event))); } -void TKeyedWriteSession::TEventsWorker::HandleReadyToAcceptEvent(std::uint64_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event) { +void TKeyedWriteSession::TEventsWorker::HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event) { Session->MessagesWorker->HandleContinuationToken(partition, std::move(event.ContinuationToken)); } -void TKeyedWriteSession::TEventsWorker::HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint64_t partition) { +void TKeyedWriteSession::TEventsWorker::HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition) { if (event.IsSuccess()) { return; } + Session->Partitions[partition].Locked_ = true; + if (event.GetStatus() == EStatus::OVERLOADED) { Session->HandleAutoPartitioning(partition); return; @@ -362,7 +460,7 @@ void TKeyedWriteSession::TEventsWorker::HandleSessionClosedEvent(TSessionClosedE Session->NonBlockingClose(); } -bool TKeyedWriteSession::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint64_t partition) { +bool TKeyedWriteSession::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition) { while (true) { auto event = wrappedSession->Session->GetEvent(false); if (!event) { @@ -389,7 +487,7 @@ bool TKeyedWriteSession::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrap return false; } -void TKeyedWriteSession::TEventsWorker::DoWork() { +std::optional> TKeyedWriteSession::TEventsWorker::DoWork() { std::unique_lock lock(Lock); while (!ReadyFutures.empty()) { @@ -407,42 +505,59 @@ void TKeyedWriteSession::TEventsWorker::DoWork() { lock.lock(); } - if (!Session->Done.load()) { - TransferEventsToOutputQueue(); + if (!Session->Done.load() && TransferEventsToOutputQueue()) { + return EventsPromise; } + + return std::nullopt; } -void TKeyedWriteSession::TEventsWorker::SubscribeToPartition(std::uint64_t partition) { +void TKeyedWriteSession::TEventsWorker::SubscribeToPartition(std::uint32_t partition) { if (auto it = Session->SplittedPartitionWorkers.find(partition); it != Session->SplittedPartitionWorkers.end()) { - Futures[partition] = NotReadyFuture; + Session->Partitions[partition].Future(NThreading::MakeFuture()); return; } auto wrappedSession = Session->SessionsWorker->GetWriteSession(partition); auto newFuture = wrappedSession->Session->WaitEvent(); - while (partition >= Futures.size()) { - Futures.push_back(NotReadyFuture); - } - newFuture.Subscribe([this, partition](const NThreading::TFuture&) { - std::lock_guard lock(Lock); - this->ReadyFutures.insert(partition); + std::weak_ptr session = Session->shared_from_this(); + std::weak_ptr self = shared_from_this(); + + newFuture.Subscribe([self, session, partition](const NThreading::TFuture&) { + auto sessionPtr = session.lock(); + if (!sessionPtr) { + return; + } + + auto selfPtr = self.lock(); + if (!selfPtr) { + return; + } + + { + std::lock_guard lock(selfPtr->Lock); + selfPtr->ReadyFutures.insert(partition); + } + sessionPtr->RunMainWorker(); }); - Futures[partition] = newFuture; + Session->Partitions[partition].Future(newFuture); } -void TKeyedWriteSession::TEventsWorker::HandleNewMessage() { +std::optional> TKeyedWriteSession::TEventsWorker::HandleNewMessage() { std::lock_guard lock(Lock); if (Session->MessagesWorker->IsMemoryUsageOK()) { AddReadyToAcceptEvent(); + return EventsPromise; } + + return std::nullopt; } void TKeyedWriteSession::TEventsWorker::AddReadyToAcceptEvent() { EventsOutputQueue.push_back(TWriteSessionEvent::TReadyToAcceptEvent(IssueContinuationToken())); - EventsPromise.TrySetValue(); } -bool TKeyedWriteSession::TEventsWorker::AddSessionClosedEvent() { +bool TKeyedWriteSession::TEventsWorker::AddSessionClosedIfNeeded() { if (!Session->Closed.load()) { return false; } @@ -459,20 +574,23 @@ bool TKeyedWriteSession::TEventsWorker::AddSessionClosedEvent() { return false; } -void TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { +bool TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { bool eventsTransferred = false; bool shouldAddReadyToAcceptEvent = false; - std::unordered_map> acks; + std::unordered_map> acks; auto messagesWorker = Session->MessagesWorker; - auto buildOutputAckEvent = [](std::deque& acksQueue, std::optional expectedSeqNo) -> TWriteSessionEvent::TAcksEvent { + auto buildOutputAckEvent = [&](std::deque& acksQueue, std::uint64_t partition, std::optional expectedSeqNo) -> TWriteSessionEvent::TAcksEvent { TWriteSessionEvent::TAcksEvent ackEvent; if (expectedSeqNo.has_value()) { - Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo); + if (acksQueue.front().SeqNo != expectedSeqNo.value()) { + LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << partition); + } + Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << Session->Partitions[partition].PartitionId_); } - auto ack = std::move(acksQueue.front()); + auto ack = std::move(acksQueue.front()); ackEvent.Acks.push_back(std::move(ack)); acksQueue.pop_front(); return ackEvent; @@ -490,12 +608,12 @@ void TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { auto remainingAcks = acks.find(head.Partition); if (remainingAcks != acks.end() && remainingAcks->second.size() > 0) { - EventsOutputQueue.push_back(buildOutputAckEvent(remainingAcks->second, head.Message.SeqNo_)); + EventsOutputQueue.push_back(buildOutputAckEvent(remainingAcks->second, head.Partition, head.SeqNo)); finishWithAck(); continue; } - const auto& eventsQueueIt = PartitionsEventQueues.find(head.Partition); + auto eventsQueueIt = PartitionsEventQueues.find(head.Partition); if (eventsQueueIt == PartitionsEventQueues.end() || eventsQueueIt->second.empty()) { // No events for this message yet, stop processing (preserve order) break; @@ -507,7 +625,7 @@ void TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { std::deque acksQueue; std::copy(acksEvent->Acks.begin(), acksEvent->Acks.end(), std::back_inserter(acksQueue)); - EventsOutputQueue.push_back(buildOutputAckEvent(acksQueue, head.Message.SeqNo_)); + EventsOutputQueue.push_back(buildOutputAckEvent(acksQueue, head.Partition, head.SeqNo)); acks[head.Partition] = std::move(acksQueue); eventsQueueIt->second.pop_front(); eventsTransferred = true; @@ -536,22 +654,32 @@ void TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { AddReadyToAcceptEvent(); } - if (eventsTransferred) { - EventsPromise.TrySetValue(); - } + return eventsTransferred; } -std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueBegin(std::uint64_t partition) { +std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueBegin(std::uint32_t partition) { auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); return queueIt->second.begin(); } -std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueEnd(std::uint64_t partition) { +std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueEnd(std::uint32_t partition) { auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); return queueIt->second.end(); } -std::optional TKeyedWriteSession::TEventsWorker::GetEventImpl(bool block) { +TKeyedWriteSession::TEventsWorker::EEventType TKeyedWriteSession::TEventsWorker::GetEventType(const TWriteSessionEvent::TEvent& event) { + if (std::holds_alternative(event)) { + return EEventType::SessionClosed; + } else if (std::holds_alternative(event)) { + return EEventType::ReadyToAccept; + } else if (std::holds_alternative(event)) { + return EEventType::Ack; + } + + Y_ABORT_UNLESS(false, "Unexpected event type"); +} + +std::optional TKeyedWriteSession::TEventsWorker::GetEventImpl(bool block, const std::vector& eventTypes) { std::unique_lock lock(Lock); if (EventsOutputQueue.empty() && block) { lock.unlock(); @@ -560,6 +688,10 @@ std::optional TKeyedWriteSession::TEventsWorker::Get } if (!EventsOutputQueue.empty()) { + if (!eventTypes.empty() && std::find(eventTypes.begin(), eventTypes.end(), GetEventType(EventsOutputQueue.front())) == eventTypes.end()) { + return std::nullopt; + } + auto event = std::move(EventsOutputQueue.front()); EventsOutputQueue.pop_front(); return event; @@ -568,29 +700,29 @@ std::optional TKeyedWriteSession::TEventsWorker::Get return std::nullopt; } -std::optional TKeyedWriteSession::TEventsWorker::GetEvent(bool block) { +std::optional TKeyedWriteSession::TEventsWorker::GetEvent(bool block, const std::vector& eventTypes) { { std::unique_lock lock(Lock); - AddSessionClosedEvent(); + AddSessionClosedIfNeeded(); } - auto event = GetEventImpl(block); + auto event = GetEventImpl(block, eventTypes); return event; } -std::vector TKeyedWriteSession::TEventsWorker::GetEvents(bool block, std::optional maxEventsCount) { +std::vector TKeyedWriteSession::TEventsWorker::GetEvents(bool block, std::optional maxEventsCount, const std::vector& eventTypes) { if (maxEventsCount.has_value() && maxEventsCount.value() == 0) { return {}; } { std::unique_lock lock(Lock); - AddSessionClosedEvent(); + AddSessionClosedIfNeeded(); } std::vector events; while (true) { - auto event = GetEventImpl(block); + auto event = GetEventImpl(block, eventTypes); if (!event) { break; } @@ -604,14 +736,10 @@ std::vector TKeyedWriteSession::TEventsWorker::GetEv return events; } -NThreading::TFuture TKeyedWriteSession::TEventsWorker::Wait() { - return NThreading::NWait::WaitAny(Futures); -} - NThreading::TFuture TKeyedWriteSession::TEventsWorker::WaitEvent() { std::unique_lock lock(Lock); - AddSessionClosedEvent(); + AddSessionClosedIfNeeded(); if (!EventsOutputQueue.empty()) { return NThreading::MakeFuture(); } @@ -624,11 +752,9 @@ NThreading::TFuture TKeyedWriteSession::TEventsWorker::WaitEvent() { return EventsFuture; } -void TKeyedWriteSession::TEventsWorker::UnsubscribeFromPartition(std::uint64_t partition) { +void TKeyedWriteSession::TEventsWorker::UnsubscribeFromPartition(std::uint32_t partition) { ReadyFutures.erase(partition); - if (partition < Futures.size()) { - Futures[partition] = NotReadyFuture; - } + Session->Partitions[partition].Future(NThreading::MakeFuture()); } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -638,25 +764,20 @@ TKeyedWriteSession::TSessionsWorker::TSessionsWorker(TKeyedWriteSession* session : Session(session) {} -TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::GetWriteSession(std::uint64_t partition, bool directToPartition) { +TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::GetWriteSession(std::uint32_t partition, bool directToPartition) { auto sessionIter = SessionsIndex.find(partition); - if (sessionIter == SessionsIndex.end()) { - return CreateWriteSession(partition, directToPartition); - } - - if (!directToPartition) { - SessionsIndex.erase(sessionIter); + if (sessionIter == SessionsIndex.end() || !directToPartition) { return CreateWriteSession(partition, directToPartition); } return sessionIter->second; } -std::string TKeyedWriteSession::TSessionsWorker::GetProducerId(std::uint64_t partitionId) { +std::string TKeyedWriteSession::TSessionsWorker::GetProducerId(std::uint32_t partitionId) { return std::format("{}_{}", Session->Settings.ProducerIdPrefix_, partitionId); } -TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::CreateWriteSession(std::uint64_t partition, bool directToPartition) { +TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::CreateWriteSession(std::uint32_t partition, bool directToPartition) { auto partitionId = Session->Partitions[partition].PartitionId_; auto producerId = GetProducerId(partitionId); TWriteSessionSettings alteredSettings = Session->Settings; @@ -667,9 +788,9 @@ TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker:: .MaxMemoryUsage(std::numeric_limits::max()) .RetryPolicy(Session->RetryPolicy) .EventHandlers(TWriteSessionSettings::TEventHandlers() - .ReadyToAcceptHandler({}) - .AcksHandler({}) - .SessionClosedHandler({})); + .ReadyToAcceptHandler({}) + .AcksHandler({}) + .SessionClosedHandler({})); if (directToPartition) { alteredSettings.DirectWriteToPartition(true); @@ -679,9 +800,10 @@ TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker:: Session->Client->CreateWriteSession(alteredSettings), partition); - SessionsIndex.emplace(partition, writeSession); - - Session->EventsWorker->SubscribeToPartition(partition); + if (directToPartition) { + SessionsIndex.emplace(partition, writeSession); + Session->EventsWorker->SubscribeToPartition(partition); + } return writeSession; } @@ -692,17 +814,8 @@ void TKeyedWriteSession::TSessionsWorker::DestroyWriteSession(TSessionsIndexIter auto closeResult = it->second->Session->Close(closeTimeout); Y_ABORT_UNLESS(!mustBeEmpty || closeResult, "There are still messages in flight"); - const std::uint64_t partition = it->second->Partition; - if (it->second->IdleSession) { - auto itIdle = IdlerSessionsIndex.find(partition); - if (itIdle != IdlerSessionsIndex.end()) { - IdlerSessions.erase(itIdle->second); - IdlerSessionsIndex.erase(itIdle); - } - it->second->IdleSession.reset(); - } + const auto partition = it->second->Partition; it = SessionsIndex.erase(it); - Session->EventsWorker->UnsubscribeFromPartition(partition); } @@ -729,14 +842,22 @@ void TKeyedWriteSession::TSessionsWorker::OnWriteToSession(WrappedWriteSessionPt } void TKeyedWriteSession::TSessionsWorker::DoWork() { - for (auto it = IdlerSessions.begin(); it != IdlerSessions.end(); ) { + while (!IdlerSessions.empty()) { + auto it = IdlerSessions.begin(); if (!(*it)->IsExpired()) { break; } - auto sessionIter = SessionsIndex.find((*it)->Session->Partition); - ++it; + const auto partition = (*it)->Session->Partition; + + // Remove idle tracking first to keep containers consistent even if the session + // is already absent from SessionsIndex. + IdlerSessions.erase(it); + IdlerSessionsIndex.erase(partition); + + auto sessionIter = SessionsIndex.find(partition); if (sessionIter != SessionsIndex.end()) { + sessionIter->second->IdleSession.reset(); DestroyWriteSession(sessionIter, TDuration::Zero()); } } @@ -750,84 +871,87 @@ TKeyedWriteSession::TMessagesWorker::TMessagesWorker(TKeyedWriteSession* session { } -void TKeyedWriteSession::TMessagesWorker::RechoosePartitionIfNeeded(TMessageInfo& message) { - const auto& partitionInfo = Session->Partitions[message.Partition]; +void TKeyedWriteSession::TMessagesWorker::RechoosePartitionIfNeeded(MessageIter message) { + const auto& partitionInfo = Session->Partitions[message->Partition]; if (partitionInfo.Children_.empty()) { return; } // this case means that partition was split, so we need to rechoose the partition for the message - auto newPartition = Session->PartitionChooser->ChoosePartition(message.Key); - if (newPartition != message.Partition) { - message.Partition = newPartition; - } + auto newPartition = Session->PartitionChooser->ChoosePartition(message->Key); + message->Partition = newPartition; } void TKeyedWriteSession::TMessagesWorker::DoWork() { auto sessionsWorker = Session->SessionsWorker; - while (!PendingMessages.empty() && MessagesToResend.empty()) { - auto& head = PendingMessages.front(); - if (Session->Partitions[head.Partition].Locked_ || Session->SplittedPartitionWorkers.contains(head.Partition)) { - break; - } - RechoosePartitionIfNeeded(head); - auto msgToSave = head; - auto wrappedSession = sessionsWorker->GetWriteSession(msgToSave.Partition); - if (!SendMessage(wrappedSession, std::move(head))) { - break; - } + auto iterateMessagesIndex = [&](std::unordered_map>& messagesIndex, auto stopCondition) { + std::vector partitionsProcessed; + for (auto& [partition, messages] : messagesIndex) { + while (!messages.empty()) { + auto head = messages.front(); + if (stopCondition(head)) { + break; + } - PendingMessages.pop_front(); - sessionsWorker->OnWriteToSession(wrappedSession); - PushInFlightMessage(msgToSave.Partition, std::move(msgToSave)); - } + auto wrappedSession = sessionsWorker->GetWriteSession(head->Partition); + if (!SendMessage(wrappedSession, *head)) { + break; + } - std::vector partitionsConsumed; - for (auto& [partition, iter] : MessagesToResend) { - auto inFlightMessagesIndexChain = InFlightMessagesIndex.find(partition); - Y_ABORT_UNLESS(inFlightMessagesIndexChain != InFlightMessagesIndex.end(), "InFlightMessagesIndex not found"); - - while (iter != inFlightMessagesIndexChain->second.end()) { - auto msgToSend = **iter; - auto wrappedSession = sessionsWorker->GetWriteSession(msgToSend.Partition); - if (!SendMessage(wrappedSession, std::move(msgToSend))) { - break; + Session->Metrics.AddWriteLag((TInstant::Now() - head->CreateTimestamp.value_or(TInstant::Now())).MilliSeconds()); + head->Sent = true; + sessionsWorker->OnWriteToSession(wrappedSession); + messages.pop_front(); } - sessionsWorker->OnWriteToSession(wrappedSession); - ++iter; + if (messages.empty()) { + partitionsProcessed.push_back(partition); + } } - if (iter == inFlightMessagesIndexChain->second.end()) { - partitionsConsumed.push_back(partition); + for (const auto& partition : partitionsProcessed) { + messagesIndex.erase(partition); } - } + }; - for (const auto& partition : partitionsConsumed) { - MessagesToResend.erase(partition); - } + iterateMessagesIndex( + MessagesToResendIndex, + [](MessageIter) { + return false; + } + ); - if (Session->MessagesNotEmptyFuture.IsReady()) { - Session->MessagesNotEmptyPromise = NThreading::NewPromise(); - Session->MessagesNotEmptyFuture = Session->MessagesNotEmptyPromise.GetFuture(); - } + iterateMessagesIndex( + PendingMessagesIndex, + [this](MessageIter head) { + return Session->Partitions[head->Partition].Locked_ || + MessagesToResendIndex.contains(head->Partition); + } + ); } -bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wrappedSession, TMessageInfo&& message) { +bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message) { + if (!wrappedSession) { + return false; + } + auto continuationToken = GetContinuationToken(message.Partition); if (!continuationToken) { return false; } - wrappedSession->Session->Write(std::move(*continuationToken), std::move(message.Message), message.Tx); + wrappedSession->Session->Write(std::move(*continuationToken), message.BuildMessage(), message.Tx); return true; } -void TKeyedWriteSession::TMessagesWorker::PushInFlightMessage(std::uint64_t partition, TMessageInfo&& message) { - InFlightMessages.push_back(std::move(message)); - auto [listIt, _] = InFlightMessagesIndex.try_emplace(partition, std::list::iterator>()); - listIt->second.push_back(std::prev(InFlightMessages.end())); +void TKeyedWriteSession::TMessagesWorker::PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message) { + auto iter = InFlightMessages.insert(InFlightMessages.end(), std::move(message)); + auto [inFlightMessagesIndexIt, _] = InFlightMessagesIndex.try_emplace(partition, std::list()); + inFlightMessagesIndexIt->second.push_back(iter); + + auto [pendingMessagesIndexIt, __] = PendingMessagesIndex.try_emplace(partition, std::list()); + pendingMessagesIndexIt->second.push_back(iter); } void TKeyedWriteSession::TMessagesWorker::HandleAck() { @@ -853,21 +977,21 @@ void TKeyedWriteSession::TMessagesWorker::PopInFlightMessage() { } } - Y_ABORT_UNLESS(it->Message.Data.size() <= MemoryUsage, "MemoryUsage is less than the size of the message"); - MemoryUsage -= it->Message.Data.size(); + Y_ABORT_UNLESS(it->Data.size() <= MemoryUsage, "MemoryUsage is less than the size of the message"); + MemoryUsage -= it->Data.size(); InFlightMessages.pop_front(); } bool TKeyedWriteSession::TMessagesWorker::IsMemoryUsageOK() const { - return MemoryUsage <= Session->Settings.MaxMemoryUsage_; + return MemoryUsage <= Session->Settings.MaxMemoryUsage_ / 2; } -void TKeyedWriteSession::TMessagesWorker::AddMessage(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx) { - PendingMessages.push_back(TMessageInfo(key, std::move(message), partition, tx)); +void TKeyedWriteSession::TMessagesWorker::AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx) { MemoryUsage += message.Data.size(); + PushInFlightMessage(partition, TMessageInfo(key, std::move(message), partition, tx)); } -std::optional TKeyedWriteSession::TMessagesWorker::GetContinuationToken(std::uint64_t partition) { +std::optional TKeyedWriteSession::TMessagesWorker::GetContinuationToken(std::uint32_t partition) { auto it = ContinuationTokens.find(partition); if (it != ContinuationTokens.end() && !it->second.empty()) { auto token = std::move(it->second.front()); @@ -881,17 +1005,13 @@ std::optional TKeyedWriteSession::TMessagesWorker::GetContin return std::nullopt; } -void TKeyedWriteSession::TMessagesWorker::HandleContinuationToken(std::uint64_t partition, TContinuationToken&& continuationToken) { +void TKeyedWriteSession::TMessagesWorker::HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken) { auto [it, _] = ContinuationTokens.try_emplace(partition, std::deque()); it->second.push_back(std::move(continuationToken)); } -NThreading::TFuture TKeyedWriteSession::TMessagesWorker::Wait() { - return Session->MessagesNotEmptyFuture; -} - bool TKeyedWriteSession::TMessagesWorker::IsQueueEmpty() const { - return PendingMessages.empty() && InFlightMessages.empty(); + return InFlightMessages.empty(); } const TKeyedWriteSession::TMessageInfo& TKeyedWriteSession::TMessagesWorker::GetFrontInFlightMessage() const { @@ -903,7 +1023,7 @@ bool TKeyedWriteSession::TMessagesWorker::HasInFlightMessages() const { return !InFlightMessages.empty(); } -void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint64_t partition, std::uint64_t afterSeqNo) { +void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo) { auto it = InFlightMessagesIndex.find(partition); if (it == InFlightMessagesIndex.end()) { return; @@ -917,11 +1037,11 @@ void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint64_t p std::vector acksToSend; while (resendIt != list.end()) { - if (!(*resendIt)->Message.SeqNo_.has_value() || (*resendIt)->Message.SeqNo_.value() > afterSeqNo) { + if (!(*resendIt)->SeqNo.has_value() || (*resendIt)->SeqNo.value() > afterSeqNo) { break; } - auto seqNo = (*resendIt)->Message.SeqNo_.value(); + auto seqNo = (*resendIt)->SeqNo.value(); if (ackQueueIt == ackQueueEnd) { // this case can happen if the message was sent, but session was closed before the ack was received TWriteSessionEvent::TWriteAck ack; @@ -947,29 +1067,59 @@ void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint64_t p } if (!acksToSend.empty()) { - LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, TStringBuilder() << "Sending acks to partition " << partition << ": " << acksToSend.size()); TWriteSessionEvent::TAcksEvent event; event.Acks = std::move(acksToSend); Session->EventsWorker->HandleAcksEvent(partition, std::move(event)); } + // IMPORTANT: do not mutate InFlightMessagesIndex while holding references/iterators to its elements. + // try_emplace()/rehash may invalidate 'it' and 'list' -> use-after-free and segfaults. + std::vector> messagesFromOldPartition; + messagesFromOldPartition.reserve(std::distance(resendIt, list.end())); + auto currentSeqNo = resendIt != list.end() ? (*resendIt)->SeqNo.value_or(0) : 0; for (auto iter = resendIt; iter != list.end(); ++iter) { + if (iter != resendIt && currentSeqNo != 0) { + Y_ABORT_UNLESS((*iter)->SeqNo.value_or(0) > currentSeqNo, "SeqNo is not increasing for partition %d", partition); + } + auto newPartition = Session->PartitionChooser->ChoosePartition((*iter)->Key); (*iter)->Partition = newPartition; + messagesFromOldPartition.emplace_back(newPartition, *iter); - auto [listIt, _] = InFlightMessagesIndex.try_emplace(newPartition, std::list::iterator>()); - listIt->second.push_back(*iter); + currentSeqNo = (*iter)->SeqNo.value_or(0); } - - if (resendIt != list.end()) { - for (const auto& child : Session->Partitions[partition].Children_) { - if (auto childList = InFlightMessagesIndex.find(child); childList != InFlightMessagesIndex.end()) { - MessagesToResend.try_emplace(child, childList->second.begin()); - } + + list.erase(resendIt, list.end()); + for (const auto& [newPartition, msgIt] : messagesFromOldPartition) { + auto [inFlightMessagesIndexChainIt, _] = InFlightMessagesIndex.try_emplace(newPartition, std::list()); + inFlightMessagesIndexChainIt->second.push_back(msgIt); + + if (msgIt->Sent) { + auto [messagesToResendChainIt, __] = MessagesToResendIndex.try_emplace(newPartition, std::list()); + messagesToResendChainIt->second.push_back(msgIt); } + } - list.erase(resendIt, list.end()); + InFlightMessagesIndex.erase(partition); +} + +void TKeyedWriteSession::TMessagesWorker::RebuildPendingMessagesIndex(std::uint32_t partition) { + auto [oldPendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(partition, std::list()); + std::unordered_map> pendingMessagesForNewPartitions; + for (auto it = oldPendingMessagesIndexChainIt->second.begin(); it != oldPendingMessagesIndexChainIt->second.end(); ++it) { + auto newPartition = Session->PartitionChooser->ChoosePartition((*it)->Key); + auto [pendingMessagesForNewPartitionsIt, __] = pendingMessagesForNewPartitions.try_emplace(newPartition, std::list()); + pendingMessagesForNewPartitionsIt->second.push_back(*it); + } + + for (const auto& [newPartition, pendingMessagesForNewPartition] : pendingMessagesForNewPartitions) { + auto [pendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(newPartition, std::list()); + for (auto reverseIt = pendingMessagesForNewPartition.rbegin(); reverseIt != pendingMessagesForNewPartition.rend(); ++reverseIt) { + pendingMessagesIndexChainIt->second.push_front(*reverseIt); + } } + + PendingMessagesIndex.erase(partition); } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1006,6 +1156,53 @@ typename TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::IRetryState::TPtr TK return std::make_unique(Session); } +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TKeyedWriteSession::Metrics + +void TKeyedWriteSession::MetricGauge::Add(std::uint64_t value) { + Sum += value; + MetricCount++; +} + +void TKeyedWriteSession::MetricGauge::Clear() { + Sum = 0; + MetricCount = 0; +} + +long double TKeyedWriteSession::MetricGauge::Average() { + if (MetricCount == 0) { + return 0; + } + + return (long double)Sum / (long double)MetricCount; +} + +TKeyedWriteSession::Metrics::Metrics(TKeyedWriteSession* session): Session(session) {} + +void TKeyedWriteSession::Metrics::AddMainWorkerTime(std::uint64_t ms) { + std::lock_guard lock(Lock); + MainWorkerTimeMs.Add(ms); +} + +void TKeyedWriteSession::Metrics::AddCycleTime(std::uint64_t ms) { + std::lock_guard lock(Lock); + CycleTimeMs.Add(ms); +} + +void TKeyedWriteSession::Metrics::AddWriteLag(std::uint64_t lagMs) { + std::lock_guard lock(Lock); + WriteLagMs.Add(lagMs); +} + +void TKeyedWriteSession::Metrics::PrintMetrics() { + std::lock_guard lock(Lock); + LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "METRICS: MainWorkerTimeMs: " << MainWorkerTimeMs.Average() << " ms, CycleTimeMs: " << CycleTimeMs.Average() << " ms, WriteLagMs: " << WriteLagMs.Average() << " ms"); + MainWorkerTimeMs.Clear(); + CycleTimeMs.Clear(); + WriteLagMs.Clear(); +} + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TKeyedWriteSession @@ -1017,6 +1214,7 @@ TKeyedWriteSession::TKeyedWriteSession( : Connections(connections), Client(client), DbDriverState(dbDriverState), + Metrics(this), Settings(settings) { if (settings.ProducerIdPrefix_.empty()) { @@ -1033,28 +1231,59 @@ TKeyedWriteSession::TKeyedWriteSession( TDescribeTopicSettings describeTopicSettings; auto topicConfig = client->DescribeTopic(settings.Path_, describeTopicSettings).GetValueSync(); - const auto& partitions = topicConfig.GetTopicDescription().GetPartitions(); - auto partitionChooserStrategy = settings.PartitionChooserStrategy_; + auto partitions = topicConfig.GetTopicDescription().GetPartitions(); + std::sort(partitions.begin(), partitions.end(), [](const auto& a, const auto& b) -> bool { + return a.GetPartitionId() < b.GetPartitionId(); + }); + auto partitionChooserStrategy = settings.PartitionChooserStrategy_; auto strategy = topicConfig.GetTopicDescription().GetPartitioningSettings().GetAutoPartitioningSettings().GetStrategy(); auto autoPartitioningEnabled = (strategy != EAutoPartitioningStrategy::Disabled && strategy != EAutoPartitioningStrategy::Unspecified); for (const auto& partition : partitions) { - if (!partition.GetActive()) { - continue; - } - auto partitionId = partition.GetPartitionId(); - PartitionIdsMapping[partitionId] = Partitions.size(); - Partitions.push_back( - TPartitionInfo() + auto fromBound = partition.GetFromBound().value_or(""); + auto toBound = partition.GetToBound(); + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Adding partition " << partitionId << " from bound " << fromBound << " to bound " << (toBound.has_value() ? toBound.value() : "null")); + Partitions[partitionId] = TPartitionInfo() .PartitionId(partitionId) - .FromBound(partition.GetFromBound().value_or("")) - .ToBound(partition.GetToBound())); + .FromBound(fromBound) + .ToBound(toBound); } - Y_ABORT_UNLESS(!Partitions.empty(), "No active partitions found"); + for (const auto& partition : partitions) { + auto children = partition.GetChildPartitionIds(); + + std::vector childrenIndices; + childrenIndices.reserve(children.size()); + for (auto child : children) { + childrenIndices.push_back(child); + } + Partitions[partition.GetPartitionId()].Children(childrenIndices); + } + + if (Settings.EventHandlers_.CommonHandler_) { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); + } else { + if (!Settings.EventHandlers_.SessionClosedHandler_) { + EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::SessionClosed); + } else { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); + } + if (!Settings.EventHandlers_.ReadyToAcceptHandler_) { + EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); + } else { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); + } + if (!Settings.EventHandlers_.AcksHandler_) { + EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::Ack); + } else { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); + } + } switch (partitionChooserStrategy) { case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound: @@ -1065,7 +1294,11 @@ TKeyedWriteSession::TKeyedWriteSession( ythrow TContractViolation("Unbounded partition is not supported for Bound partition chooser strategy"); } - PartitionsIndex[Partitions[i].FromBound_] = i; + if (!Partitions[i].Children_.empty()) { + continue; + } + + PartitionsIndex[Partitions[i].FromBound_] = Partitions[i].PartitionId_; } break; case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash: @@ -1073,7 +1306,13 @@ TKeyedWriteSession::TKeyedWriteSession( ythrow TContractViolation("Hash partition chooser strategy is not supported for topic with auto partitioning"); } - PartitionChooser = std::make_unique(this); + std::vector partitionsIds; + partitionsIds.reserve(partitions.size()); + for (const auto& partition : partitions) { + partitionsIds.push_back(partition.GetPartitionId()); + } + + PartitionChooser = std::make_unique(std::move(partitionsIds)); break; } @@ -1081,8 +1320,6 @@ TKeyedWriteSession::TKeyedWriteSession( CloseFuture = ClosePromise.GetFuture(); ShutdownPromise = NThreading::NewPromise(); ShutdownFuture = ShutdownPromise.GetFuture(); - MessagesNotEmptyPromise = NThreading::NewPromise(); - MessagesNotEmptyFuture = MessagesNotEmptyPromise.GetFuture(); SessionsWorker = std::make_shared(this); MessagesWorker = std::make_shared(this); @@ -1092,18 +1329,26 @@ TKeyedWriteSession::TKeyedWriteSession( // Start handlers executor for user callbacks (Acks/ReadyToAccept/SessionClosed/Common). Settings.EventHandlers_.HandlersExecutor_->Start(); - RunUserEventLoop(); - NextFuture = Next(false); - NextFuture.Subscribe([this](const NThreading::TFuture&) { + CloseFuture.Subscribe([this](const NThreading::TFuture&) { RunMainWorker(); }); + + RunMainWorker(); + + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Keyed write session created"); } -const std::vector& TKeyedWriteSession::GetPartitions() const { - return Partitions; +std::vector TKeyedWriteSession::GetPartitions() const { + std::vector partitions; + partitions.reserve(Partitions.size()); + for (const auto& [partitionId, partitionInfo] : Partitions) { + partitions.push_back(partitionInfo); + } + return partitions; } void TKeyedWriteSession::Write(TContinuationToken&&, const std::string& key, TWriteMessage&& message, TTransactionBase* tx) { + std::optional> eventsPromise; { std::lock_guard lock(GlobalLock); if (Closed.load()) { @@ -1121,10 +1366,14 @@ void TKeyedWriteSession::Write(TContinuationToken&&, const std::string& key, TWr auto partition = PartitionChooser->ChoosePartition(key); MessagesWorker->AddMessage(key, std::move(message), partition, tx); - EventsWorker->HandleNewMessage(); + eventsPromise = EventsWorker->HandleNewMessage(); + RunUserEventLoop(); } - MessagesNotEmptyPromise.TrySetValue(); + RunMainWorker(); + if (eventsPromise) { + eventsPromise->TrySetValue(); + } } bool TKeyedWriteSession::Close(TDuration closeTimeout) { @@ -1165,11 +1414,19 @@ NThreading::TFuture TKeyedWriteSession::WaitEvent() { } std::optional TKeyedWriteSession::GetEvent(bool block) { - return EventsWorker->GetEvent(block); + if (Settings.EventHandlers_.CommonHandler_) { + return std::nullopt; + } + + return EventsWorker->GetEvent(block, EventTypesWithoutHandlers); } std::vector TKeyedWriteSession::GetEvents(bool block, std::optional maxEventsCount) { - return EventsWorker->GetEvents(block, maxEventsCount); + if (Settings.EventHandlers_.CommonHandler_) { + return {}; + } + + return EventsWorker->GetEvents(block, maxEventsCount, EventTypesWithoutHandlers); } TDuration TKeyedWriteSession::GetCloseTimeout() { @@ -1181,47 +1438,30 @@ TDuration TKeyedWriteSession::GetCloseTimeout() { return CloseDeadline - now; } -void TKeyedWriteSession::RunSplittedPartitionWorkers() { +bool TKeyedWriteSession::RunSplittedPartitionWorkers() { if (SplittedPartitionWorkers.empty()) { - return; + return false; } - std::vector toRemove; + bool needRerun = false; for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { if (splittedPartitionWorker->IsDone()) { - toRemove.push_back(partition); continue; } splittedPartitionWorker->DoWork(); + needRerun = needRerun || splittedPartitionWorker->IsInit(); + needRerun = needRerun || splittedPartitionWorker->IsDone(); } - for (const auto& partition : toRemove) { - SplittedPartitionWorkers.erase(partition); - } -} - -NThreading::TFuture TKeyedWriteSession::Next(bool isClosed) { - std::vector> futures{ - EventsWorker->Wait(), - MessagesWorker->Wait() - }; - - for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { - futures.push_back(splittedPartitionWorker->Wait()); - } - - if (!isClosed) { - futures.push_back(CloseFuture); - } - - return NThreading::NWait::WaitAny(futures); + return needRerun; } void TKeyedWriteSession::RunUserEventLoop() { if (!Settings.EventHandlers_.AcksHandler_ && !Settings.EventHandlers_.ReadyToAcceptHandler_ && - !Settings.EventHandlers_.SessionClosedHandler_) { + !Settings.EventHandlers_.SessionClosedHandler_ && + !Settings.EventHandlers_.CommonHandler_) { return; } @@ -1231,7 +1471,7 @@ void TKeyedWriteSession::RunUserEventLoop() { } while (true) { - auto event = GetEvent(false); + auto event = EventsWorker->GetEvent(false, EventTypesWithHandlers); if (!event) { break; } @@ -1242,7 +1482,7 @@ void TKeyedWriteSession::RunUserEventLoop() { [this, ev = std::move(*readyToAcceptEvent)]() mutable { Settings.EventHandlers_.ReadyToAcceptHandler_(ev); }); - } else { + } else if (Settings.EventHandlers_.CommonHandler_) { handlersExecutor->Post( [this, ev = std::move(*event)]() mutable { Settings.EventHandlers_.CommonHandler_(ev); @@ -1257,7 +1497,7 @@ void TKeyedWriteSession::RunUserEventLoop() { [this, ev = std::move(*acksEvent)]() mutable { Settings.EventHandlers_.AcksHandler_(ev); }); - } else { + } else if (Settings.EventHandlers_.CommonHandler_) { handlersExecutor->Post( [this, ev = std::move(*event)]() mutable { Settings.EventHandlers_.CommonHandler_(ev); @@ -1297,39 +1537,119 @@ void TKeyedWriteSession::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrap } } - if (!receivedSessionClosedEvent) { - LOG_LAZY(DbDriverState->Log, TLOG_ERR, TStringBuilder() << "Failed to get session closed event"); + if (!receivedSessionClosedEvent || receivedSessionClosedEvent->GetStatus() == EStatus::SUCCESS || receivedSessionClosedEvent->GetStatus() == EStatus::OVERLOADED) { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Failed to get session closed event"); EventsWorker->HandleSessionClosedEvent(std::move(*sessionClosedEvent), wrappedSession->Partition); } else { EventsWorker->HandleSessionClosedEvent(std::move(*receivedSessionClosedEvent), wrappedSession->Partition); } } +TStringBuilder TKeyedWriteSession::LogPrefix() { + return TStringBuilder() << " SessionId: " << Settings.SessionId_ << " Epoch: " << Epoch.load() << " "; +} + +void TKeyedWriteSession::NextEpoch() { + auto maxEpoch = MAX_EPOCH - 1; + if (Epoch.compare_exchange_weak(maxEpoch, 0)) { + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Epoch overflow, resetting to 0"); + return; + } + + Epoch.fetch_add(1); +} + void TKeyedWriteSession::RunMainWorker() { - RunSplittedPartitionWorkers(); - { - std::unique_lock lock(GlobalLock); - EventsWorker->DoWork(); - if (!Done.load()) { - SessionsWorker->DoWork(); - MessagesWorker->DoWork(); + // This function is both "request to run" and the runner itself. + // We must handle two properties: + // - TFuture::Subscribe may call back synchronously when future is already ready. + // - A callback may race with the runner trying to go idle (avoid lost wakeups). + enum : std::uint8_t { + Idle = 0, + Running = 1, + Rerun = 2, + }; + + // Try to become the runner. If already running, just request a rerun. + std::uint8_t state = MainWorkerState.load(std::memory_order_acquire); + for (;;) { + if (state & Running) { + if (MainWorkerState.compare_exchange_weak(state, std::uint8_t(state | Rerun), + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return; + } + continue; + } else { + if (MainWorkerState.compare_exchange_weak(state, Running, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + break; // we are the runner now + } + continue; } } - RunUserEventLoop(); - auto isClosed = Closed.load(); - auto closeTimeout = GetCloseTimeout(); - if (isClosed && (Done.load() || MessagesWorker->IsQueueEmpty() || closeTimeout == TDuration::Zero())) { - ShutdownPromise.TrySetValue(); - EventsWorker->EventsPromise.TrySetValue(); - ClosePromise.TrySetValue(); - return; - } + NextEpoch(); + + auto startWorkerTime = TInstant::Now(); + // Runner loop: process, arm subscription, then either go idle or loop again. + for (;;) { + auto startIter = TInstant::Now(); + // Clear rerun request for this iteration. + MainWorkerState.fetch_and(std::uint8_t(~Rerun), std::memory_order_acq_rel); + bool needRerun = false; + std::optional> eventsPromise; + + { + std::unique_lock lock(GlobalLock); + eventsPromise = EventsWorker->DoWork(); + RunUserEventLoop(); + needRerun = RunSplittedPartitionWorkers(); + if (!Done.load()) { + SessionsWorker->DoWork(); + MessagesWorker->DoWork(); + } + } - NextFuture = Next(isClosed); - NextFuture.Subscribe([this](const NThreading::TFuture&) { - RunMainWorker(); - }); + if (eventsPromise) { + eventsPromise->TrySetValue(); + } + + const auto isClosed = Closed.load(); + const auto closeTimeout = GetCloseTimeout(); + if (isClosed && (Done.load() || MessagesWorker->IsQueueEmpty() || closeTimeout == TDuration::Zero())) { + ShutdownPromise.TrySetValue(); + EventsWorker->EventsPromise.TrySetValue(); + ClosePromise.TrySetValue(); + MainWorkerState.store(Idle, std::memory_order_release); + return; + } + + if (needRerun) { + // we need this case to start resending messages if there are any + Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); + continue; + } + + // Try to go idle. If someone requested rerun concurrently, keep running. + std::uint8_t cur = MainWorkerState.load(std::memory_order_acquire); + for (;;) { + if (cur & Rerun) { + Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); + break; // continue outer loop + } + if (MainWorkerState.compare_exchange_weak(cur, Idle, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + auto workerFinished = TInstant::Now(); + Metrics.AddCycleTime((workerFinished - startIter).MilliSeconds()); + Metrics.AddMainWorkerTime((workerFinished - startWorkerTime).MilliSeconds()); + return; // successfully went idle + } + } + // Rerun was requested; continue the loop without recursion. + } } TInstant TKeyedWriteSession::GetCloseDeadline() { @@ -1337,17 +1657,16 @@ TInstant TKeyedWriteSession::GetCloseDeadline() { return CloseDeadline; } -void TKeyedWriteSession::HandleAutoPartitioning(std::uint64_t partition) { - auto splittedPartitionWorker = std::make_shared(this, Partitions[partition].PartitionId_, partition); +void TKeyedWriteSession::HandleAutoPartitioning(std::uint32_t partition) { + auto splittedPartitionWorker = std::make_shared(this, partition); SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); } -std::string TKeyedWriteSession::GetProducerId(std::uint64_t partition) { +std::string TKeyedWriteSession::GetProducerId(std::uint32_t partition) { return std::format("{}_{}", Settings.ProducerIdPrefix_, partition); } TWriterCounters::TPtr TKeyedWriteSession::GetCounters() { - // what should we return here? return nullptr; } @@ -1367,14 +1686,13 @@ std::uint32_t TKeyedWriteSession::TBoundPartitionChooser::ChoosePartition(const return std::prev(lowerBound)->second; } -TKeyedWriteSession::THashPartitionChooser::THashPartitionChooser(TKeyedWriteSession* session) - : Session(session) -{ -} +TKeyedWriteSession::THashPartitionChooser::THashPartitionChooser(std::vector&& partitions) + : Partitions(std::move(partitions)) +{} std::uint32_t TKeyedWriteSession::THashPartitionChooser::ChoosePartition(const std::string_view key) { - std::uint64_t hash = MurmurHash(key.data(), key.size()); - return hash % Session->Partitions.size(); + auto hash = MurmurHash(key.data(), key.size()); + return Partitions[hash % Partitions.size()]; } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/src/client/topic/impl/write_session.h b/src/client/topic/impl/write_session.h index d28df6435b8..61e314129a0 100644 --- a/src/client/topic/impl/write_session.h +++ b/src/client/topic/impl/write_session.h @@ -81,21 +81,25 @@ class TKeyedWriteSession : public IKeyedWriteSession, FLUENT_SETTING(std::uint32_t, PartitionId); FLUENT_SETTING(std::vector, Children); FLUENT_SETTING_DEFAULT(bool, Locked, false); + FLUENT_SETTING_DEFAULT(NThreading::TFuture, Future, NThreading::MakeFuture()); }; struct TMessageInfo { - TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx) - : Key(key) - , Message(std::move(message)) - , Partition(partition) - , Tx(tx) - {} + TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx); std::string Key; - TWriteMessage Message; - std::uint32_t Partition; + std::string Data; + std::optional Codec; + uint32_t OriginalSize = 0; + std::optional SeqNo; + std::optional CreateTimestamp; + TMessageMeta MessageMeta; + std::optional> TxInMessage; TTransactionBase* Tx; - bool MovedToNewPartition = false; + std::uint32_t Partition; + bool Sent = false; + + TWriteMessage BuildMessage() const; }; struct TIdleSession; @@ -106,7 +110,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::uint64_t QueueSize = 0; std::shared_ptr IdleSession = nullptr; - TWriteSessionWrapper(WriteSessionPtr session, std::uint64_t partition); + TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition); bool IsQueueEmpty() const; bool AddToQueue(std::uint64_t delta); @@ -164,26 +168,26 @@ class TKeyedWriteSession : public IKeyedWriteSession, struct TSessionsWorker { TSessionsWorker(TKeyedWriteSession* session); - WrappedWriteSessionPtr GetWriteSession(std::uint64_t partition, bool directToPartition = true); + WrappedWriteSessionPtr GetWriteSession(std::uint32_t partition, bool directToPartition = true); void OnReadFromSession(WrappedWriteSessionPtr wrappedSession); void OnWriteToSession(WrappedWriteSessionPtr wrappedSession); void DoWork(); private: void AddIdleSession(WrappedWriteSessionPtr wrappedSession, TInstant emptySince, TDuration idleTimeout); - void RemoveIdleSession(std::uint64_t partition); - WrappedWriteSessionPtr CreateWriteSession(std::uint64_t partition, bool directToPartition = true); + void RemoveIdleSession(std::uint32_t partition); + WrappedWriteSessionPtr CreateWriteSession(std::uint32_t partition, bool directToPartition = true); - using TSessionsIndexIterator = std::unordered_map::iterator; + using TSessionsIndexIterator = std::unordered_map::iterator; void DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout, bool mustBeEmpty = true); - std::string GetProducerId(std::uint64_t partitionId); + std::string GetProducerId(std::uint32_t partitionId); TKeyedWriteSession* Session; std::set IdlerSessions; using IdlerSessionsIterator = std::set::iterator; - std::unordered_map IdlerSessionsIndex; - std::unordered_map SessionsIndex; + std::unordered_map IdlerSessionsIndex; + std::unordered_map SessionsIndex; }; struct TMessagesWorker { @@ -191,45 +195,47 @@ class TKeyedWriteSession : public IKeyedWriteSession, void DoWork(); - void AddMessage(const std::string& key, TWriteMessage&& message, std::uint64_t partition, TTransactionBase* tx); - void ScheduleResendMessages(std::uint64_t partition, std::uint64_t afterSeqNo); + void AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx); + void ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo); + void RebuildPendingMessagesIndex(std::uint32_t partition); void HandleAck(); - void HandleContinuationToken(std::uint64_t partition, TContinuationToken&& continuationToken); + void HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken); bool IsMemoryUsageOK() const; - NThreading::TFuture Wait(); bool IsQueueEmpty() const; bool HasInFlightMessages() const; const TMessageInfo& GetFrontInFlightMessage() const; private: - void PushInFlightMessage(std::uint64_t partition, TMessageInfo&& message); + using MessageIter = std::list::iterator; + + void PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message); void PopInFlightMessage(); - bool SendMessage(WrappedWriteSessionPtr wrappedSession, TMessageInfo&& message); - std::optional GetContinuationToken(std::uint64_t partition); - void RechoosePartitionIfNeeded(TMessageInfo& message); + bool SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message); + std::optional GetContinuationToken(std::uint32_t partition); + void RechoosePartitionIfNeeded(MessageIter message); TKeyedWriteSession* Session; - std::list PendingMessages; std::list InFlightMessages; - std::unordered_map::iterator>> InFlightMessagesIndex; - - using InFlightMessagesIndexIter = std::list::iterator>::iterator; - std::unordered_map MessagesToResend; - std::unordered_map> ContinuationTokens; + std::unordered_map> InFlightMessagesIndex; + std::unordered_map> PendingMessagesIndex; + std::unordered_map> MessagesToResendIndex; + std::unordered_map> ContinuationTokens; std::uint64_t MemoryUsage = 0; + + friend class TKeyedWriteSession; }; - struct TSplittedPartitionWorker { + struct TSplittedPartitionWorker : public std::enable_shared_from_this { private: enum class EState { - Init, - PendingDescribe, - GotDescribe, - PendingMaxSeqNo, - GotMaxSeqNo, - Done, + Init = 0, + PendingDescribe = 1, + GotDescribe = 2, + PendingMaxSeqNo = 3, + GotMaxSeqNo = 4, + Done = 5, }; void MoveTo(EState state); @@ -238,58 +244,62 @@ class TKeyedWriteSession : public IKeyedWriteSession, void HandleDescribeResult(); public: - TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId, std::uint64_t partitionIdx); + TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId); void DoWork(); - NThreading::TFuture Wait(); bool IsDone(); + bool IsInit(); + std::string GetStateName() const; private: TKeyedWriteSession* Session; NThreading::TFuture DescribeTopicFuture; EState State = EState::Init; std::uint32_t PartitionId; - std::uint64_t PartitionIdx; std::uint64_t MaxSeqNo = 0; std::vector WriteSessions; std::vector> GetMaxSeqNoFutures; std::mutex Lock; std::uint64_t NotReadyFutures = 0; + size_t Retries = 0; }; - struct TEventsWorker { + struct TEventsWorker : public std::enable_shared_from_this { + enum class EEventType { + SessionClosed = 0, + ReadyToAccept = 1, + Ack = 2, + }; + TEventsWorker(TKeyedWriteSession* session); - void DoWork(); - NThreading::TFuture Wait(); + std::optional> DoWork(); NThreading::TFuture WaitEvent(); - void UnsubscribeFromPartition(std::uint64_t partition); - void SubscribeToPartition(std::uint64_t partition); - void HandleNewMessage(); + void UnsubscribeFromPartition(std::uint32_t partition); + void SubscribeToPartition(std::uint32_t partition); + std::optional> HandleNewMessage(); void HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event); - std::optional GetEvent(bool block); - std::vector GetEvents(bool block, std::optional maxEventsCount = std::nullopt); - std::list::iterator AckQueueBegin(std::uint64_t partition); - std::list::iterator AckQueueEnd(std::uint64_t partition); + std::optional GetEvent(bool block, const std::vector& eventTypes = {}); + std::vector GetEvents(bool block, std::optional maxEventsCount = std::nullopt, const std::vector& eventTypes = {}); + std::list::iterator AckQueueBegin(std::uint32_t partition); + std::list::iterator AckQueueEnd(std::uint32_t partition); private: - void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint64_t partition); - void HandleReadyToAcceptEvent(std::uint64_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); - bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint64_t partition); - void TransferEventsToOutputQueue(); + void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition); + void HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); + bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition); + bool TransferEventsToOutputQueue(); void AddReadyToAcceptEvent(); - bool AddSessionClosedEvent(); - std::optional GetEventImpl(bool block); + bool AddSessionClosedIfNeeded(); + std::optional GetEventImpl(bool block, const std::vector& eventTypes = {}); + EEventType GetEventType(const TWriteSessionEvent::TEvent& event); TKeyedWriteSession* Session; - std::vector> Futures; - std::unordered_set ReadyFutures; - std::unordered_map> PartitionsEventQueues; + std::unordered_set ReadyFutures; + std::unordered_map> PartitionsEventQueues; std::list EventsOutputQueue; std::mutex Lock; - NThreading::TFuture NotReadyFuture; - NThreading::TPromise NotReadyPromise; NThreading::TPromise EventsPromise; NThreading::TFuture EventsFuture; @@ -314,10 +324,36 @@ class TKeyedWriteSession : public IKeyedWriteSession, }; struct THashPartitionChooser : IPartitionChooser { - THashPartitionChooser(TKeyedWriteSession* session); + THashPartitionChooser(std::vector&& partitions); std::uint32_t ChoosePartition(const std::string_view key) override; private: + std::vector Partitions; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + + struct MetricGauge { + std::uint64_t MetricCount = 0; + std::uint64_t Sum = 0; + + long double Average(); + void Add(std::uint64_t value); + void Clear(); + }; + + struct Metrics { + Metrics(TKeyedWriteSession* session); + + MetricGauge MainWorkerTimeMs; + MetricGauge CycleTimeMs; + MetricGauge WriteLagMs; + std::mutex Lock; TKeyedWriteSession* Session; + + void AddMainWorkerTime(std::uint64_t ms); + void AddCycleTime(std::uint64_t ms); + void AddWriteLag(std::uint64_t lagMs); + void PrintMetrics(); }; //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -330,13 +366,11 @@ class TKeyedWriteSession : public IKeyedWriteSession, TDuration GetCloseTimeout(); - std::string GetProducerId(std::uint64_t partition); + std::string GetProducerId(std::uint32_t partition); - void HandleAutoPartitioning(std::uint64_t partition); + void HandleAutoPartitioning(std::uint32_t partition); - void RunSplittedPartitionWorkers(); - - NThreading::TFuture Next(bool isClosed); + bool RunSplittedPartitionWorkers(); void RunUserEventLoop(); @@ -344,6 +378,10 @@ class TKeyedWriteSession : public IKeyedWriteSession, void GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent = std::nullopt); + TStringBuilder LogPrefix(); + + void NextEpoch(); + public: TKeyedWriteSession(const TKeyedWriteSessionSettings& settings, std::shared_ptr client, @@ -363,7 +401,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, TWriterCounters::TPtr GetCounters() override; - const std::vector& GetPartitions() const; + std::vector GetPartitions() const; ~TKeyedWriteSession(); @@ -372,8 +410,9 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::shared_ptr Client; TDbDriverStatePtr DbDriverState; - std::vector Partitions; - std::unordered_map PartitionIdsMapping; + Metrics Metrics; + + std::unordered_map Partitions; std::map PartitionsIndex; TKeyedWriteSessionSettings Settings; @@ -381,16 +420,13 @@ class TKeyedWriteSession : public IKeyedWriteSession, NThreading::TPromise ClosePromise; NThreading::TFuture CloseFuture; - NThreading::TFuture NextFuture; NThreading::TPromise ShutdownPromise; NThreading::TFuture ShutdownFuture; - NThreading::TPromise MessagesNotEmptyPromise; - NThreading::TFuture MessagesNotEmptyFuture; std::mutex GlobalLock; std::atomic_bool Closed = false; std::atomic_bool Done = false; - TInstant CloseDeadline = TInstant::Now(); + TInstant CloseDeadline = TInstant::Now(); std::unique_ptr PartitionChooser; @@ -398,9 +434,19 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::shared_ptr EventsWorker; std::shared_ptr SessionsWorker; - std::unordered_map> SplittedPartitionWorkers; + std::unordered_map> SplittedPartitionWorkers; std::shared_ptr MessagesWorker; std::shared_ptr RetryPolicy; + + // TFuture::Subscribe may invoke callback synchronously when the future is already ready. + // Also, callbacks may arrive concurrently with the attempt to go idle. + // Use a small state machine to avoid re-entrancy and lost wakeups. + std::atomic MainWorkerState = 0; + std::atomic Epoch = 0; + static constexpr size_t MAX_EPOCH = 1'000'000'000; + + std::vector EventTypesWithoutHandlers; + std::vector EventTypesWithHandlers; }; //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index b0f04b02117..c43f79a8697 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -1004,7 +1004,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < messages; ++i) { auto token = getReadyToken(); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - TWriteMessage msg("payload"); + std::string payload = "payload"; + TWriteMessage msg(payload); msg.SeqNo(i + 1); session->Write(std::move(*token), "key-" + ToString(i), std::move(msg)); } @@ -1048,7 +1049,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TKeyedWriteSessionEventLoop eventLoop(session); auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - TWriteMessage msg("msg0"); + std::string payload = "msg0"; + TWriteMessage msg(payload); msg.SeqNo(0); session->Write(std::move(*token), "key", std::move(msg)); @@ -1102,14 +1104,16 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - TWriteMessage msg("msg0"); + std::string payload = "msg0"; + TWriteMessage msg(payload); msg.SeqNo(seqNo++); session->Write(std::move(*token), key0, std::move(msg)); } for (ui64 i = 0; i < count1; ++i) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - TWriteMessage msg("msg1"); + std::string payload = "msg1"; + TWriteMessage msg(payload); msg.SeqNo(seqNo++); session->Write(std::move(*token), key1, std::move(msg)); } @@ -1192,7 +1196,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - TWriteMessage msg("msg"); + std::string payload = "msg"; + TWriteMessage msg(payload); msg.SeqNo(i + 1); session->Write(std::move(*token), key, std::move(msg)); } @@ -1238,7 +1243,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto key = CreateGuidAsString(); auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - auto msg = TWriteMessage("data"); + std::string payload = "data"; + TWriteMessage msg(payload); msg.SeqNo(i); session->Write(std::move(*token), key, std::move(msg)); } @@ -1282,7 +1288,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); const ui64 seqNo = nextSeqNo.fetch_add(1); - auto msg = TWriteMessage("data"); + std::string payload = "data"; + TWriteMessage msg(payload); msg.SeqNo(seqNo); session->Write(std::move(*token), key, std::move(msg)); } @@ -1317,7 +1324,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto key = CreateGuidAsString(); auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - auto msg = TWriteMessage("data"); + std::string payload = "data"; + TWriteMessage msg(payload); msg.SeqNo(seqNo++); session->Write(std::move(*token), key, std::move(msg)); } @@ -1331,7 +1339,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto key = CreateGuidAsString(); auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - auto msg = TWriteMessage("data"); + std::string payload = "data"; + TWriteMessage msg(payload); msg.SeqNo(seqNo++); session->Write(std::move(*token), key, std::move(msg)); } @@ -1367,7 +1376,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); auto key = CreateGuidAsString(); - auto msg = TWriteMessage("data"); + std::string payload = "data"; + TWriteMessage msg(payload); msg.SeqNo(i); session->Write(std::move(*token), key, std::move(msg)); } @@ -1379,6 +1389,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { NKikimr::NPQ::NTest::SplitPartition(setup, ++txId, 0, "a"); }); + writer.join(); + splitter.join(); + UNIT_ASSERT(eventLoop.WaitForAcks(messages, TDuration::Seconds(60))); eventLoop.CheckAcksOrder(); UNIT_ASSERT(session->Close(TDuration::Seconds(30))); @@ -1406,14 +1419,16 @@ Y_UNIT_TEST_SUITE(BasicUsage) { // Write several messages with different keys size_t seqNo = 1; for (int i = 0; i < 5; ++i) { - TWriteMessage msg("message1-" + ToString(i)); + std::string payload = "message1-" + ToString(i); + TWriteMessage msg(payload); msg.SeqNo(seqNo++); bool res = session->Write(key1, std::move(msg)); UNIT_ASSERT(res); } for (int i = 0; i < 5; ++i) { - TWriteMessage msg("message2-" + ToString(i)); + std::string payload = "message2-" + ToString(i); + TWriteMessage msg(payload); msg.SeqNo(seqNo++); bool res = session->Write(key2, std::move(msg)); UNIT_ASSERT(res); @@ -1440,7 +1455,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const ui64 messages = 10; for (ui64 i = 0; i < messages; ++i) { - TWriteMessage msg("payload-" + ToString(i)); + std::string payload = "payload-" + ToString(i); + TWriteMessage msg(payload); bool res = session->Write("key-" + ToString(i % 3), std::move(msg)); UNIT_ASSERT(res); } @@ -1469,7 +1485,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < 1000; ++i) { auto key = CreateGuidAsString(); - TWriteMessage msg("payload-" + ToString(seqNo)); + std::string payload = "payload-" + ToString(seqNo); + TWriteMessage msg(payload); msg.SeqNo(seqNo++); bool res = session->Write(key, std::move(msg)); UNIT_ASSERT(res); @@ -1500,7 +1517,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (int i = 0; i < 1000; ++i) { auto token = eventLoop.GetContinuationToken(TDuration::Seconds(10)); UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - TWriteMessage msg("message-" + ToString(i)); + std::string payload = "message-" + ToString(i); + TWriteMessage msg(payload); msg.SeqNo(i + 1); session->Write(std::move(*token), "key1", std::move(msg)); } diff --git a/src/client/topic/ut/ut_utils/event_loop.cpp b/src/client/topic/ut/ut_utils/event_loop.cpp index f6130422adc..233b2863e71 100644 --- a/src/client/topic/ut/ut_utils/event_loop.cpp +++ b/src/client/topic/ut/ut_utils/event_loop.cpp @@ -75,8 +75,17 @@ void TKeyedWriteSessionEventLoop::CheckAcksOrder() { std::lock_guard lock(Lock_); size_t expectedAck = 1; UNIT_ASSERT_C(AckedSeqNos_.size() == AckOrder_.size(), TStringBuilder() << "Unexpected number of acks: got " << AckOrder_.size() << ", expected " << AckedSeqNos_.size()); + size_t index = 0; for (const auto& ack : AckOrder_) { - UNIT_ASSERT_VALUES_EQUAL_C(ack, expectedAck, TStringBuilder() << "Unexpected ack order: got " << ack << ", expected " << expectedAck); + TStringBuilder sb; + if (ack != expectedAck) { + for (size_t i = std::min(size_t(0), index - 10); i < std::min(index + 10, AckOrder_.size()); i++) { + sb << "Ack " << i << ": " << AckOrder_[i] << " "; + } + + sb << "Unexpected ack order: got " << ack << ", expected " << expectedAck; + } + UNIT_ASSERT_VALUES_EQUAL_C(ack, expectedAck, sb); expectedAck++; } } From 2746ce9895fa8b760ac025e5026d6c59fa08a379 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 13:40:22 +0000 Subject: [PATCH 24/93] Update import generation: 35 --- .github/import_generation.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/import_generation.txt b/.github/import_generation.txt index 8f92bfdd497..7facc89938b 100644 --- a/.github/import_generation.txt +++ b/.github/import_generation.txt @@ -1 +1 @@ -35 +36 From eaeaa93fb570c146340e5f99e248b806dfc1e75d Mon Sep 17 00:00:00 2001 From: Bulat Gayazov Date: Fri, 13 Feb 2026 16:05:33 +0000 Subject: [PATCH 25/93] Added future/subscription lib --- .github/scripts/copy_sources.sh | 1 + library/cpp/threading/future/CMakeLists.txt | 2 + .../future/subscription/CMakeLists.txt | 31 ++ .../threading/future/subscription/README.md | 104 +++++ .../future/subscription/subscription-inl.h | 118 +++++ .../future/subscription/subscription.cpp | 65 +++ .../future/subscription/subscription.h | 186 ++++++++ .../future/subscription/subscription_ut.cpp | 432 ++++++++++++++++++ .../threading/future/subscription/ut/ya.make | 11 + .../cpp/threading/future/subscription/wait.h | 119 +++++ .../future/subscription/wait_all.cpp | 1 + .../threading/future/subscription/wait_all.h | 26 ++ .../future/subscription/wait_all_inl.h | 80 ++++ .../subscription/wait_all_or_exception.cpp | 1 + .../subscription/wait_all_or_exception.h | 28 ++ .../subscription/wait_all_or_exception_inl.h | 79 ++++ .../subscription/wait_all_or_exception_ut.cpp | 167 +++++++ .../future/subscription/wait_all_ut.cpp | 161 +++++++ .../future/subscription/wait_any.cpp | 1 + .../threading/future/subscription/wait_any.h | 26 ++ .../future/subscription/wait_any_inl.h | 64 +++ .../future/subscription/wait_any_ut.cpp | 166 +++++++ .../future/subscription/wait_ut_common.cpp | 26 ++ .../future/subscription/wait_ut_common.h | 56 +++ .../cpp/threading/future/subscription/ya.make | 18 + .../yc_private/accessservice/sensitive.proto | 29 ++ .../impl/internal/common/client_pid.cpp | 2 + src/client/topic/impl/CMakeLists.txt | 1 + 28 files changed, 2001 insertions(+) create mode 100644 library/cpp/threading/future/subscription/CMakeLists.txt create mode 100644 library/cpp/threading/future/subscription/README.md create mode 100644 library/cpp/threading/future/subscription/subscription-inl.h create mode 100644 library/cpp/threading/future/subscription/subscription.cpp create mode 100644 library/cpp/threading/future/subscription/subscription.h create mode 100644 library/cpp/threading/future/subscription/subscription_ut.cpp create mode 100644 library/cpp/threading/future/subscription/ut/ya.make create mode 100644 library/cpp/threading/future/subscription/wait.h create mode 100644 library/cpp/threading/future/subscription/wait_all.cpp create mode 100644 library/cpp/threading/future/subscription/wait_all.h create mode 100644 library/cpp/threading/future/subscription/wait_all_inl.h create mode 100644 library/cpp/threading/future/subscription/wait_all_or_exception.cpp create mode 100644 library/cpp/threading/future/subscription/wait_all_or_exception.h create mode 100644 library/cpp/threading/future/subscription/wait_all_or_exception_inl.h create mode 100644 library/cpp/threading/future/subscription/wait_all_or_exception_ut.cpp create mode 100644 library/cpp/threading/future/subscription/wait_all_ut.cpp create mode 100644 library/cpp/threading/future/subscription/wait_any.cpp create mode 100644 library/cpp/threading/future/subscription/wait_any.h create mode 100644 library/cpp/threading/future/subscription/wait_any_inl.h create mode 100644 library/cpp/threading/future/subscription/wait_any_ut.cpp create mode 100644 library/cpp/threading/future/subscription/wait_ut_common.cpp create mode 100644 library/cpp/threading/future/subscription/wait_ut_common.h create mode 100644 library/cpp/threading/future/subscription/ya.make create mode 100644 src/api/client/yc_private/accessservice/sensitive.proto diff --git a/.github/scripts/copy_sources.sh b/.github/scripts/copy_sources.sh index b5076f0d34d..f3778bf21d8 100755 --- a/.github/scripts/copy_sources.sh +++ b/.github/scripts/copy_sources.sh @@ -25,6 +25,7 @@ rm -r $tmp_dir/tests/unit/client/draft mkdir -p $tmp_dir/src/api/client/yc_private mkdir -p $tmp_dir/src/api/client/yc_public +cp -r $1/ydb/public/api/client/yc_private/accessservice/sensitive.proto $tmp_dir/src/api/client/yc_private/accessservice/sensitive.proto cp -r $1/ydb/public/api/client/yc_private/iam $tmp_dir/src/api/client/yc_private cp -r $1/ydb/public/api/client/yc_private/operation $tmp_dir/src/api/client/yc_private cp -r $1/ydb/public/api/client/yc_public/common $tmp_dir/src/api/client/yc_public diff --git a/library/cpp/threading/future/CMakeLists.txt b/library/cpp/threading/future/CMakeLists.txt index 3ae66ecf4dc..c34885dea26 100644 --- a/library/cpp/threading/future/CMakeLists.txt +++ b/library/cpp/threading/future/CMakeLists.txt @@ -48,3 +48,5 @@ if (YDB_SDK_TESTS) unit ) endif() + +add_subdirectory(subscription) diff --git a/library/cpp/threading/future/subscription/CMakeLists.txt b/library/cpp/threading/future/subscription/CMakeLists.txt new file mode 100644 index 00000000000..972ecd40b81 --- /dev/null +++ b/library/cpp/threading/future/subscription/CMakeLists.txt @@ -0,0 +1,31 @@ +_ydb_sdk_add_library(threading-future-subscription) + +target_link_libraries(threading-future-subscription PUBLIC + yutil + threading-future +) + +target_sources(threading-future-subscription + PRIVATE + subscription.cpp + wait_all.cpp + wait_all_or_exception.cpp + wait_any.cpp +) + +_ydb_sdk_install_targets(TARGETS threading-future-subscription) + +if (YDB_SDK_TESTS) + add_ydb_test(NAME future-subscription-ut + SOURCES + subscription_ut.cpp + wait_all_ut.cpp + wait_all_or_exception_ut.cpp + wait_any_ut.cpp + wait_ut_common.cpp + LINK_LIBRARIES + threading-future-subscription + LABELS + unit + ) +endif() diff --git a/library/cpp/threading/future/subscription/README.md b/library/cpp/threading/future/subscription/README.md new file mode 100644 index 00000000000..6f547926854 --- /dev/null +++ b/library/cpp/threading/future/subscription/README.md @@ -0,0 +1,104 @@ +Subscriptions manager and wait primitives library +================================================= + +Wait primitives +--------------- + +All wait primitives are futures those being signaled when some or all of theirs dependencies are signaled. +Wait privimitives could be constructed either from an initializer_list or from a standard container of futures. + +1. WaitAll is signaled when all its dependencies are signaled: + + ```C++ + #include + + auto w = NWait::WaitAll({ future1, future2, ..., futureN }); + ... + w.Wait(); // wait for all futures + ``` + +2. WaitAny is signaled when any of its dependencies is signaled: + + ```C++ + #include + + auto w = NWait::WaitAny(TVector>{ future1, future2, ..., futureN }); + ... + w.Wait(); // wait for any future + ``` + +3. WaitAllOrException is signaled when all its dependencies are signaled with values or any dependency is signaled with an exception: + + ```C++ + #include + + auto w = NWait::WaitAllOrException(TVector>{ future1, future2, ..., futureN }); + ... + w.Wait(); // wait for all values or for an exception + ``` + +Subscriptions manager +--------------------- + +The subscription manager can manage multiple links beetween futures and callbacks. Multiple managed subscriptions to a single future shares just a single underlying subscription to the future. That allows dynamic creation and deletion of subscriptions and efficient implementation of different wait primitives. +The subscription manager could be used in the following way: + +1. Subscribe to a single future: + + ```C++ + #include + + TFuture LongOperation(); + + ... + auto future = LongRunnigOperation(); + auto m = MakeSubsriptionManager(); + auto id = m->Subscribe(future, [](TFuture const& f) { + try { + auto value = f.GetValue(); + ... + } catch (...) { + ... // handle exception + } + }); + if (id.has_value()) { + ... // Callback will run asynchronously + } else { + ... // Future has been signaled already. The callback has been invoked synchronously + } + ``` + + Note that a callback could be invoked synchronously during a Subscribe call. In this case the returned optional will have no value. + +2. Unsubscribe from a single future: + + ```C++ + // id holds the subscription id from a previous Subscribe call + m->Unsubscribe(id.value()); + ``` + + There is no need to call Unsubscribe if the callback has been called. In this case Unsubscribe will do nothing. And it is safe to call Unsubscribe with the same id multiple times. + +3. Subscribe a single callback to multiple futures: + + ```C++ + auto ids = m->Subscribe({ future1, future2, ..., futureN }, [](auto&& f) { ... }); + ... + ``` + + Futures could be passed to Subscribe method either via an initializer_list or via a standard container like vector or list. Subscribe method accept an optional boolean parameter revertOnSignaled. If the parameter is false (default) then all subscriptions will be performed regardless of the futures states and the returned vector will have a subscription id for each future (even if callback has been executed synchronously for some futures). Otherwise the method will stop on the first signaled future (the callback will be synchronously called for it), no subscriptions will be created and an empty vector will be returned. + +4. Unsubscribe multiple subscriptions: + + ```C++ + // ids is the vector or subscription ids + m->Unsubscribe(ids); + ``` + + The vector of IDs could be a result of a previous Subscribe call or an arbitrary set of IDs of previously created subscriptions. + +5. If you do not want to instantiate a new instance of the subscription manager it is possible to use the default instance: + + ```C++ + auto m = TSubscriptionManager::Default(); + ``` diff --git a/library/cpp/threading/future/subscription/subscription-inl.h b/library/cpp/threading/future/subscription/subscription-inl.h new file mode 100644 index 00000000000..a9d3b3114c7 --- /dev/null +++ b/library/cpp/threading/future/subscription/subscription-inl.h @@ -0,0 +1,118 @@ +#pragma once + +#if !defined(INCLUDE_LIBRARY_THREADING_FUTURE_SUBSCRIPTION_INL_H) +#error "you should never include subscription-inl.h directly" +#endif + +namespace NThreading { + +namespace NPrivate { + +template +TFutureStateId CheckedStateId(TFuture const& future) { + auto const id = future.StateId(); + if (id.Defined()) { + return *id; + } + ythrow TFutureException() << "Future state should be initialized"; +} + +} + +template +inline TSubscriptionManager::TSubscription::TSubscription(TFuture future, F&& callback, TCallbackExecutor&& executor) + : Callback( + [future = std::move(future), callback = std::forward(callback), executor = std::forward(executor)]() mutable { + executor(std::as_const(future), callback); + }) +{ +} + +template +inline std::optional TSubscriptionManager::Subscribe(TFuture const& future, F&& callback, TCallbackExecutor&& executor) { + auto stateId = NPrivate::CheckedStateId(future); + with_lock(Lock) { + auto const status = TrySubscribe(future, std::forward(callback), stateId, std::forward(executor)); + switch (status) { + case ECallbackStatus::Subscribed: + return TSubscriptionId(stateId, Revision); + case ECallbackStatus::ExecutedSynchronously: + return {}; + default: + Y_ABORT("Unexpected callback status"); + } + } +} + +template +inline TVector TSubscriptionManager::Subscribe(TFutures const& futures, F&& callback, bool revertOnSignaled + , TCallbackExecutor&& executor) +{ + return SubscribeImpl(futures, std::forward(callback), revertOnSignaled, std::forward(executor)); +} + +template +inline TVector TSubscriptionManager::Subscribe(std::initializer_list const> futures, F&& callback + , bool revertOnSignaled, TCallbackExecutor&& executor) +{ + return SubscribeImpl(futures, std::forward(callback), revertOnSignaled, std::forward(executor)); +} + +template +inline TSubscriptionManager::ECallbackStatus TSubscriptionManager::TrySubscribe(TFuture const& future, F&& callback, TFutureStateId stateId + , TCallbackExecutor&& executor) +{ + TSubscription subscription(future, std::forward(callback), std::forward(executor)); + auto const it = Subscriptions.find(stateId); + auto const revision = ++Revision; + if (it == std::end(Subscriptions)) { + auto const success = Subscriptions.emplace(stateId, THashMap{ { revision, std::move(subscription) } }).second; + Y_ABORT_UNLESS(success); + auto self = TSubscriptionManagerPtr(this); + future.Subscribe([self, stateId](TFuture const&) { self->OnCallback(stateId); }); + if (Subscriptions.find(stateId) == std::end(Subscriptions)) { + return ECallbackStatus::ExecutedSynchronously; + } + } else { + Y_ABORT_UNLESS(it->second.emplace(revision, std::move(subscription)).second); + } + return ECallbackStatus::Subscribed; +} + +template +inline TVector TSubscriptionManager::SubscribeImpl(TFutures const& futures, F const& callback, bool revertOnSignaled + , TCallbackExecutor const& executor) +{ + TVector results; + results.reserve(std::size(futures)); + // resolve all state ids to minimize processing under the lock + for (auto const& f : futures) { + results.push_back(TSubscriptionId(NPrivate::CheckedStateId(f), 0)); + } + with_lock(Lock) { + size_t i = 0; + for (auto const& f : futures) { + auto& r = results[i]; + auto const status = TrySubscribe(f, callback, r.StateId(), executor); + switch (status) { + case ECallbackStatus::Subscribed: + break; + case ECallbackStatus::ExecutedSynchronously: + if (revertOnSignaled) { + // revert + results.crop(i); + UnsubscribeImpl(results); + return {}; + } + break; + default: + Y_ABORT("Unexpected callback status"); + } + r.SetSubId(Revision); + ++i; + } + } + return results; +} + +} diff --git a/library/cpp/threading/future/subscription/subscription.cpp b/library/cpp/threading/future/subscription/subscription.cpp new file mode 100644 index 00000000000..e3cb3052c8d --- /dev/null +++ b/library/cpp/threading/future/subscription/subscription.cpp @@ -0,0 +1,65 @@ +#include "subscription.h" + +namespace NThreading { + +bool operator==(TSubscriptionId const& l, TSubscriptionId const& r) noexcept { + return l.StateId() == r.StateId() && l.SubId() == r.SubId(); +} + +bool operator!=(TSubscriptionId const& l, TSubscriptionId const& r) noexcept { + return !(l == r); +} + +void TSubscriptionManager::TSubscription::operator()() { + Callback(); +} + +TSubscriptionManagerPtr TSubscriptionManager::NewInstance() { + return new TSubscriptionManager(); +} + +TSubscriptionManagerPtr TSubscriptionManager::Default() { + static auto instance = NewInstance(); + return instance; +} + +void TSubscriptionManager::Unsubscribe(TSubscriptionId id) { + with_lock(Lock) { + UnsubscribeImpl(id); + } +} + +void TSubscriptionManager::Unsubscribe(TVector const& ids) { + with_lock(Lock) { + UnsubscribeImpl(ids); + } +} + +void TSubscriptionManager::OnCallback(TFutureStateId stateId) noexcept { + THashMap subscriptions; + with_lock(Lock) { + auto const it = Subscriptions.find(stateId); + Y_ABORT_UNLESS(it != Subscriptions.end(), "The callback has been triggered more than once"); + subscriptions.swap(it->second); + Subscriptions.erase(it); + } + for (auto& [_, subscription] : subscriptions) { + subscription(); + } +} + +void TSubscriptionManager::UnsubscribeImpl(TSubscriptionId id) { + auto const it = Subscriptions.find(id.StateId()); + if (it == std::end(Subscriptions)) { + return; + } + it->second.erase(id.SubId()); +} + +void TSubscriptionManager::UnsubscribeImpl(TVector const& ids) { + for (auto const& id : ids) { + UnsubscribeImpl(id); + } +} + +} diff --git a/library/cpp/threading/future/subscription/subscription.h b/library/cpp/threading/future/subscription/subscription.h new file mode 100644 index 00000000000..afe5eda7111 --- /dev/null +++ b/library/cpp/threading/future/subscription/subscription.h @@ -0,0 +1,186 @@ +#pragma once + +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace NThreading { + +namespace NPrivate { + +struct TNoexceptExecutor { + template + void operator()(TFuture const& future, F&& callee) const noexcept { + return callee(future); + } +}; + +} + +class TSubscriptionManager; + +using TSubscriptionManagerPtr = TIntrusivePtr; + +//! A subscription id +class TSubscriptionId { +private: + TFutureStateId StateId_; + ui64 SubId_; // Secondary id to make the whole subscription id unique + + friend class TSubscriptionManager; + +public: + TFutureStateId StateId() const noexcept { + return StateId_; + } + + ui64 SubId() const noexcept { + return SubId_; + } + +private: + TSubscriptionId(TFutureStateId stateId, ui64 subId) + : StateId_(stateId) + , SubId_(subId) + { + } + + void SetSubId(ui64 subId) noexcept { + SubId_ = subId; + } +}; + +bool operator==(TSubscriptionId const& l, TSubscriptionId const& r) noexcept; +bool operator!=(TSubscriptionId const& l, TSubscriptionId const& r) noexcept; + +//! The subscription manager manages subscriptions to futures +/** It provides an ability to create (and drop) multiple subscriptions to any future + with just a single underlying subscription per future. + + When a future is signaled all its subscriptions are removed. + So, there no need to call Unsubscribe for subscriptions to already signaled futures. + + Warning!!! For correct operation this class imposes the following requirement to futures/promises: + Any used future must be signaled (value or exception set) before the future state destruction. + Otherwise subscriptions and futures may happen. + Current future design does not provide the required guarantee. But that should be fixed soon. +**/ +class TSubscriptionManager final : public TAtomicRefCount { +private: + //! A single subscription + class TSubscription { + private: + std::function Callback; + + public: + template + TSubscription(TFuture future, F&& callback, TCallbackExecutor&& executor); + + void operator()(); + }; + + struct TFutureStateIdHash { + size_t operator()(TFutureStateId const id) const noexcept { + auto const value = id.Value(); + return ::hash()(value); + } + }; + +private: + THashMap, TFutureStateIdHash> Subscriptions; + ui64 Revision = 0; + TMutex Lock; + +public: + //! Creates a new subscription manager instance + static TSubscriptionManagerPtr NewInstance(); + + //! The default subscription manager instance + static TSubscriptionManagerPtr Default(); + + //! Attempts to subscribe the callback to the future + /** Subscription should succeed if the future is not signaled yet. + Otherwise the callback will be called synchronously and nullopt will be returned + + @param future - The future to subscribe to + @param callback - The callback to attach + @return The subscription id on success, nullopt if the future has been signaled already + **/ + template + std::optional Subscribe(TFuture const& future, F&& callback + , TCallbackExecutor&& executor = NPrivate::TNoexceptExecutor()); + + //! Drops the subscription with the given id + /** @param id - The subscription id + **/ + void Unsubscribe(TSubscriptionId id); + + //! Attempts to subscribe the callback to the set of futures + /** @param futures - The futures to subscribe to + @param callback - The callback to attach + @param revertOnSignaled - Shows whether to stop and revert the subscription process if one of the futures is in signaled state + @return The vector of subscription ids if no revert happened or an empty vector otherwise + A subscription id will be valid even if a corresponding future has been signaled + **/ + template + TVector Subscribe(TFutures const& futures, F&& callback, bool revertOnSignaled = false + , TCallbackExecutor&& executor = NPrivate::TNoexceptExecutor()); + + //! Attempts to subscribe the callback to the set of futures + /** @param futures - The futures to subscribe to + @param callback - The callback to attach + @param revertOnSignaled - Shows whether to stop and revert the subscription process if one of the futures is in signaled state + @return The vector of subscription ids if no revert happened or an empty vector otherwise + A subscription id will be valid even if a corresponding future has been signaled + **/ + template + TVector Subscribe(std::initializer_list const> futures, F&& callback, bool revertOnSignaled = false + , TCallbackExecutor&& executor = NPrivate::TNoexceptExecutor()); + + //! Drops the subscriptions with the given ids + /** @param ids - The subscription ids + **/ + void Unsubscribe(TVector const& ids); + +private: + enum class ECallbackStatus { + Subscribed, //! A subscription has been created. The callback will be called asynchronously. + ExecutedSynchronously //! A callback has been called synchronously. No subscription has been created + }; + +private: + //! .ctor + TSubscriptionManager() = default; + //! Processes a callback from a future + void OnCallback(TFutureStateId stateId) noexcept; + //! Attempts to create a subscription + /** This method should be called under the lock + **/ + template + ECallbackStatus TrySubscribe(TFuture const& future, F&& callback, TFutureStateId stateId, TCallbackExecutor&& executor); + //! Batch subscribe implementation + template + TVector SubscribeImpl(TFutures const& futures, F const& callback, bool revertOnSignaled + , TCallbackExecutor const& executor); + //! Unsubscribe implementation + /** This method should be called under the lock + **/ + void UnsubscribeImpl(TSubscriptionId id); + //! Batch unsubscribe implementation + /** This method should be called under the lock + **/ + void UnsubscribeImpl(TVector const& ids); +}; + +} + +#define INCLUDE_LIBRARY_THREADING_FUTURE_SUBSCRIPTION_INL_H +#include "subscription-inl.h" +#undef INCLUDE_LIBRARY_THREADING_FUTURE_SUBSCRIPTION_INL_H diff --git a/library/cpp/threading/future/subscription/subscription_ut.cpp b/library/cpp/threading/future/subscription/subscription_ut.cpp new file mode 100644 index 00000000000..d018ea15cc2 --- /dev/null +++ b/library/cpp/threading/future/subscription/subscription_ut.cpp @@ -0,0 +1,432 @@ +#include "subscription.h" + +#include + +using namespace NThreading; + +Y_UNIT_TEST_SUITE(TSubscriptionManagerTest) { + + Y_UNIT_TEST(TestSubscribeUnsignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount = 0; + auto id = m->Subscribe(p.GetFuture(), [&callCount](auto&&) { ++callCount; } ); + UNIT_ASSERT(id.has_value()); + UNIT_ASSERT_EQUAL(callCount, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount, 1); + } + + Y_UNIT_TEST(TestSubscribeSignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto f = MakeFuture(); + + size_t callCount = 0; + auto id = m->Subscribe(f, [&callCount](auto&&) { ++callCount; } ); + UNIT_ASSERT(!id.has_value()); + UNIT_ASSERT_EQUAL(callCount, 1); + } + + Y_UNIT_TEST(TestSubscribeUnsignaledAndSignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount1, 1); + + size_t callCount2 = 0; + auto id2 = m->Subscribe(p.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + UNIT_ASSERT(!id2.has_value()); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount1, 1); + } + + Y_UNIT_TEST(TestSubscribeUnsubscribeUnsignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount = 0; + auto id = m->Subscribe(p.GetFuture(), [&callCount](auto&&) { ++callCount; } ); + UNIT_ASSERT(id.has_value()); + UNIT_ASSERT_EQUAL(callCount, 0); + + m->Unsubscribe(id.value()); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount, 0); + } + + Y_UNIT_TEST(TestSubscribeUnsignaledUnsubscribeSignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount = 0; + auto id = m->Subscribe(p.GetFuture(), [&callCount](auto&&) { ++callCount; } ); + UNIT_ASSERT(id.has_value()); + UNIT_ASSERT_EQUAL(callCount, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount, 1); + + m->Unsubscribe(id.value()); + UNIT_ASSERT_EQUAL(callCount, 1); + } + + Y_UNIT_TEST(TestUnsubscribeTwice) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount = 0; + auto id = m->Subscribe(p.GetFuture(), [&callCount](auto&&) { ++callCount; } ); + UNIT_ASSERT(id.has_value()); + UNIT_ASSERT_EQUAL(callCount, 0); + + m->Unsubscribe(id.value()); + UNIT_ASSERT_EQUAL(callCount, 0); + m->Unsubscribe(id.value()); + UNIT_ASSERT_EQUAL(callCount, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount, 0); + } + + Y_UNIT_TEST(TestSubscribeOneUnsignaledManyTimes) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(p.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(p.GetFuture(), [&callCount3](auto&&) { ++callCount3; } ); + + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT(id2.has_value()); + UNIT_ASSERT(id3.has_value()); + UNIT_ASSERT_UNEQUAL(id1.value(), id2.value()); + UNIT_ASSERT_UNEQUAL(id2.value(), id3.value()); + UNIT_ASSERT_UNEQUAL(id3.value(), id1.value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + } + + Y_UNIT_TEST(TestSubscribeOneSignaledManyTimes) { + auto m = TSubscriptionManager::NewInstance(); + auto f = MakeFuture(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(f, [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(f, [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(f, [&callCount3](auto&&) { ++callCount3; } ); + + UNIT_ASSERT(!id1.has_value()); + UNIT_ASSERT(!id2.has_value()); + UNIT_ASSERT(!id3.has_value()); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + } + + Y_UNIT_TEST(TestSubscribeUnsubscribeOneUnsignaledManyTimes) { + auto m = TSubscriptionManager::NewInstance(); + auto p = NewPromise(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(p.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(p.GetFuture(), [&callCount3](auto&&) { ++callCount3; } ); + size_t callCount4 = 0; + auto id4 = m->Subscribe(p.GetFuture(), [&callCount4](auto&&) { ++callCount4; } ); + + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT(id2.has_value()); + UNIT_ASSERT(id3.has_value()); + UNIT_ASSERT(id4.has_value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 0); + + m->Unsubscribe(id3.value()); + m->Unsubscribe(id1.value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 0); + + p.SetValue(); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 1); + } + + Y_UNIT_TEST(TestSubscribeManyUnsignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p1.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(p2.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(p1.GetFuture(), [&callCount3](auto&&) { ++callCount3; } ); + + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT(id2.has_value()); + UNIT_ASSERT(id3.has_value()); + UNIT_ASSERT_UNEQUAL(id1.value(), id2.value()); + UNIT_ASSERT_UNEQUAL(id2.value(), id3.value()); + UNIT_ASSERT_UNEQUAL(id3.value(), id1.value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + + p1.SetValue(33); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 1); + + p2.SetValue(111); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + } + + Y_UNIT_TEST(TestSubscribeManySignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto f1 = MakeFuture(0); + auto f2 = MakeFuture(1); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(f1, [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(f2, [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(f2, [&callCount3](auto&&) { ++callCount3; } ); + + UNIT_ASSERT(!id1.has_value()); + UNIT_ASSERT(!id2.has_value()); + UNIT_ASSERT(!id3.has_value()); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + } + + Y_UNIT_TEST(TestSubscribeManyMixed) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p1.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(p2.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(f, [&callCount3](auto&&) { ++callCount3; } ); + + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT(id2.has_value()); + UNIT_ASSERT(!id3.has_value()); + UNIT_ASSERT_UNEQUAL(id1.value(), id2.value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 1); + + p1.SetValue(45); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 1); + + p2.SetValue(-7); + UNIT_ASSERT_EQUAL(callCount1, 1); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + } + + Y_UNIT_TEST(TestSubscribeUnsubscribeMany) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto p3 = NewPromise(); + + size_t callCount1 = 0; + auto id1 = m->Subscribe(p1.GetFuture(), [&callCount1](auto&&) { ++callCount1; } ); + size_t callCount2 = 0; + auto id2 = m->Subscribe(p2.GetFuture(), [&callCount2](auto&&) { ++callCount2; } ); + size_t callCount3 = 0; + auto id3 = m->Subscribe(p3.GetFuture(), [&callCount3](auto&&) { ++callCount3; } ); + size_t callCount4 = 0; + auto id4 = m->Subscribe(p2.GetFuture(), [&callCount4](auto&&) { ++callCount4; } ); + size_t callCount5 = 0; + auto id5 = m->Subscribe(p1.GetFuture(), [&callCount5](auto&&) { ++callCount5; } ); + + UNIT_ASSERT(id1.has_value()); + UNIT_ASSERT(id2.has_value()); + UNIT_ASSERT(id3.has_value()); + UNIT_ASSERT(id4.has_value()); + UNIT_ASSERT(id5.has_value()); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 0); + UNIT_ASSERT_EQUAL(callCount5, 0); + + m->Unsubscribe(id1.value()); + p1.SetValue(-1); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 0); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 0); + UNIT_ASSERT_EQUAL(callCount5, 1); + + m->Unsubscribe(id4.value()); + p2.SetValue(23); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 0); + UNIT_ASSERT_EQUAL(callCount4, 0); + UNIT_ASSERT_EQUAL(callCount5, 1); + + p3.SetValue(100500); + UNIT_ASSERT_EQUAL(callCount1, 0); + UNIT_ASSERT_EQUAL(callCount2, 1); + UNIT_ASSERT_EQUAL(callCount3, 1); + UNIT_ASSERT_EQUAL(callCount4, 0); + UNIT_ASSERT_EQUAL(callCount5, 1); + } + + Y_UNIT_TEST(TestBulkSubscribeManyUnsignaled) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + + size_t callCount = 0; + auto ids = m->Subscribe({ p1.GetFuture(), p2.GetFuture(), p1.GetFuture() }, [&callCount](auto&&) { ++callCount; }); + + UNIT_ASSERT_EQUAL(ids.size(), 3); + UNIT_ASSERT_UNEQUAL(ids[0], ids[1]); + UNIT_ASSERT_UNEQUAL(ids[1], ids[2]); + UNIT_ASSERT_UNEQUAL(ids[2], ids[0]); + UNIT_ASSERT_EQUAL(callCount, 0); + + p1.SetValue(33); + UNIT_ASSERT_EQUAL(callCount, 2); + + p2.SetValue(111); + UNIT_ASSERT_EQUAL(callCount, 3); + } + + Y_UNIT_TEST(TestBulkSubscribeManySignaledNoRevert) { + auto m = TSubscriptionManager::NewInstance(); + auto f1 = MakeFuture(0); + auto f2 = MakeFuture(1); + + size_t callCount = 0; + auto ids = m->Subscribe({ f1, f2, f1 }, [&callCount](auto&&) { ++callCount; }); + + UNIT_ASSERT_EQUAL(ids.size(), 3); + UNIT_ASSERT_UNEQUAL(ids[0], ids[1]); + UNIT_ASSERT_UNEQUAL(ids[1], ids[2]); + UNIT_ASSERT_UNEQUAL(ids[2], ids[0]); + UNIT_ASSERT_EQUAL(callCount, 3); + } + + Y_UNIT_TEST(TestBulkSubscribeManySignaledRevert) { + auto m = TSubscriptionManager::NewInstance(); + auto f1 = MakeFuture(0); + auto f2 = MakeFuture(1); + + size_t callCount = 0; + auto ids = m->Subscribe({ f1, f2, f1 }, [&callCount](auto&&) { ++callCount; }, true); + + UNIT_ASSERT(ids.empty()); + UNIT_ASSERT_EQUAL(callCount, 1); + } + + Y_UNIT_TEST(TestBulkSubscribeManyMixedNoRevert) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + + size_t callCount = 0; + auto ids = m->Subscribe({ p1.GetFuture(), p2.GetFuture(), f }, [&callCount](auto&&) { ++callCount; } ); + + UNIT_ASSERT_EQUAL(ids.size(), 3); + UNIT_ASSERT_UNEQUAL(ids[0], ids[1]); + UNIT_ASSERT_UNEQUAL(ids[1], ids[2]); + UNIT_ASSERT_UNEQUAL(ids[2], ids[0]); + UNIT_ASSERT_EQUAL(callCount, 1); + + p1.SetValue(45); + UNIT_ASSERT_EQUAL(callCount, 2); + + p2.SetValue(-7); + UNIT_ASSERT_EQUAL(callCount, 3); + } + + Y_UNIT_TEST(TestBulkSubscribeManyMixedRevert) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(); + + size_t callCount = 0; + auto ids = m->Subscribe({ p1.GetFuture(), f, p2.GetFuture() }, [&callCount](auto&&) { ++callCount; }, true); + + UNIT_ASSERT(ids.empty()); + UNIT_ASSERT_EQUAL(callCount, 1); + + p1.SetValue(); + p2.SetValue(); + UNIT_ASSERT_EQUAL(callCount, 1); + } + + Y_UNIT_TEST(TestBulkSubscribeUnsubscribeMany) { + auto m = TSubscriptionManager::NewInstance(); + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto p3 = NewPromise(); + + size_t callCount = 0; + auto ids = m->Subscribe( + TVector>{ p1.GetFuture(), p2.GetFuture(), p3.GetFuture(), p2.GetFuture(), p1.GetFuture() } + , [&callCount](auto&&) { ++callCount; } ); + + UNIT_ASSERT_EQUAL(ids.size(), 5); + UNIT_ASSERT_EQUAL(callCount, 0); + + m->Unsubscribe(TVector{ ids[0], ids[3] }); + UNIT_ASSERT_EQUAL(callCount, 0); + + p1.SetValue(-1); + UNIT_ASSERT_EQUAL(callCount, 1); + + p2.SetValue(23); + UNIT_ASSERT_EQUAL(callCount, 2); + + p3.SetValue(100500); + UNIT_ASSERT_EQUAL(callCount, 3); + } +} diff --git a/library/cpp/threading/future/subscription/ut/ya.make b/library/cpp/threading/future/subscription/ut/ya.make new file mode 100644 index 00000000000..9b7e371509b --- /dev/null +++ b/library/cpp/threading/future/subscription/ut/ya.make @@ -0,0 +1,11 @@ +UNITTEST_FOR(library/cpp/threading/future/subscription) + +SRCS( + subscription_ut.cpp + wait_all_ut.cpp + wait_all_or_exception_ut.cpp + wait_any_ut.cpp + wait_ut_common.cpp +) + +END() diff --git a/library/cpp/threading/future/subscription/wait.h b/library/cpp/threading/future/subscription/wait.h new file mode 100644 index 00000000000..533bab9d8d9 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait.h @@ -0,0 +1,119 @@ +#pragma once + +#include "subscription.h" + +#include +#include +#include + + +#include + +namespace NThreading::NPrivate { + +template +class TWait : public TThrRefBase { +private: + TSubscriptionManagerPtr Manager; + TVector Subscriptions; + bool Unsubscribed = false; + +protected: + TAdaptiveLock Lock; + TPromise Promise; + +public: + template + static TFuture Make(TFutures const& futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + TIntrusivePtr w(new TDerived(std::move(manager))); + w->Subscribe(futures, std::forward(executor)); + return w->Promise.GetFuture(); + } + +protected: + TWait(TSubscriptionManagerPtr manager) + : Manager(std::move(manager)) + , Subscriptions() + , Unsubscribed(false) + , Lock() + , Promise(NewPromise()) + { + Y_ENSURE(Manager != nullptr); + } + +protected: + //! Unsubscribes all existing subscriptions + /** Lock should be acquired! + **/ + void Unsubscribe() noexcept { + if (Unsubscribed) { + return; + } + Unsubscribe(Subscriptions); + Subscriptions.clear(); + } + +private: + //! Performs a subscription to the given futures + /** Lock should not be acquired! + @param future - The futures to subscribe to + @param callback - The callback to call for each future + **/ + template + void Subscribe(TFutures const& futures, TCallbackExecutor&& executor) { + auto self = TIntrusivePtr(static_cast(this)); + self->BeforeSubscribe(futures); + auto callback = [self = std::move(self)](const auto& future) mutable { + self->Set(future); + }; + auto subscriptions = Manager->Subscribe(futures, callback, TDerived::RevertOnSignaled, std::forward(executor)); + if (subscriptions.empty()) { + return; + } + with_lock (Lock) { + if (Unsubscribed) { + Unsubscribe(subscriptions); + } else { + Subscriptions = std::move(subscriptions); + } + } + } + + void Unsubscribe(TVector& subscriptions) noexcept { + Manager->Unsubscribe(subscriptions); + Unsubscribed = true; + } +}; + +template +TFuture Wait(TFutures const& futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + switch (std::size(futures)) { + case 0: + return MakeFuture(); + case 1: + return std::begin(futures)->IgnoreResult(); + default: + return TWaiter::Make(futures, std::move(manager), std::forward(executor)); + } +} + +template +TFuture Wait(std::initializer_list const> futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + switch (std::size(futures)) { + case 0: + return MakeFuture(); + case 1: + return std::begin(futures)->IgnoreResult(); + default: + return TWaiter::Make(futures, std::move(manager), std::forward(executor)); + } +} + + +template +TFuture Wait(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return TWaiter::Make(std::initializer_list const>({ future1, future2 }), std::move(manager) + , std::forward(executor)); +} + +} diff --git a/library/cpp/threading/future/subscription/wait_all.cpp b/library/cpp/threading/future/subscription/wait_all.cpp new file mode 100644 index 00000000000..10e7ee75984 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all.cpp @@ -0,0 +1 @@ +#include "wait_all.h" diff --git a/library/cpp/threading/future/subscription/wait_all.h b/library/cpp/threading/future/subscription/wait_all.h new file mode 100644 index 00000000000..8c1e6fea3a5 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all.h @@ -0,0 +1,26 @@ +#pragma once + +#include "wait.h" + +namespace NThreading::NWait { + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAll(TFutures const& futures, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAll(std::initializer_list const> futures, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAll(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +} + +#define INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_INL_H +#include "wait_all_inl.h" +#undef INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_INL_H diff --git a/library/cpp/threading/future/subscription/wait_all_inl.h b/library/cpp/threading/future/subscription/wait_all_inl.h new file mode 100644 index 00000000000..a3b665f6427 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_inl.h @@ -0,0 +1,80 @@ +#pragma once + +#if !defined(INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_INL_H) +#error "you should never include wait_all_inl.h directly" +#endif + +#include "subscription.h" + +#include + +namespace NThreading::NWait { + +namespace NPrivate { + +class TWaitAll final : public NThreading::NPrivate::TWait { +private: + size_t Count = 0; + std::exception_ptr Exception; + + static constexpr bool RevertOnSignaled = false; + + using TBase = NThreading::NPrivate::TWait; + friend TBase; + +private: + TWaitAll(TSubscriptionManagerPtr manager) + : TBase(std::move(manager)) + , Count(0) + , Exception() + { + } + + template + void BeforeSubscribe(TFutures const& futures) { + Count = std::size(futures); + Y_ENSURE(Count > 0, "It is meaningless to use this class with empty futures set"); + } + + template + void Set(TFuture const& future) { + with_lock (TBase::Lock) { + if (!Exception) { + try { + future.TryRethrow(); + } catch (...) { + Exception = std::current_exception(); + } + } + + if (--Count == 0) { + // there is no need to call Unsubscribe here since all futures are signaled + Y_ASSERT(!TBase::Promise.HasValue() && !TBase::Promise.HasException()); + if (Exception) { + TBase::Promise.SetException(std::move(Exception)); + } else { + TBase::Promise.SetValue(); + } + } + } + } +}; + +} + +template +TFuture WaitAll(TFutures const& futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} + +template +TFuture WaitAll(std::initializer_list const> futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} + +template +TFuture WaitAll(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(future1, future2, std::move(manager), std::forward(executor)); +} + +} diff --git a/library/cpp/threading/future/subscription/wait_all_or_exception.cpp b/library/cpp/threading/future/subscription/wait_all_or_exception.cpp new file mode 100644 index 00000000000..0c73ddeb84a --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_or_exception.cpp @@ -0,0 +1 @@ +#include "wait_all_or_exception.h" diff --git a/library/cpp/threading/future/subscription/wait_all_or_exception.h b/library/cpp/threading/future/subscription/wait_all_or_exception.h new file mode 100644 index 00000000000..10bba2bffa7 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_or_exception.h @@ -0,0 +1,28 @@ +#pragma once + +#include "wait.h" + +namespace NThreading::NWait { + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAllOrException(TFutures const& futures, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAllOrException(std::initializer_list const> futures + , TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAllOrException(TFuture const& future1, TFuture const& future2 + , TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +} + +#define INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_OR_EXCEPTION_INL_H +#include "wait_all_or_exception_inl.h" +#undef INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_OR_EXCEPTION_INL_H diff --git a/library/cpp/threading/future/subscription/wait_all_or_exception_inl.h b/library/cpp/threading/future/subscription/wait_all_or_exception_inl.h new file mode 100644 index 00000000000..fcd9782d543 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_or_exception_inl.h @@ -0,0 +1,79 @@ +#pragma once + +#if !defined(INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ALL_OR_EXCEPTION_INL_H) +#error "you should never include wait_all_or_exception_inl.h directly" +#endif + +#include "subscription.h" + +#include + +namespace NThreading::NWait { + +namespace NPrivate { + +class TWaitAllOrException final : public NThreading::NPrivate::TWait +{ +private: + size_t Count = 0; + + static constexpr bool RevertOnSignaled = false; + + using TBase = NThreading::NPrivate::TWait; + friend TBase; + +private: + TWaitAllOrException(TSubscriptionManagerPtr manager) + : TBase(std::move(manager)) + , Count(0) + { + } + + template + void BeforeSubscribe(TFutures const& futures) { + Count = std::size(futures); + Y_ENSURE(Count > 0, "It is meaningless to use this class with empty futures set"); + } + + template + void Set(TFuture const& future) { + with_lock (TBase::Lock) { + try { + future.TryRethrow(); + if (--Count == 0) { + // there is no need to call Unsubscribe here since all futures are signaled + TBase::Promise.SetValue(); + } + } catch (...) { + Y_ASSERT(!TBase::Promise.HasValue()); + TBase::Unsubscribe(); + if (!TBase::Promise.HasException()) { + TBase::Promise.SetException(std::current_exception()); + } + } + } + } +}; + +} + +template +TFuture WaitAllOrException(TFutures const& futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} + +template +TFuture WaitAllOrException(std::initializer_list const> futures, TSubscriptionManagerPtr manager + , TCallbackExecutor&& executor) +{ + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} +template +TFuture WaitAllOrException(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager + , TCallbackExecutor&& executor) +{ + return NThreading::NPrivate::Wait(future1, future2, std::move(manager) + , std::forward(executor)); +} + +} diff --git a/library/cpp/threading/future/subscription/wait_all_or_exception_ut.cpp b/library/cpp/threading/future/subscription/wait_all_or_exception_ut.cpp new file mode 100644 index 00000000000..34ae9edb4e6 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_or_exception_ut.cpp @@ -0,0 +1,167 @@ +#include "wait_all_or_exception.h" +#include "wait_ut_common.h" + +#include +#include + +#include +#include + +using namespace NThreading; + +Y_UNIT_TEST_SUITE(TWaitAllOrExceptionTest) { + + Y_UNIT_TEST(TestTwoUnsignaled) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAllOrException(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(10); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + p2.SetValue(1); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestTwoUnsignaledWithException) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAllOrException(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception"; + p1.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p2.SetValue(-11); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaled) { + auto p = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAllOrException(p.GetFuture(), f); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p.SetValue(); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaledWithException) { + auto p = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAllOrException(f, p.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 2"; + p.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestEmptyInitializer) { + auto w = NWait::WaitAllOrException(std::initializer_list const>({})); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestEmptyVector) { + auto w = NWait::WaitAllOrException(TVector>()); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithInitializer) { + auto p = NewPromise(); + auto w = NWait::WaitAllOrException({ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p.SetValue(1); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithVector) { + auto p = NewPromise(); + auto w = NWait::WaitAllOrException(TVector>{ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 3"; + p.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestManyWithInitializer) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + auto w = NWait::WaitAllOrException({ p1.GetFuture(), f, p2.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(10); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + p2.SetValue(-3); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestManyWithVector) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + auto w = NWait::WaitAllOrException(TVector>{ p1.GetFuture(), f, p2.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 4"; + p1.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p2.SetValue(34); + } + + Y_UNIT_TEST(TestManyWithVectorAndIntialError) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + constexpr TStringBuf message = "Test exception 5"; + auto f = MakeErrorFuture(std::make_exception_ptr(yexception() << message)); + auto w = NWait::WaitAllOrException(TVector>{ p1.GetFuture(), p2.GetFuture(), f }); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p1.SetValue(); + p2.SetValue(); + } + + Y_UNIT_TEST(TestManyStress) { + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAllOrException(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAllOrException(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(22); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + auto e = std::make_exception_ptr(yexception() << "Test exception 6"); + std::atomic index = 0; + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAllOrException(futures); } + , [e, &index](size_t size) { + auto exceptionIndex = size / 2; + index = 0; + return [e, exceptionIndex, &index](auto&& p) { + if (index++ == exceptionIndex) { + p.SetException(e); + } else { + p.SetValue(); + } + }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasException()); }); + } + +} diff --git a/library/cpp/threading/future/subscription/wait_all_ut.cpp b/library/cpp/threading/future/subscription/wait_all_ut.cpp new file mode 100644 index 00000000000..3bc9762671c --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_all_ut.cpp @@ -0,0 +1,161 @@ +#include "wait_all.h" +#include "wait_ut_common.h" + +#include +#include + +#include +#include + +using namespace NThreading; + +Y_UNIT_TEST_SUITE(TWaitAllTest) { + + Y_UNIT_TEST(TestTwoUnsignaled) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAll(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(10); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + p2.SetValue(1); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestTwoUnsignaledWithException) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAll(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception"; + p1.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p2.SetValue(-11); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaled) { + auto p = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAll(p.GetFuture(), f); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p.SetValue(); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaledWithException) { + auto p = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAll(f, p.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 2"; + p.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestEmptyInitializer) { + auto w = NWait::WaitAll(std::initializer_list const>({})); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestEmptyVector) { + auto w = NWait::WaitAll(TVector>()); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithInitializer) { + auto p = NewPromise(); + auto w = NWait::WaitAll({ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p.SetValue(1); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithVector) { + auto p = NewPromise(); + auto w = NWait::WaitAll(TVector>{ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 3"; + p.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestManyWithInitializer) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + auto w = NWait::WaitAll({ p1.GetFuture(), f, p2.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(10); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + p2.SetValue(-3); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestManyWithVector) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + auto w = NWait::WaitAll(TVector>{ p1.GetFuture(), f, p2.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 4"; + p1.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p2.SetValue(34); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestManyStress) { + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAll(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(42); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAll(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + auto e = std::make_exception_ptr(yexception() << "Test exception 5"); + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAll(futures); } + , [e](size_t) { + return [e](auto&& p) { p.SetException(e); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasException()); }); + e = std::make_exception_ptr(yexception() << "Test exception 6"); + std::atomic index = 0; + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAll(futures); } + , [e, &index](size_t size) { + auto exceptionIndex = size / 2; + index = 0; + return [e, exceptionIndex, &index](auto&& p) { + if (index++ == exceptionIndex) { + p.SetException(e); + } else { + p.SetValue(index); + } + }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasException()); }); + } + +} diff --git a/library/cpp/threading/future/subscription/wait_any.cpp b/library/cpp/threading/future/subscription/wait_any.cpp new file mode 100644 index 00000000000..57cc1b2c253 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_any.cpp @@ -0,0 +1 @@ +#include "wait_any.h" diff --git a/library/cpp/threading/future/subscription/wait_any.h b/library/cpp/threading/future/subscription/wait_any.h new file mode 100644 index 00000000000..969e307a897 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_any.h @@ -0,0 +1,26 @@ +#pragma once + +#include "wait.h" + +namespace NThreading::NWait { + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAny(TFutures const& futures, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAny(std::initializer_list const> futures, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +template +[[nodiscard("This method creates TFuture, wait for it")]] +TFuture WaitAny(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager = TSubscriptionManager::Default() + , TCallbackExecutor&& executor = TCallbackExecutor()); + +} + +#define INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ANY_INL_H +#include "wait_any_inl.h" +#undef INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ANY_INL_H diff --git a/library/cpp/threading/future/subscription/wait_any_inl.h b/library/cpp/threading/future/subscription/wait_any_inl.h new file mode 100644 index 00000000000..e80822bfc9c --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_any_inl.h @@ -0,0 +1,64 @@ +#pragma once + +#if !defined(INCLUDE_LIBRARY_THREADING_FUTURE_WAIT_ANY_INL_H) +#error "you should never include wait_any_inl.h directly" +#endif + +#include "subscription.h" + +#include + +namespace NThreading::NWait { + +namespace NPrivate { + +class TWaitAny final : public NThreading::NPrivate::TWait { +private: + static constexpr bool RevertOnSignaled = true; + + using TBase = NThreading::NPrivate::TWait; + friend TBase; + +private: + TWaitAny(TSubscriptionManagerPtr manager) + : TBase(std::move(manager)) + { + } + + template + void BeforeSubscribe(TFutures const& futures) { + Y_ENSURE(std::size(futures) > 0, "Futures set cannot be empty"); + } + + template + void Set(TFuture const& future) { + with_lock (TBase::Lock) { + TBase::Unsubscribe(); + try { + future.TryRethrow(); + TBase::Promise.TrySetValue(); + } catch (...) { + TBase::Promise.TrySetException(std::current_exception()); + } + } + } +}; + +} + +template +TFuture WaitAny(TFutures const& futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} + +template +TFuture WaitAny(std::initializer_list const> futures, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(futures, std::move(manager), std::forward(executor)); +} + +template +TFuture WaitAny(TFuture const& future1, TFuture const& future2, TSubscriptionManagerPtr manager, TCallbackExecutor&& executor) { + return NThreading::NPrivate::Wait(future1, future2, std::move(manager), std::forward(executor)); +} + +} diff --git a/library/cpp/threading/future/subscription/wait_any_ut.cpp b/library/cpp/threading/future/subscription/wait_any_ut.cpp new file mode 100644 index 00000000000..262080e8d12 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_any_ut.cpp @@ -0,0 +1,166 @@ +#include "wait_any.h" +#include "wait_ut_common.h" + +#include +#include + +#include + +using namespace NThreading; + +Y_UNIT_TEST_SUITE(TWaitAnyTest) { + + Y_UNIT_TEST(TestTwoUnsignaled) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAny(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(10); + UNIT_ASSERT(w.HasValue()); + p2.SetValue(1); + } + + Y_UNIT_TEST(TestTwoUnsignaledWithException) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto w = NWait::WaitAny(p1.GetFuture(), p2.GetFuture()); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception"; + p2.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p1.SetValue(-11); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaled) { + auto p = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAny(p.GetFuture(), f); + UNIT_ASSERT(w.HasValue()); + + p.SetValue(); + } + + Y_UNIT_TEST(TestOneUnsignaledOneSignaledWithException) { + auto p = NewPromise(); + constexpr TStringBuf message = "Test exception 2"; + auto f = MakeErrorFuture(std::make_exception_ptr(yexception() << message)); + auto w = NWait::WaitAny(f, p.GetFuture()); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p.SetValue(); + } + + Y_UNIT_TEST(TestEmptyInitializer) { + auto w = NWait::WaitAny(std::initializer_list const>({})); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestEmptyVector) { + auto w = NWait::WaitAny(TVector>()); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithInitializer) { + auto p = NewPromise(); + auto w = NWait::WaitAny({ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p.SetValue(1); + UNIT_ASSERT(w.HasValue()); + } + + Y_UNIT_TEST(TestOneUnsignaledWithVector) { + auto p = NewPromise(); + auto w = NWait::WaitAny(TVector>{ p.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 3"; + p.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + } + + Y_UNIT_TEST(TestManyUnsignaledWithInitializer) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto p3 = NewPromise(); + auto w = NWait::WaitAny({ p1.GetFuture(), p2.GetFuture(), p3.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + p1.SetValue(42); + UNIT_ASSERT(w.HasValue()); + + p2.SetValue(-3); + p3.SetValue(12); + } + + Y_UNIT_TEST(TestManyMixedWithInitializer) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(42); + auto w = NWait::WaitAny({ p1.GetFuture(), f, p2.GetFuture() }); + UNIT_ASSERT(w.HasValue()); + + p1.SetValue(10); + p2.SetValue(-3); + } + + + Y_UNIT_TEST(TestManyUnsignaledWithVector) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto p3 = NewPromise(); + auto w = NWait::WaitAny(TVector>{ p1.GetFuture(), p2.GetFuture(), p3.GetFuture() }); + UNIT_ASSERT(!w.HasValue() && !w.HasException()); + + constexpr TStringBuf message = "Test exception 4"; + p2.SetException(std::make_exception_ptr(yexception() << message)); + UNIT_ASSERT_EXCEPTION_SATISFIES(w.TryRethrow(), yexception, [message](auto const& e) { + return message == e.what(); + }); + + p1.SetValue(); + p3.SetValue(); + } + + + Y_UNIT_TEST(TestManyMixedWithVector) { + auto p1 = NewPromise(); + auto p2 = NewPromise(); + auto f = MakeFuture(); + auto w = NWait::WaitAny(TVector>{ p1.GetFuture(), p2.GetFuture(), f }); + UNIT_ASSERT(w.HasValue()); + + p1.SetValue(); + p2.SetValue(); + } + + Y_UNIT_TEST(TestManyStress) { + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAny(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAny(futures); } + , [](size_t) { + return [](auto&& p) { p.SetValue(22); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasValue()); }); + auto e = std::make_exception_ptr(yexception() << "Test exception 5"); + NTest::TestManyStress([](auto&& futures) { return NWait::WaitAny(futures); } + , [e](size_t) { + return [e](auto&& p) { p.SetException(e); }; + } + , [](auto&& waiter) { UNIT_ASSERT(waiter.HasException()); }); + } + +} diff --git a/library/cpp/threading/future/subscription/wait_ut_common.cpp b/library/cpp/threading/future/subscription/wait_ut_common.cpp new file mode 100644 index 00000000000..9f961e73036 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_ut_common.cpp @@ -0,0 +1,26 @@ +#include "wait_ut_common.h" + +#include +#include +#include + +namespace NThreading::NTest::NPrivate { + +void ExecuteAndWait(TVector> jobs, TFuture waiter, size_t threads) { + Y_ENSURE(threads > 0); + Shuffle(jobs.begin(), jobs.end()); + auto pool = CreateThreadPool(threads); + TManualEvent start; + for (auto& j : jobs) { + pool->SafeAddFunc( + [&start, job = std::move(j)]() { + start.WaitI(); + job(); + }); + } + start.Signal(); + waiter.Wait(); + pool->Stop(); +} + +} diff --git a/library/cpp/threading/future/subscription/wait_ut_common.h b/library/cpp/threading/future/subscription/wait_ut_common.h new file mode 100644 index 00000000000..99530dd1f67 --- /dev/null +++ b/library/cpp/threading/future/subscription/wait_ut_common.h @@ -0,0 +1,56 @@ +#pragma once + +#include +#include + +#include + +#include +#include + +namespace NThreading::NTest { + +namespace NPrivate { + +void ExecuteAndWait(TVector> jobs, TFuture waiter, size_t threads); + +template +void SetConcurrentAndWait(TPromises&& promises, FSetter&& setter, TFuture waiter, size_t threads = 8) { + TVector> jobs; + jobs.reserve(std::size(promises)); + for (auto& p : promises) { + jobs.push_back([p, setter]() mutable {setter(p); }); + } + ExecuteAndWait(std::move(jobs), std::move(waiter), threads); +} + +template +auto MakePromise() { + if constexpr (std::is_same_v) { + return NewPromise(); + } + return NewPromise(); +} + +} + +template +void TestManyStress(FWaiterFactory&& waiterFactory, FSetterFactory&& setterFactory, FChecker&& checker) { + for (size_t i : { 1, 2, 4, 8, 16, 32, 64, 128, 256 }) { + TVector> promises; + TVector> futures; + promises.reserve(i); + futures.reserve(i); + for (size_t j = 0; j < i; ++j) { + auto promise = NPrivate::MakePromise(); + futures.push_back(promise.GetFuture()); + promises.push_back(std::move(promise)); + } + auto waiter = waiterFactory(futures); + NPrivate::SetConcurrentAndWait(std::move(promises), [valueSetter = setterFactory(i)](auto&& p) { valueSetter(p); } + , waiter); + checker(waiter); + } +} + +} diff --git a/library/cpp/threading/future/subscription/ya.make b/library/cpp/threading/future/subscription/ya.make new file mode 100644 index 00000000000..759c80f3394 --- /dev/null +++ b/library/cpp/threading/future/subscription/ya.make @@ -0,0 +1,18 @@ +LIBRARY() + +SRCS( + subscription.cpp + wait_all.cpp + wait_all_or_exception.cpp + wait_any.cpp +) + +PEERDIR( + library/cpp/threading/future +) + +END() + +RECURSE_FOR_TESTS( + ut +) diff --git a/src/api/client/yc_private/accessservice/sensitive.proto b/src/api/client/yc_private/accessservice/sensitive.proto new file mode 100644 index 00000000000..009620bd991 --- /dev/null +++ b/src/api/client/yc_private/accessservice/sensitive.proto @@ -0,0 +1,29 @@ +syntax = "proto3"; + +// Based on: +// https://bb.yandexcloud.net/projects/CLOUD/repos/cloud-go/browse/private-api/yandex/cloud/priv/sensitive.proto + +package yandex.cloud; + +import "google/protobuf/descriptor.proto"; + +option go_package = "cloud/proto_extensions"; + +enum SensitiveType { + SENSITIVE_TYPE_UNSPECIFIED = 0; + SENSITIVE_CRC = 1; + SENSITIVE_IAM_TOKEN = 2; + SENSITIVE_REMOVE = 3; + SENSITIVE_YANDEX_PASSPORT_OAUTH_TOKEN = 4; + SENSITIVE_IAM_COOKIE = 5; + SENSITIVE_REFRESH_TOKEN = 6; + SENSITIVE_SESSION_TOKEN = 7; +} + +extend google.protobuf.FieldOptions { + // novikoff: + // Sensitive fields are hidden in logs + // For now could be applied only to string fields + bool sensitive = 110601; + SensitiveType sensitive_type = 110602; +} diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index cb5cc8eb0a5..4bc2136a99b 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,6 +3,8 @@ #include +#include + #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. diff --git a/src/client/topic/impl/CMakeLists.txt b/src/client/topic/impl/CMakeLists.txt index 6f3bc99989b..96fc3b6ba8b 100644 --- a/src/client/topic/impl/CMakeLists.txt +++ b/src/client/topic/impl/CMakeLists.txt @@ -9,6 +9,7 @@ target_link_libraries(client-ydb_topic-impl PUBLIC persqueue-obfuscate api-grpc-draft api-grpc + threading-future-subscription impl-internal-make_request client-ydb_common_client-impl client-ydb_driver From a9073ef96579ffada90532eeba365acbbb1196dd Mon Sep 17 00:00:00 2001 From: Bulat Gayazov Date: Fri, 13 Feb 2026 16:15:06 +0000 Subject: [PATCH 26/93] Fixed topic cmake build --- .github/scripts/copy_sources.sh | 5 +- src/client/topic/impl/write_session.cpp | 64 ++++++++++++------------- src/client/topic/impl/write_session.h | 15 +++--- 3 files changed, 39 insertions(+), 45 deletions(-) diff --git a/.github/scripts/copy_sources.sh b/.github/scripts/copy_sources.sh index f3778bf21d8..3d648a7776c 100755 --- a/.github/scripts/copy_sources.sh +++ b/.github/scripts/copy_sources.sh @@ -7,7 +7,6 @@ echo "Copying sources..." cp -r $1/ydb/public/sdk/cpp/* $tmp_dir echo "tmp_dir: $tmp_dir" -rm -r $tmp_dir/client rm -r $tmp_dir/src/client/arrow rm -r $tmp_dir/src/client/cms rm -r $tmp_dir/src/client/config @@ -23,6 +22,7 @@ rm -r $tmp_dir/include/ydb-cpp-sdk/client/draft rm -r $tmp_dir/tests/unit/client/draft mkdir -p $tmp_dir/src/api/client/yc_private +mkdir -p $tmp_dir/src/api/client/yc_private/accessservice mkdir -p $tmp_dir/src/api/client/yc_public cp -r $1/ydb/public/api/client/yc_private/accessservice/sensitive.proto $tmp_dir/src/api/client/yc_private/accessservice/sensitive.proto @@ -34,7 +34,7 @@ cp -r $1/ydb/public/api/grpc $tmp_dir/src/api cp -r $1/ydb/public/api/protos $tmp_dir/src/api rm -r $tmp_dir/src/api/protos/out -rm $tmp_dir/include/ydb-cpp-sdk/type_switcher.h $tmp_dir/include/ydb-cpp-sdk/client/proto/private.h $tmp_dir/src/version.h +rm $tmp_dir/include/ydb-cpp-sdk/type_switcher.h $tmp_dir/src/version.h cp -r $2/util $tmp_dir cp -r $2/library $tmp_dir @@ -58,7 +58,6 @@ cp $2/tests/slo_workloads/.dockerignore $tmp_dir/tests/slo_workloads cp $2/tests/slo_workloads/Dockerfile $tmp_dir/tests/slo_workloads cp $2/include/ydb-cpp-sdk/type_switcher.h $tmp_dir/include/ydb-cpp-sdk/type_switcher.h -cp $2/include/ydb-cpp-sdk/client/proto/private.h $tmp_dir/include/ydb-cpp-sdk/client/proto/private.h cp $2/src/version.h $tmp_dir/src/version.h cd $2 diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index d9070eb7e5c..dd93430b12d 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -1,5 +1,6 @@ #include "write_session.h" +#include #include #include #include @@ -101,8 +102,8 @@ std::string TKeyedWriteSessionSettings::DefaultPartitioningKeyHasher(const std:: const std::uint64_t lo = MurmurHash(key.data(), key.size(), std::uint64_t{0}); const std::uint64_t hi = MurmurHash(key.data(), key.size(), std::uint64_t{0x9E3779B97F4A7C15ull}); // fixed seed - const std::uint64_t hiBe = y_absl::gntohll(hi); - const std::uint64_t loBe = y_absl::gntohll(lo); + const std::uint64_t hiBe = InetToHost(hi); + const std::uint64_t loBe = InetToHost(lo); std::string out; out.resize(16); @@ -317,7 +318,7 @@ void TKeyedWriteSession::TSplittedPartitionWorker::HandleDescribeResult() { if (newPartitionsIds.empty()) { // describe response is incomplete, we need to resend describe request MoveTo(EState::Init); - Y_ABORT_UNLESS(++Retries < 20, "Too many retries for partition %lu", PartitionId); + Y_ABORT_UNLESS(++Retries < 40, "Too many retries for partition %u", PartitionId); LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Describe response is incomplete, we need to resend describe request for partition " << PartitionId); return; } @@ -434,7 +435,7 @@ TKeyedWriteSession::TEventsWorker::TEventsWorker(TKeyedWriteSession* session) } void TKeyedWriteSession::TEventsWorker::HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); queueIt->second.push_back(TWriteSessionEvent::TEvent(std::move(event))); } @@ -658,12 +659,12 @@ bool TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { } std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueBegin(std::uint32_t partition) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); return queueIt->second.begin(); } std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueEnd(std::uint32_t partition) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition, std::list()); + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); return queueIt->second.end(); } @@ -947,10 +948,10 @@ bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wra void TKeyedWriteSession::TMessagesWorker::PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message) { auto iter = InFlightMessages.insert(InFlightMessages.end(), std::move(message)); - auto [inFlightMessagesIndexIt, _] = InFlightMessagesIndex.try_emplace(partition, std::list()); + auto [inFlightMessagesIndexIt, _] = InFlightMessagesIndex.try_emplace(partition); inFlightMessagesIndexIt->second.push_back(iter); - auto [pendingMessagesIndexIt, __] = PendingMessagesIndex.try_emplace(partition, std::list()); + auto [pendingMessagesIndexIt, __] = PendingMessagesIndex.try_emplace(partition); pendingMessagesIndexIt->second.push_back(iter); } @@ -1006,7 +1007,7 @@ std::optional TKeyedWriteSession::TMessagesWorker::GetContin } void TKeyedWriteSession::TMessagesWorker::HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken) { - auto [it, _] = ContinuationTokens.try_emplace(partition, std::deque()); + auto [it, _] = ContinuationTokens.try_emplace(partition); it->second.push_back(std::move(continuationToken)); } @@ -1091,11 +1092,11 @@ void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint32_t p list.erase(resendIt, list.end()); for (const auto& [newPartition, msgIt] : messagesFromOldPartition) { - auto [inFlightMessagesIndexChainIt, _] = InFlightMessagesIndex.try_emplace(newPartition, std::list()); + auto [inFlightMessagesIndexChainIt, _] = InFlightMessagesIndex.try_emplace(newPartition); inFlightMessagesIndexChainIt->second.push_back(msgIt); if (msgIt->Sent) { - auto [messagesToResendChainIt, __] = MessagesToResendIndex.try_emplace(newPartition, std::list()); + auto [messagesToResendChainIt, __] = MessagesToResendIndex.try_emplace(newPartition); messagesToResendChainIt->second.push_back(msgIt); } } @@ -1104,16 +1105,16 @@ void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint32_t p } void TKeyedWriteSession::TMessagesWorker::RebuildPendingMessagesIndex(std::uint32_t partition) { - auto [oldPendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(partition, std::list()); + auto [oldPendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(partition); std::unordered_map> pendingMessagesForNewPartitions; for (auto it = oldPendingMessagesIndexChainIt->second.begin(); it != oldPendingMessagesIndexChainIt->second.end(); ++it) { auto newPartition = Session->PartitionChooser->ChoosePartition((*it)->Key); - auto [pendingMessagesForNewPartitionsIt, __] = pendingMessagesForNewPartitions.try_emplace(newPartition, std::list()); + auto [pendingMessagesForNewPartitionsIt, __] = pendingMessagesForNewPartitions.try_emplace(newPartition); pendingMessagesForNewPartitionsIt->second.push_back(*it); } for (const auto& [newPartition, pendingMessagesForNewPartition] : pendingMessagesForNewPartitions) { - auto [pendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(newPartition, std::list()); + auto [pendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(newPartition); for (auto reverseIt = pendingMessagesForNewPartition.rbegin(); reverseIt != pendingMessagesForNewPartition.rend(); ++reverseIt) { pendingMessagesIndexChainIt->second.push_front(*reverseIt); } @@ -1159,17 +1160,17 @@ typename TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::IRetryState::TPtr TK //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TKeyedWriteSession::Metrics -void TKeyedWriteSession::MetricGauge::Add(std::uint64_t value) { +void TKeyedWriteSession::TMetricGauge::Add(std::uint64_t value) { Sum += value; MetricCount++; } -void TKeyedWriteSession::MetricGauge::Clear() { +void TKeyedWriteSession::TMetricGauge::Clear() { Sum = 0; MetricCount = 0; } -long double TKeyedWriteSession::MetricGauge::Average() { +long double TKeyedWriteSession::TMetricGauge::Average() { if (MetricCount == 0) { return 0; } @@ -1177,24 +1178,24 @@ long double TKeyedWriteSession::MetricGauge::Average() { return (long double)Sum / (long double)MetricCount; } -TKeyedWriteSession::Metrics::Metrics(TKeyedWriteSession* session): Session(session) {} +TKeyedWriteSession::TMetrics::TMetrics(TKeyedWriteSession* session): Session(session) {} -void TKeyedWriteSession::Metrics::AddMainWorkerTime(std::uint64_t ms) { +void TKeyedWriteSession::TMetrics::AddMainWorkerTime(std::uint64_t ms) { std::lock_guard lock(Lock); MainWorkerTimeMs.Add(ms); } -void TKeyedWriteSession::Metrics::AddCycleTime(std::uint64_t ms) { +void TKeyedWriteSession::TMetrics::AddCycleTime(std::uint64_t ms) { std::lock_guard lock(Lock); CycleTimeMs.Add(ms); } -void TKeyedWriteSession::Metrics::AddWriteLag(std::uint64_t lagMs) { +void TKeyedWriteSession::TMetrics::AddWriteLag(std::uint64_t lagMs) { std::lock_guard lock(Lock); WriteLagMs.Add(lagMs); } -void TKeyedWriteSession::Metrics::PrintMetrics() { +void TKeyedWriteSession::TMetrics::PrintMetrics() { std::lock_guard lock(Lock); LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "METRICS: MainWorkerTimeMs: " << MainWorkerTimeMs.Average() << " ms, CycleTimeMs: " << CycleTimeMs.Average() << " ms, WriteLagMs: " << WriteLagMs.Average() << " ms"); MainWorkerTimeMs.Clear(); @@ -1268,19 +1269,13 @@ TKeyedWriteSession::TKeyedWriteSession( EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); } else { - if (!Settings.EventHandlers_.SessionClosedHandler_) { - EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::SessionClosed); - } else { + if (Settings.EventHandlers_.SessionClosedHandler_) { EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); } - if (!Settings.EventHandlers_.ReadyToAcceptHandler_) { - EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); - } else { + if (Settings.EventHandlers_.ReadyToAcceptHandler_) { EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); } - if (!Settings.EventHandlers_.AcksHandler_) { - EventTypesWithoutHandlers.push_back(TEventsWorker::EEventType::Ack); - } else { + if (Settings.EventHandlers_.AcksHandler_) { EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); } } @@ -1418,7 +1413,7 @@ std::optional TKeyedWriteSession::GetEvent(bool bloc return std::nullopt; } - return EventsWorker->GetEvent(block, EventTypesWithoutHandlers); + return EventsWorker->GetEvent(block); } std::vector TKeyedWriteSession::GetEvents(bool block, std::optional maxEventsCount) { @@ -1426,7 +1421,7 @@ std::vector TKeyedWriteSession::GetEvents(bool block return {}; } - return EventsWorker->GetEvents(block, maxEventsCount, EventTypesWithoutHandlers); + return EventsWorker->GetEvents(block, maxEventsCount); } TDuration TKeyedWriteSession::GetCloseTimeout() { @@ -1658,6 +1653,7 @@ TInstant TKeyedWriteSession::GetCloseDeadline() { } void TKeyedWriteSession::HandleAutoPartitioning(std::uint32_t partition) { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "HandleAutoPartitioning: " << partition); auto splittedPartitionWorker = std::make_shared(this, partition); SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); } @@ -1888,4 +1884,4 @@ TWriterCounters::TPtr TSimpleBlockingKeyedWriteSession::GetCounters() { return nullptr; } -} // namespace NYdb::inline V3::NTopic +} // namespace NYdb::inline V3::NTopic \ No newline at end of file diff --git a/src/client/topic/impl/write_session.h b/src/client/topic/impl/write_session.h index 61e314129a0..441d675a508 100644 --- a/src/client/topic/impl/write_session.h +++ b/src/client/topic/impl/write_session.h @@ -332,7 +332,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - struct MetricGauge { + struct TMetricGauge { std::uint64_t MetricCount = 0; std::uint64_t Sum = 0; @@ -341,12 +341,12 @@ class TKeyedWriteSession : public IKeyedWriteSession, void Clear(); }; - struct Metrics { - Metrics(TKeyedWriteSession* session); + struct TMetrics { + TMetrics(TKeyedWriteSession* session); - MetricGauge MainWorkerTimeMs; - MetricGauge CycleTimeMs; - MetricGauge WriteLagMs; + TMetricGauge MainWorkerTimeMs; + TMetricGauge CycleTimeMs; + TMetricGauge WriteLagMs; std::mutex Lock; TKeyedWriteSession* Session; @@ -410,7 +410,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::shared_ptr Client; TDbDriverStatePtr DbDriverState; - Metrics Metrics; + TMetrics Metrics; std::unordered_map Partitions; std::map PartitionsIndex; @@ -445,7 +445,6 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::atomic Epoch = 0; static constexpr size_t MAX_EPOCH = 1'000'000'000; - std::vector EventTypesWithoutHandlers; std::vector EventTypesWithHandlers; }; From cb8db7400ee84433af7cc290833cd00e557ebf7c Mon Sep 17 00:00:00 2001 From: Bulat Date: Sat, 14 Feb 2026 02:55:48 +0300 Subject: [PATCH 27/93] Update version to v3.14.0 --- src/version.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/version.h b/src/version.h index ff7de6c1f8f..1b6feda781e 100644 --- a/src/version.h +++ b/src/version.h @@ -2,7 +2,7 @@ namespace NYdb { -inline const char* YDB_SDK_VERSION = "3.13.0"; +inline const char* YDB_SDK_VERSION = "3.14.0"; inline const char* YDB_CERTIFICATE_FILE_KEY = "ydb_root_ca_v3.pem"; } // namespace NYdb From 8fa5e72b093db86107ac207f8fe06d06c9a92727 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Tue, 10 Mar 2026 07:48:30 +0000 Subject: [PATCH 28/93] LOGBROKER-10206 Fix after review (#34064) --- .github/last_commit.txt | 2 +- src/client/impl/internal/common/client_pid.cpp | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 1604fb18241..aee4631adc9 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -3b235ed1f2fc3977cfc6f99a74123c0097ef9795 +de0728225782a78a37252da6c80da1f88359924f diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index 4bc2136a99b..cb5cc8eb0a5 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,8 +3,6 @@ #include -#include - #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. From 932ddcd4ac88bbcd259037ca625b1e3525855208 Mon Sep 17 00:00:00 2001 From: Maksim Date: Tue, 10 Mar 2026 07:48:37 +0000 Subject: [PATCH 29/93] [NBS-6956] Add nbs partition in-mem (#33486) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_nbs.proto | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index aee4631adc9..0612a065c90 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -de0728225782a78a37252da6c80da1f88359924f +b3a37eff6a25f54114cceb2a3e3e945a70f766da diff --git a/src/api/protos/draft/ydb_nbs.proto b/src/api/protos/draft/ydb_nbs.proto index b1b41915040..97d7077a96c 100644 --- a/src/api/protos/draft/ydb_nbs.proto +++ b/src/api/protos/draft/ydb_nbs.proto @@ -13,6 +13,12 @@ import "src/api/protos/ydb_operation.proto"; //////////////////////////////////////////////////////////////////////////////// // Partition create request/response. +enum StorageMediaKind { + STORAGE_MEDIA_DEFAULT = 0; + STORAGE_MEDIA_SSD = 1; + STORAGE_MEDIA_MEMORY = 2; +} + message CreatePartitionRequest { Ydb.Operations.OperationParams operation_params = 1; @@ -24,6 +30,9 @@ message CreatePartitionRequest { // Maximum number of blocks stored in partition. uint64 BlocksCount = 4; + + // Storage media type + StorageMediaKind StorageMedia = 5; } message CreatePartitionResponse { From d212e36a435cee3400d0ae929c4e041cf9775d40 Mon Sep 17 00:00:00 2001 From: ubyte Date: Tue, 10 Mar 2026 07:48:44 +0000 Subject: [PATCH 30/93] PQv1: passthrough the '_advanced_monitoring' attribute (#34128) --- .github/last_commit.txt | 2 +- src/client/persqueue_public/impl/persqueue.cpp | 2 ++ src/client/persqueue_public/impl/persqueue_impl.h | 1 + src/client/persqueue_public/include/control_plane.h | 4 ++++ 4 files changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 0612a065c90..1752c9c4604 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -b3a37eff6a25f54114cceb2a3e3e945a70f766da +fd07c393b5c0c58c315dd64c9fcb369d0bf1aff8 diff --git a/src/client/persqueue_public/impl/persqueue.cpp b/src/client/persqueue_public/impl/persqueue.cpp index fc8970bb685..4829a76a6a6 100644 --- a/src/client/persqueue_public/impl/persqueue.cpp +++ b/src/client/persqueue_public/impl/persqueue.cpp @@ -124,6 +124,8 @@ TDescribeTopicResult::TTopicSettings::TTopicSettings(const Ydb::PersQueue::V1::T AbcSlug_ = pair.second; } else if (pair.first == "_federation_account") { FederationAccount_ = pair.second; + } else if (pair.first == "_advanced_monitoring") { + AdvancedMonitoringSettings_ = pair.second; } } for (const auto& readRule : settings.read_rules()) { diff --git a/src/client/persqueue_public/impl/persqueue_impl.h b/src/client/persqueue_public/impl/persqueue_impl.h index c2a08680134..01baccc86f9 100644 --- a/src/client/persqueue_public/impl/persqueue_impl.h +++ b/src/client/persqueue_public/impl/persqueue_impl.h @@ -101,6 +101,7 @@ class TPersQueueClient::TImpl : public TClientImplCommon, AbcSlug); GETTER(std::optional, FederationAccount); GETTER(std::optional, MetricsLevel); + GETTER(std::optional, AdvancedMonitoringSettings); const std::vector& ReadRules() const { return ReadRules_; @@ -148,6 +149,7 @@ struct TDescribeTopicResult : public TStatus { std::optional AbcSlug_; std::string FederationAccount_; std::optional MetricsLevel_; + std::optional AdvancedMonitoringSettings_; std::optional MaxPartitionsCount_; std::optional StabilizationWindow_; @@ -254,6 +256,7 @@ struct TTopicSettings : public TOperationRequestSettings { FLUENT_SETTING_OPTIONAL(std::string, AbcSlug); FLUENT_SETTING_OPTIONAL(std::string, FederationAccount); FLUENT_SETTING_OPTIONAL(ui32, MetricsLevel); + FLUENT_SETTING_OPTIONAL(std::string, AdvancedMonitoringSettings); //TODO: FLUENT_SETTING_VECTOR FLUENT_SETTING_DEFAULT(std::vector, ReadRules, {}); @@ -279,6 +282,7 @@ struct TTopicSettings : public TOperationRequestSettings { AbcSlug_ = settings.AbcSlug(); FederationAccount_ = settings.FederationAccount(); MetricsLevel_ = settings.MetricsLevel(); + AdvancedMonitoringSettings_ = settings.AdvancedMonitoringSettings(); ReadRules_.clear(); for (const auto& readRule : settings.ReadRules()) { ReadRules_.push_back({}); From 9dc92c2cac8fcb80c8fc3a769fa0dff019b997c0 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Tue, 10 Mar 2026 07:48:51 +0000 Subject: [PATCH 31/93] LOGBROKER-10206 Add include (#34147) --- .github/last_commit.txt | 2 +- src/client/topic/impl/write_session.cpp | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 1752c9c4604..60aa8939a11 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -fd07c393b5c0c58c315dd64c9fcb369d0bf1aff8 +458428195c3350fa928f8b93264d43507c823615 diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index dd93430b12d..0170131de30 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -1,3 +1,4 @@ +#include #include "write_session.h" #include From 7a421cda0597fd9269b0de7274f30dab4cde5192 Mon Sep 17 00:00:00 2001 From: ubyte Date: Tue, 10 Mar 2026 07:48:58 +0000 Subject: [PATCH 32/93] Consumers partition metrics: add unittest for the '_advanced_monitoring' attribute read/write operations (#34230) --- .github/last_commit.txt | 2 +- .../persqueue_public/ut/basic_usage_ut.cpp | 44 ++++++++++++++++++- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 60aa8939a11..5b5d75eb747 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -458428195c3350fa928f8b93264d43507c823615 +aa6b205fe3a51e97857940d24ad73619f4b81b67 diff --git a/src/client/persqueue_public/ut/basic_usage_ut.cpp b/src/client/persqueue_public/ut/basic_usage_ut.cpp index d0fd95a2e59..5be709615a2 100644 --- a/src/client/persqueue_public/ut/basic_usage_ut.cpp +++ b/src/client/persqueue_public/ut/basic_usage_ut.cpp @@ -6,10 +6,11 @@ #include #include +#include +#include +#include #include #include -#include -#include using namespace NThreading; using namespace NKikimr; @@ -589,5 +590,44 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_VALUES_EQUAL_C(descr.TopicSettings().MetricsLevel().value(), metricsLevel, descr.GetIssues().ToOneLineString()); } } + + Y_UNIT_TEST(CreateTopicWithCustomAdvancedMonitoringSettings) { + auto setup = std::make_shared(TEST_CASE_NAME); + auto& client = setup->GetPersQueueClient(); + const TString name = "test-topic-" + ToString(TInstant::Now().Seconds()); + const TString path = setup->GetServer().ServerSettings.PQConfig.GetRoot() + "/" + ::NPersQueue::BuildFullTopicName(name, setup->GetLocalCluster()); + + const TVector rrSettings{ + {NYdb::NPersQueue::TReadRuleSettings{}.ConsumerName("shared/user")}, + }; + { + NPersQueue::TCreateTopicSettings settings{}; + settings.ReadRules(rrSettings); + const auto creat = client.CreateTopic(path, settings).GetValueSync(); + UNIT_ASSERT_C(creat.IsSuccess(), creat.GetIssues().ToOneLineString()); + } + { + const auto descr = client.DescribeTopic(path).GetValueSync(); + UNIT_ASSERT_C(descr.IsSuccess(), descr.GetIssues().ToOneLineString()); + UNIT_ASSERT_C(!descr.TopicSettings().AdvancedMonitoringSettings().has_value(), descr.GetIssues().ToOneLineString()); + } + for (const bool enabled: {true, false}) { + NPersQueue::TAlterTopicSettings settings{}; + settings.ReadRules(rrSettings); + const std::string m = R"-({"shared/user":{"metrics_level":3,"monitoring_project_id":"mon"}})-"; + if (enabled) { + settings.AdvancedMonitoringSettings(m); + } + const auto alter = client.AlterTopic(path, settings).GetValueSync(); + UNIT_ASSERT_C(alter.IsSuccess(), LabeledOutput(enabled, alter.GetIssues().ToOneLineString())); + + const auto descr = client.DescribeTopic(path).GetValueSync(); + UNIT_ASSERT_C(descr.IsSuccess(), LabeledOutput(enabled, descr.GetIssues().ToOneLineString())); + UNIT_ASSERT_VALUES_EQUAL_C(descr.TopicSettings().AdvancedMonitoringSettings().has_value(), enabled, LabeledOutput(enabled, descr.GetIssues().ToOneLineString())); + if (enabled) { + UNIT_ASSERT_VALUES_EQUAL_C(descr.TopicSettings().AdvancedMonitoringSettings().value(), m, LabeledOutput(enabled, descr.GetIssues().ToOneLineString())); + } + } + } } } // namespace NYdb::NPersQueue::NTests From 7fba2504fe14a71bde8ad36e135b72ebfdd7a662 Mon Sep 17 00:00:00 2001 From: ubyte Date: Tue, 10 Mar 2026 07:49:05 +0000 Subject: [PATCH 33/93] remove unused TThrRefBase base class from the public interface ISimpleBlockingWriteSession (#33600) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/topic/write_session.h | 2 +- src/client/persqueue_public/include/write_session.h | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 5b5d75eb747..af5de897cef 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -aa6b205fe3a51e97857940d24ad73619f4b81b67 +17d32dc597c509083938e2be2f19d38d2225e43a diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index b3bba9904d8..ceef9701fe9 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -235,7 +235,7 @@ struct TWriteMessage { }; //! Simple write session. Does not need event handlers. Does not provide Events, ContinuationTokens, write Acks. -class ISimpleBlockingWriteSession : public TThrRefBase { +class ISimpleBlockingWriteSession { public: //! Write single message. Blocks for up to blockTimeout if inflight is full or memoryUsage is exceeded; //! return - true if write succeeded, false if message was not enqueued for write within blockTimeout. diff --git a/src/client/persqueue_public/include/write_session.h b/src/client/persqueue_public/include/write_session.h index 98852c46488..5f7c0d49d00 100644 --- a/src/client/persqueue_public/include/write_session.h +++ b/src/client/persqueue_public/include/write_session.h @@ -137,7 +137,7 @@ struct TWriteSessionSettings : public TRequestSettings { }; //! Simple write session. Does not need event handlers. Does not provide Events, ContinuationTokens, write Acks. -class ISimpleBlockingWriteSession : public TThrRefBase { +class ISimpleBlockingWriteSession { public: //! Write single message. Blocks for up to blockTimeout if inflight is full or memoryUsage is exceeded; //! return - true if write succeeded, false if message was not enqueued for write within blockTimeout. From ee6a96233df98ec3100968e16a8f4c188738bb2d Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Tue, 10 Mar 2026 07:49:11 +0000 Subject: [PATCH 34/93] LOGBROKER-10206 Add metrics & some fixes (#34200) --- .github/last_commit.txt | 2 +- src/client/topic/impl/write_session.cpp | 103 +++++++++++++++++++++++- src/client/topic/impl/write_session.h | 21 ++++- src/client/topic/ut/basic_usage_ut.cpp | 27 +++++++ 4 files changed, 147 insertions(+), 6 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index af5de897cef..e63e80fc5f8 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -17d32dc597c509083938e2be2f19d38d2225e43a +939db69461aa364b0509804001a2f71ba623d9e5 diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 0170131de30..698c5122291 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -126,6 +126,10 @@ bool TKeyedWriteSession::TPartitionInfo::InRange(const std::string_view key) con return true; } +bool TKeyedWriteSession::TPartitionInfo::IsSplitted() const { + return !Children_.empty(); +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TKeyedWriteSession::TMessageInfo @@ -293,9 +297,19 @@ void TKeyedWriteSession::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t bool TKeyedWriteSession::TSplittedPartitionWorker::IsDone() { std::lock_guard lock(Lock); + DoneAt = TInstant::Now(); return State == EState::Done; } +bool TKeyedWriteSession::TSplittedPartitionWorker::CanBeRemoved() { + std::lock_guard lock(Lock); + if (State != EState::Done) { + return false; + } + + return TInstant::Now() - DoneAt > TDuration::Seconds(10); +} + bool TKeyedWriteSession::TSplittedPartitionWorker::IsInit() { std::lock_guard lock(Lock); return State == EState::Init; @@ -515,7 +529,7 @@ std::optional> TKeyedWriteSession::TEventsWorker::DoW } void TKeyedWriteSession::TEventsWorker::SubscribeToPartition(std::uint32_t partition) { - if (auto it = Session->SplittedPartitionWorkers.find(partition); it != Session->SplittedPartitionWorkers.end()) { + if (Session->Partitions[partition].IsSplitted() || Session->SplittedPartitionWorkers.contains(partition)) { Session->Partitions[partition].Future(NThreading::MakeFuture()); return; } @@ -552,11 +566,13 @@ std::optional> TKeyedWriteSession::TEventsWorker::Han return EventsPromise; } + Session->Metrics.IncBufferFull(); return std::nullopt; } void TKeyedWriteSession::TEventsWorker::AddReadyToAcceptEvent() { EventsOutputQueue.push_back(TWriteSessionEvent::TReadyToAcceptEvent(IssueContinuationToken())); + Session->Metrics.IncContinuationTokensSent(); } bool TKeyedWriteSession::TEventsWorker::AddSessionClosedIfNeeded() { @@ -654,6 +670,7 @@ bool TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { if (shouldAddReadyToAcceptEvent) { AddReadyToAcceptEvent(); + Session->Metrics.IncContinuationTokensSent(); } return eventsTransferred; @@ -851,6 +868,9 @@ void TKeyedWriteSession::TSessionsWorker::DoWork() { } const auto partition = (*it)->Session->Partition; + if (Session->Partitions[partition].Locked_) { + continue; + } // Remove idle tracking first to keep containers consistent even if the session // is already absent from SessionsIndex. @@ -860,7 +880,7 @@ void TKeyedWriteSession::TSessionsWorker::DoWork() { auto sessionIter = SessionsIndex.find(partition); if (sessionIter != SessionsIndex.end()) { sessionIter->second->IdleSession.reset(); - DestroyWriteSession(sessionIter, TDuration::Zero()); + DestroyWriteSession(sessionIter, TDuration::Zero(), !Session->SplittedPartitionWorkers.contains(partition)); } } } @@ -943,6 +963,7 @@ bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wra return false; } + Session->Metrics.IncOutgoingMessages(); wrappedSession->Session->Write(std::move(*continuationToken), message.BuildMessage(), message.Tx); return true; } @@ -1163,12 +1184,14 @@ typename TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::IRetryState::TPtr TK void TKeyedWriteSession::TMetricGauge::Add(std::uint64_t value) { Sum += value; + Max = std::max(Max, value); MetricCount++; } void TKeyedWriteSession::TMetricGauge::Clear() { Sum = 0; MetricCount = 0; + Max = 0; } long double TKeyedWriteSession::TMetricGauge::Average() { @@ -1179,6 +1202,14 @@ long double TKeyedWriteSession::TMetricGauge::Average() { return (long double)Sum / (long double)MetricCount; } +std::uint64_t TKeyedWriteSession::TMetricGauge::GetMax() const { + return Max; +} + +std::uint64_t TKeyedWriteSession::TMetricGauge::GetSum() const { + return Sum; +} + TKeyedWriteSession::TMetrics::TMetrics(TKeyedWriteSession* session): Session(session) {} void TKeyedWriteSession::TMetrics::AddMainWorkerTime(std::uint64_t ms) { @@ -1196,12 +1227,49 @@ void TKeyedWriteSession::TMetrics::AddWriteLag(std::uint64_t lagMs) { WriteLagMs.Add(lagMs); } +void TKeyedWriteSession::TMetrics::IncContinuationTokensSent() { + std::lock_guard lock(Lock); + ContinuationTokensSent.Add(1); +} + +void TKeyedWriteSession::TMetrics::IncBufferFull() { + std::lock_guard lock(Lock); + BufferFull.Add(1); +} + +void TKeyedWriteSession::TMetrics::IncIncomingMessages() { + std::lock_guard lock(Lock); + IncomingMessages.Add(1); +} + +void TKeyedWriteSession::TMetrics::IncOutgoingMessages() { + std::lock_guard lock(Lock); + OutgoingMessages.Add(1); +} + void TKeyedWriteSession::TMetrics::PrintMetrics() { std::lock_guard lock(Lock); - LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "METRICS: MainWorkerTimeMs: " << MainWorkerTimeMs.Average() << " ms, CycleTimeMs: " << CycleTimeMs.Average() << " ms, WriteLagMs: " << WriteLagMs.Average() << " ms"); + LOG_LAZY( + Session->DbDriverState->Log, + TLOG_ERR, + Session->LogPrefix() + << "METRICS: average MainWorkerTimeMs: " << MainWorkerTimeMs.Average() + << " ms, average CycleTimeMs: " << CycleTimeMs.Average() + << " ms, average WriteLagMs: " << WriteLagMs.Average() << " ms, " + << "max MainWorkerTimeMs: " << MainWorkerTimeMs.GetMax() << " ms, " + << "max CycleTimeMs: " << CycleTimeMs.GetMax() << " ms, " + << "max WriteLagMs: " << WriteLagMs.GetMax() << " ms, " + << "ContinuationTokensSent: " << ContinuationTokensSent.GetSum() << " tokens, " + << "BufferFull: " << BufferFull.GetSum() << " times, " + << "IncomingMessages: " << IncomingMessages.GetSum() << " messages, " + << "OutgoingMessages: " << OutgoingMessages.GetSum() << " messages"); MainWorkerTimeMs.Clear(); CycleTimeMs.Clear(); WriteLagMs.Clear(); + ContinuationTokensSent.Clear(); + BufferFull.Clear(); + IncomingMessages.Clear(); + OutgoingMessages.Clear(); } @@ -1343,10 +1411,19 @@ std::vector TKeyedWriteSession::GetPartition return partitions; } +std::unordered_map TKeyedWriteSession::GetPartitionsMap() const { + return Partitions; +} + +std::map TKeyedWriteSession::GetPartitionsIndex() const { + return PartitionsIndex; +} + void TKeyedWriteSession::Write(TContinuationToken&&, const std::string& key, TWriteMessage&& message, TTransactionBase* tx) { std::optional> eventsPromise; { std::lock_guard lock(GlobalLock); + Metrics.IncIncomingMessages(); if (Closed.load()) { return; } @@ -1435,13 +1512,15 @@ TDuration TKeyedWriteSession::GetCloseTimeout() { } bool TKeyedWriteSession::RunSplittedPartitionWorkers() { - if (SplittedPartitionWorkers.empty()) { + if (SplittedPartitionWorkers.empty() && ReadySplittedPartitionWorkers.empty()) { return false; } bool needRerun = false; + std::unordered_map> readySplittedPartitionWorkers; for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { if (splittedPartitionWorker->IsDone()) { + readySplittedPartitionWorkers[partition] = splittedPartitionWorker; continue; } @@ -1450,6 +1529,22 @@ bool TKeyedWriteSession::RunSplittedPartitionWorkers() { needRerun = needRerun || splittedPartitionWorker->IsDone(); } + for (const auto& [partition, splittedPartitionWorker] : readySplittedPartitionWorkers) { + ReadySplittedPartitionWorkers[partition] = splittedPartitionWorker; + SplittedPartitionWorkers.erase(partition); + } + + std::vector partitionsToRemove; + for (const auto& [partition, splittedPartitionWorker] : ReadySplittedPartitionWorkers) { + if (splittedPartitionWorker->CanBeRemoved()) { + partitionsToRemove.push_back(partition); + } + } + + for (const auto& partition : partitionsToRemove) { + ReadySplittedPartitionWorkers.erase(partition); + } + return needRerun; } diff --git a/src/client/topic/impl/write_session.h b/src/client/topic/impl/write_session.h index 441d675a508..5d47017f21b 100644 --- a/src/client/topic/impl/write_session.h +++ b/src/client/topic/impl/write_session.h @@ -75,6 +75,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, using TSelf = TPartitionInfo; bool InRange(const std::string_view key) const; + bool IsSplitted() const; FLUENT_SETTING(std::string, FromBound); FLUENT_SETTING(std::optional, ToBound); @@ -248,6 +249,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, void DoWork(); bool IsDone(); bool IsInit(); + bool CanBeRemoved(); std::string GetStateName() const; private: @@ -261,6 +263,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::mutex Lock; std::uint64_t NotReadyFutures = 0; size_t Retries = 0; + TInstant DoneAt = TInstant::Max(); }; struct TEventsWorker : public std::enable_shared_from_this { @@ -335,9 +338,12 @@ class TKeyedWriteSession : public IKeyedWriteSession, struct TMetricGauge { std::uint64_t MetricCount = 0; std::uint64_t Sum = 0; + std::uint64_t Max = 0; long double Average(); void Add(std::uint64_t value); + std::uint64_t GetMax() const; + std::uint64_t GetSum() const; void Clear(); }; @@ -347,12 +353,20 @@ class TKeyedWriteSession : public IKeyedWriteSession, TMetricGauge MainWorkerTimeMs; TMetricGauge CycleTimeMs; TMetricGauge WriteLagMs; + TMetricGauge ContinuationTokensSent; + TMetricGauge BufferFull; + TMetricGauge IncomingMessages; + TMetricGauge OutgoingMessages; std::mutex Lock; TKeyedWriteSession* Session; void AddMainWorkerTime(std::uint64_t ms); void AddCycleTime(std::uint64_t ms); void AddWriteLag(std::uint64_t lagMs); + void IncContinuationTokensSent(); + void IncBufferFull(); + void IncIncomingMessages(); + void IncOutgoingMessages(); void PrintMetrics(); }; @@ -403,6 +417,10 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::vector GetPartitions() const; + std::unordered_map GetPartitionsMap() const; + + std::map GetPartitionsIndex() const; + ~TKeyedWriteSession(); private: @@ -413,7 +431,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, TMetrics Metrics; std::unordered_map Partitions; - std::map PartitionsIndex; + std::map PartitionsIndex; TKeyedWriteSessionSettings Settings; ESeqNoStrategy SeqNoStrategy = ESeqNoStrategy::NotInitialized; @@ -435,6 +453,7 @@ class TKeyedWriteSession : public IKeyedWriteSession, std::shared_ptr EventsWorker; std::shared_ptr SessionsWorker; std::unordered_map> SplittedPartitionWorkers; + std::unordered_map> ReadySplittedPartitionWorkers; std::shared_ptr MessagesWorker; std::shared_ptr RetryPolicy; diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index c43f79a8697..35fdf76c816 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -1717,6 +1717,33 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } } + TKeyedWriteSessionSettings writeSettings3 = writeSettings1; + writeSettings3.ProducerIdPrefix("autopartitioning_keyed_3"); + auto session3 = client.CreateKeyedWriteSession(writeSettings3); + + auto partitionsMap1 = dynamic_cast(session1.get())->GetPartitionsMap(); + auto partitionsMap3 = dynamic_cast(session3.get())->GetPartitionsMap(); + + for (const auto& [partitionId, partitionInfo] : partitionsMap1) { + auto partitionInfo3 = partitionsMap3.find(partitionId); + UNIT_ASSERT_C(partitionInfo3 != partitionsMap3.end(), + "Partition " << partitionId << " is not in session3"); + UNIT_ASSERT_C(partitionInfo.FromBound_ == partitionInfo3->second.FromBound_, + "From bound is not equal for partition " << partitionId); + UNIT_ASSERT_C(partitionInfo.ToBound_ == partitionInfo3->second.ToBound_, + "To bound is not equal for partition " << partitionId); + } + + auto partitionsIndex1 = dynamic_cast(session1.get())->GetPartitionsIndex(); + auto partitionsIndex3 = dynamic_cast(session3.get())->GetPartitionsIndex(); + for (const auto& [key, partitionId] : partitionsIndex1) { + auto partitionId3 = partitionsIndex3.find(key); + UNIT_ASSERT_C(partitionId3 != partitionsIndex3.end(), + "Key " << key << " is not in session3"); + UNIT_ASSERT_EQUAL_C(partitionId, partitionId3->second, + "Partition id is not equal for key " << key); + } + UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), 14, "Expected exactly 14 distinct acks, each seqNo exactly once; got " << ackedSeqNos.size()); auto sessionPartitions = dynamic_cast(session1.get())->GetPartitions(); From 0675af0bd8e62e59da6c194c55e64182fe74bd75 Mon Sep 17 00:00:00 2001 From: BarkovBG Date: Tue, 10 Mar 2026 07:49:18 +0000 Subject: [PATCH 35/93] NBS2: add "get load actor adapter actor id" request to dstool; return the test; remove actor logging from dbg (#34092) --- .github/last_commit.txt | 2 +- src/api/grpc/draft/ydb_nbs_v1.proto | 3 +++ src/api/protos/draft/ydb_nbs.proto | 18 ++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e63e80fc5f8..6fb9d8baf9a 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -939db69461aa364b0509804001a2f71ba623d9e5 +24a37d26ad5efc64648c333c16d3237049d1dcb9 diff --git a/src/api/grpc/draft/ydb_nbs_v1.proto b/src/api/grpc/draft/ydb_nbs_v1.proto index 5e5c333e014..9d8569b7b46 100644 --- a/src/api/grpc/draft/ydb_nbs_v1.proto +++ b/src/api/grpc/draft/ydb_nbs_v1.proto @@ -16,6 +16,9 @@ service NbsService { // Delete NBS partition rpc DeletePartition(Ydb.Nbs.DeletePartitionRequest) returns (Ydb.Nbs.DeletePartitionResponse); + // Get load actor adapter actor id for partition + rpc GetLoadActorAdapterActorId(Ydb.Nbs.GetLoadActorAdapterActorIdRequest) returns (Ydb.Nbs.GetLoadActorAdapterActorIdResponse); + // List NBS partitions rpc ListPartitions(Ydb.Nbs.ListPartitionsRequest) returns (Ydb.Nbs.ListPartitionsResponse); diff --git a/src/api/protos/draft/ydb_nbs.proto b/src/api/protos/draft/ydb_nbs.proto index 97d7077a96c..f323b291b7d 100644 --- a/src/api/protos/draft/ydb_nbs.proto +++ b/src/api/protos/draft/ydb_nbs.proto @@ -61,6 +61,24 @@ message DeletePartitionResult { string TabletId = 1; } +//////////////////////////////////////////////////////////////////////////////// +// Get load actor adapter actor id request/response. + +message GetLoadActorAdapterActorIdRequest { + Ydb.Operations.OperationParams operation_params = 1; + // Partition tablet id (actor id string). + string TabletId = 2; +} + +message GetLoadActorAdapterActorIdResponse { + Ydb.Operations.Operation operation = 1; +} + +message GetLoadActorAdapterActorIdResult { + // Load actor adapter actor id string. + string ActorId = 1; +} + //////////////////////////////////////////////////////////////////////////////// // Partition list request/response. From 550c27742538403528daa2311bb35886fdb0d951 Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Tue, 10 Mar 2026 07:49:25 +0000 Subject: [PATCH 36/93] =?UTF-8?q?Fixed=20sanitizer=20error=20in=20ydb/core?= =?UTF-8?q?/persqueue/pqtablet/partition/mlp/ut=E2=80=A6=20(#34357)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/last_commit.txt | 2 +- .../persqueue_public/ut/ut_utils/test_server.h | 12 +++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 6fb9d8baf9a..828b5a7612b 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -24a37d26ad5efc64648c333c16d3237049d1dcb9 +774e3910f10e59d3bd9837dfe75a1025fed962b0 diff --git a/src/client/persqueue_public/ut/ut_utils/test_server.h b/src/client/persqueue_public/ut/ut_utils/test_server.h index 80abcba47fc..a2c731384c6 100644 --- a/src/client/persqueue_public/ut/ut_utils/test_server.h +++ b/src/client/persqueue_public/ut/ut_utils/test_server.h @@ -44,6 +44,14 @@ class TTestServer { { } + TTestServer(TTestServer&&) = default; + TTestServer& operator=(TTestServer&&) = default; + + ~TTestServer() { + ShutdownGRpc(); + ShutdownServer(); + } + void StartServer(bool doClientInit = true, TMaybe databaseName = Nothing()) { Log.SetFormatter([](ELogPriority priority, TStringBuf message) { return TStringBuilder() << TInstant::Now() << " " << priority << ": " << message << Endl; @@ -64,7 +72,9 @@ class TTestServer { } void ShutdownGRpc() { - CleverServer->ShutdownGRpc(); + if (CleverServer) { + CleverServer->ShutdownGRpc(); + } } void EnableGRpc() { From 5c095883461a97755a04bf660e98e1ffe6ad6f68 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Tue, 10 Mar 2026 07:49:32 +0000 Subject: [PATCH 37/93] LOGBROKER-7430 fixed partition reading hanging in topic sdk (#34362) --- .github/last_commit.txt | 2 +- .../persqueue_public/ut/read_session_ut.cpp | 9 +- src/client/topic/impl/common.h | 11 ++ src/client/topic/impl/read_session_impl.h | 62 ++++++-- src/client/topic/impl/read_session_impl.ipp | 143 ++++++++++++++++-- 5 files changed, 193 insertions(+), 34 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 828b5a7612b..1b44aa6987a 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -774e3910f10e59d3bd9837dfe75a1025fed962b0 +9c0a1bceb7dcc45e1e32aa9e33cfba9b55ec3b64 diff --git a/src/client/persqueue_public/ut/read_session_ut.cpp b/src/client/persqueue_public/ut/read_session_ut.cpp index 32d8d55149c..071ca522a9d 100644 --- a/src/client/persqueue_public/ut/read_session_ut.cpp +++ b/src/client/persqueue_public/ut/read_session_ut.cpp @@ -1904,14 +1904,15 @@ Y_UNIT_TEST_SUITE(ReadSessionImplTest) { 0); std::atomic ready = true; + std::atomic abandoned = false; - stream->InsertDataEvent(0, 0, data, ready); + stream->InsertDataEvent(0, 0, data, ready, abandoned); stream->InsertEvent(TServiceEvent{stream, 0, 0, 0, {}}); - stream->InsertDataEvent(0, 0, data, ready); - stream->InsertDataEvent(0, 0, data, ready); + stream->InsertDataEvent(0, 0, data, ready, abandoned); + stream->InsertDataEvent(0, 0, data, ready, abandoned); stream->InsertEvent(TServiceEvent{stream, 0, 0, 0, {}}); stream->InsertEvent(TServiceEvent{stream, 0, 0, 0, {}}); - stream->InsertDataEvent(0, 0, data, ready); + stream->InsertDataEvent(0, 0, data, ready, abandoned); TDeferredActions actions; diff --git a/src/client/topic/impl/common.h b/src/client/topic/impl/common.h index 0445589cbda..b429750aeaa 100644 --- a/src/client/topic/impl/common.h +++ b/src/client/topic/impl/common.h @@ -394,6 +394,9 @@ class TBaseSessionEventsQueue : public ISignalable { } } + virtual void SelfCheck() { + } + public: NThreading::TFuture WaitEvent() { { @@ -401,6 +404,11 @@ class TBaseSessionEventsQueue : public ISignalable { if (HasEventsImpl()) { return NThreading::MakeFuture(); // Signalled } else { + if (const auto now = TInstant::Now(); now - LastSelfCheckAt > TDuration::Seconds(10)) { + LastSelfCheckAt = now; + SelfCheck(); + } + Y_ABORT_UNLESS(Waiter.Valid()); auto res = Waiter.GetFuture(); return res; @@ -421,6 +429,9 @@ class TBaseSessionEventsQueue : public ISignalable { std::mutex Mutex; std::optional CloseEvent; std::atomic Closed = false; + +private: + TInstant LastSelfCheckAt = TInstant::Now(); }; } // namespace NYdb::NTopic diff --git a/src/client/topic/impl/read_session_impl.h b/src/client/topic/impl/read_session_impl.h index ea06593aecc..a7809af3eb0 100644 --- a/src/client/topic/impl/read_session_impl.h +++ b/src/client/topic/impl/read_session_impl.h @@ -175,6 +175,7 @@ class TDeferredActions { void Reconnect(); void SignalWaiters(); void StartSessions(); + void DestroyDecompressionInfos(); private: // Read. @@ -236,6 +237,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this& deferred); @@ -302,17 +305,12 @@ class TDataDecompressionInfo : public std::enable_shared_from_this(GetServerMessage().batches_size()); - } - - bool HasReadyUnreadData() const; - void PutDecompressionError(std::exception_ptr error, size_t batch, size_t message); std::exception_ptr GetDecompressionError(size_t batch, size_t message); void OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount); void OnUserRetrievedEvent(i64 decompressedDataSize, size_t messagesCount); + void OnTaskCanceled(i64 sourceSize, size_t messagesCount); private: // Special struct for marking (batch/message) as ready. @@ -320,6 +318,7 @@ class TDataDecompressionInfo : public std::enable_shared_from_this Ready = false; + std::atomic Abandoned = false; // Marked by true either when decompression is completed or message is not needed anymore }; struct TDecompressionTask { @@ -371,7 +370,6 @@ class TDataDecompressionInfo : public std::enable_shared_from_this SourceDataNotProcessed = 0; std::pair CurrentDecompressingMessage = {0, 0}; // (Batch, Message) std::deque ReadyThresholds; - std::pair CurrentReadingMessage = {0, 0}; // (Batch, Message) // Decompression exceptions. // Optimization for rare using. @@ -389,11 +387,12 @@ class TDataDecompressionInfo : public std::enable_shared_from_this class TDataDecompressionEvent { public: - TDataDecompressionEvent(size_t batch, size_t message, TDataDecompressionInfoPtr parent, std::atomic& ready) : + TDataDecompressionEvent(size_t batch, size_t message, TDataDecompressionInfoPtr parent, std::atomic& ready, std::atomic& abandoned) : Batch{batch}, Message{message}, Parent{std::move(parent)}, - Ready{ready} + Ready{ready}, + Abandoned{abandoned} { } @@ -401,6 +400,15 @@ class TDataDecompressionEvent { return Ready; } + bool SetAbandoned() { + if (bool expected = false; Abandoned.compare_exchange_strong(expected, true)) { + return true; + } + + Y_ASSERT(Ready); + return false; + } + void TakeData(TIntrusivePtr> partitionStream, std::vector::TMessage>& messages, std::vector::TCompressedMessage>& compressedMessages, @@ -416,6 +424,7 @@ class TDataDecompressionEvent { size_t Message; TDataDecompressionInfoPtr Parent; std::atomic& Ready; + std::atomic& Abandoned; }; template @@ -480,12 +489,14 @@ struct TRawPartitionStreamEvent { TRawPartitionStreamEvent(size_t batch, size_t message, TDataDecompressionInfoPtr parent, - std::atomic &ready) + std::atomic &ready, + std::atomic& abandoned) : Event(std::in_place_type_t>(), batch, message, std::move(parent), - ready) + ready, + abandoned) { } @@ -499,6 +510,11 @@ struct TRawPartitionStreamEvent { return std::holds_alternative>(Event); } + TDataDecompressionEvent& GetDataEvent() { + Y_ASSERT(IsDataEvent()); + return std::get>(Event); + } + const TDataDecompressionEvent& GetDataEvent() const { Y_ASSERT(IsDataEvent()); return std::get>(Event); @@ -579,6 +595,10 @@ class TRawPartitionStreamEventQueue { std::vector::TCompressedMessage>& compressedMessages, TUserRetrievedEventsInfoAccumulator& accumulator); + ui64 GetReadyEventsCount() const { + return Ready.size(); + } + private: static void GetDataEventImpl(TIntrusivePtr> partitionStream, size_t& maxEventsCount, @@ -717,9 +737,10 @@ class TPartitionStreamImpl : public TAPartitionStream { void InsertDataEvent(size_t batch, size_t message, TDataDecompressionInfoPtr parent, - std::atomic &ready) + std::atomic& ready, + std::atomic& abandoned) { - EventsQueue.emplace_back(batch, message, std::move(parent), ready); + EventsQueue.emplace_back(batch, message, std::move(parent), ready, abandoned); } bool HasEvents() const { @@ -825,6 +846,10 @@ class TPartitionStreamImpl : public TAPartitionStream { return Lock; } + ui64 GetReadyEventsCount() const { + return EventsQueue.GetReadyEventsCount(); + } + private: const TKey Key; ui64 AssignId; @@ -920,7 +945,8 @@ class TReadSessionEventsQueue: public TBaseSessionEventsQueue parent, - std::atomic &ready); + std::atomic& ready, + std::atomic& abandoned); void SignalEventImpl(TIntrusivePtr> partitionStream, TDeferredActions& deferred, @@ -1085,6 +1111,10 @@ class TReadSessionEventsQueue: public TBaseSessionEventsQueueLockShared()->SelfCheck(); + } + private: bool HasEventCallbacks; TCallbackContextPtr CbContext; @@ -1183,6 +1213,7 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext& deferred); void OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount, i64 serverBytesSize = 0); + void OnDecompressionTaskCanceled(i64 sourceSize, size_t messagesCount, i64 serverBytesSize); TReadSessionEventsQueue* GetEventsQueue() { return EventsQueue.get(); @@ -1242,6 +1273,8 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext client); + void SelfCheck(); + private: void BreakConnectionAndReconnectImpl(TPlainStatus&& status, TDeferredActions& deferred); @@ -1433,6 +1466,7 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext::IRetryState::TPtr RetryState; // Current retry state (if now we are (re)connecting). size_t ConnectionAttemptsDone = 0; + TInstant LastActiveTime = TInstant::Now(); // Memory usage. i64 CompressedDataSize = 0; diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index e49e76b4f03..6b659254767 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -211,7 +211,15 @@ void TRawPartitionStreamEventQueue::DeleteNotReadyTail(TDe std::deque> ready; auto i = NotReady.begin(); - for (; (i != NotReady.end()) && i->IsReady(); ++i) { + for (; i != NotReady.end(); ++i) { + if (!i->IsReady()) { + // Cancel inflight decompression tasks if any + if (!i->IsDataEvent() || i->GetDataEvent().SetAbandoned()) { + break; + } + // Message was decompressed and become ready + } + ready.push_back(std::move(*i)); } @@ -846,6 +854,7 @@ void TSingleClusterReadSessionImpl::OnUserRetrievedEvent(i Y_ABORT_UNLESS(decompressedSize <= DecompressedDataSize); DecompressedDataSize -= decompressedSize; + LastActiveTime = TInstant::Now(); ContinueReadingDataImpl(); StartDecompressionTasksImpl(deferred); @@ -899,6 +908,7 @@ void TSingleClusterReadSessionImpl::ReadFromProcessorImpl( } }; + LastActiveTime = TInstant::Now(); deferred.DeferReadFromProcessor(Processor, ServerMessage.get(), std::move(callback)); } } @@ -1749,6 +1759,7 @@ void TSingleClusterReadSessionImpl::StartDecompressionTask deferred); DecompressedDataSize += sentToDecompress; if (current.BatchInfo->AllDecompressionTasksStarted()) { + deferred.DeferDestroyDecompressionInfos({current.BatchInfo}); DecompressionQueue.pop_front(); } else { break; @@ -1828,8 +1839,9 @@ void TSingleClusterReadSessionImpl::OnDataDecompressed(i64 UpdateMemoryUsageStatisticsImpl(); CompressedDataSize -= sourceSize; DecompressedDataSize += decompressedSize - estimatedDecompressedSize; + LastActiveTime = TInstant::Now(); constexpr double weight = 0.6; - if (sourceSize > 0) { + if (sourceSize > 0 && decompressedSize > 0) { AverageCompressionRatio = weight * static_cast(decompressedSize) / static_cast(sourceSize) + (1 - weight) * AverageCompressionRatio; } if (Aborting) { @@ -1843,6 +1855,34 @@ void TSingleClusterReadSessionImpl::OnDataDecompressed(i64 StartDecompressionTasksImpl(deferred); } +template +void TSingleClusterReadSessionImpl::OnDecompressionTaskCanceled(i64 sourceSize, size_t messagesCount, i64 serverBytesSize) { + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "The application data decompression is cancelled. Number of messages " << messagesCount << ", size " << sourceSize << " bytes"); + + *Settings.Counters_->BytesReadCompressed += sourceSize; + *Settings.Counters_->MessagesRead += messagesCount; + *Settings.Counters_->BytesInflightCompressed -= sourceSize; + *Settings.Counters_->BytesInflightTotal -= sourceSize; + *Settings.Counters_->MessagesInflight -= messagesCount; + + TDeferredActions deferred; + std::lock_guard guard(Lock); + UpdateMemoryUsageStatisticsImpl(); + + CompressedDataSize -= sourceSize; + LastActiveTime = TInstant::Now(); + if (Aborting) { + return; + } + if constexpr (!UseMigrationProtocol) { + LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "Returning serverBytesSize = " << serverBytesSize << " to budget"); + ReadSizeBudget += serverBytesSize; + } + + ContinueReadingDataImpl(); + StartDecompressionTasksImpl(deferred); +} + template void TSingleClusterReadSessionImpl::Abort() { LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "Abort session to cluster"); @@ -2141,6 +2181,48 @@ bool TSingleClusterReadSessionImpl::AllParentSessionsHasBe return true; } +template +void TSingleClusterReadSessionImpl::SelfCheck() { + const auto delta = TInstant::Now() - LastActiveTime; + if (delta < TDuration::Minutes(1)) { + // Session ok, we got at least one event from server since last 1 minute + return; + } + + if (WaitingReadResponse) { + // We sent to server non zero memory budget, but don't get read response since last 1 minute + LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "[SelfCheck] There is no server events since last: " << delta << ", most likely there is no data in topic partitions"); + return; + } + + if (DecompressionTasksInflight) { + LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "[SelfCheck] There is still inflight decompression tasks since last: " << delta << ", read session was stopped by back pressure, most likely decompression is too slow"); + return; + } + + ui64 readyEventsCount = 0; + std::unordered_set handledPartitionStreams; + for (const auto& [assignId, stream] : PartitionStreams) { + handledPartitionStreams.emplace(assignId); + readyEventsCount += stream->GetReadyEventsCount(); + } + + // Some streams may have some ready events after finish + for (const auto& decompressionItem : DecompressionQueue) { + if (const auto& stream = *decompressionItem.PartitionStream; handledPartitionStreams.emplace(stream.GetAssignId()).second) { + readyEventsCount += stream.GetReadyEventsCount(); + } + } + + if (readyEventsCount) { + LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "[SelfCheck] There is still " << readyEventsCount << " pending ready events since last: " << delta << ", read session was stopped by back pressure, most likely extraction pipeline is too slow"); + return; + } + + // No read from server inflight, no ready events and no decompression tasks inflight => most likely hanging + LOG_LAZY(Log, TLOG_WARNING, GetLogPrefix() << "[SelfCheck] There is no ready events / inflight decompression since last: " << delta << ", most likely hanged after stop by back pressure"); +} + template <> inline void TSingleClusterReadSessionImpl::ConfirmPartitionStreamEnd(TPartitionStreamImpl* partitionStream, std::span childIds) { { @@ -2335,6 +2417,7 @@ bool TReadSessionEventsQueue::PushEvent(TIntrusivePtr(event)) { stream->DeleteNotReadyTail(deferred); + SignalReadyEventsImpl(stream, deferred); } if (!HasDataEventCallback() && !std::holds_alternative>(event)) { @@ -2384,14 +2467,15 @@ bool TReadSessionEventsQueue::PushDataEvent(TIntrusivePtr< size_t batch, size_t message, TDataDecompressionInfoPtr parent, - std::atomic& ready) + std::atomic& ready, + std::atomic& abandoned) { std::lock_guard guard(TParent::Mutex); if (this->Closed) { return false; } - partitionStream->InsertDataEvent(batch, message, parent, ready); + partitionStream->InsertDataEvent(batch, message, parent, ready, abandoned); return true; } @@ -2732,6 +2816,18 @@ TDataDecompressionInfo::~TDataDecompressionInfo() } } +template +void TDataDecompressionInfo::Cleanup() { + auto session = CbContext->LockShared(); + Y_ASSERT(session); + + while (!Tasks.empty()) { + const auto& task = Tasks.front(); + OnTaskCanceled(task.AddedDataSize(), task.AddedMessagesCount()); + Tasks.pop_front(); + } +} + template void TDataDecompressionInfo::BuildBatchesMeta() { BatchesMeta.reserve(ServerMessage.batches_size()); @@ -2849,7 +2945,8 @@ void TDataDecompressionInfo::PlanDecompressionTasks(double CurrentDecompressingMessage.first, CurrentDecompressingMessage.second, TDataDecompressionInfo::shared_from_this(), - ReadyThresholds.back().Ready); + ReadyThresholds.back().Ready, + ReadyThresholds.back().Abandoned); if (!pushRes) { session->AbortImpl(); return; @@ -2977,15 +3074,6 @@ void TDataDecompressionEvent::TakeData(TIntrusivePtr -bool TDataDecompressionInfo::HasReadyUnreadData() const { - std::optional> threshold = GetReadyThreshold(); - if (!threshold) { - return false; - } - return CurrentReadingMessage <= *threshold; -} - template void TDataDecompressionInfo::OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount) { @@ -3010,6 +3098,18 @@ void TDataDecompressionInfo::OnUserRetrievedEvent(i64 deco } } +template +void TDataDecompressionInfo::OnTaskCanceled(i64 sourceSize, size_t messagesCount) +{ + SourceDataNotProcessed -= sourceSize; + CompressedDataSize -= sourceSize; + MessagesInflight -= messagesCount; + + if (auto session = CbContext->LockShared()) { + session->OnDecompressionTaskCanceled(sourceSize, messagesCount, ServerBytesSize.exchange(0)); + } +} + template void TDataDecompressionInfo::TDecompressionTask::Add(size_t batch, size_t message, size_t sourceDataSize, @@ -3109,6 +3209,11 @@ void TDataDecompressionInfo::TDecompressionTask::operator( if (auto session = parent->CbContext->LockShared()) { session->GetEventsQueue()->SignalReadyEvents(PartitionStream); } + + if (bool expected = false; !Ready->Abandoned.compare_exchange_strong(expected, true)) { + // Message is dropped due to partition stream cancellation, we should release decompressed memory + parent->OnUserRetrievedEvent(DecompressedSize, messagesProcessed); + } } template @@ -3230,7 +3335,7 @@ void TDeferredActions::DeferSignalWaiter(TWaiter&& waiter) template void TDeferredActions::DeferDestroyDecompressionInfos(std::vector>&& infos) { - DecompressionInfos = std::move(infos); + DecompressionInfos.insert(DecompressionInfos.end(), infos.begin(), infos.end()); } template @@ -3244,6 +3349,7 @@ void TDeferredActions::DoActions() { Reconnect(); SignalWaiters(); StartSessions(); + DestroyDecompressionInfos(); } template @@ -3329,4 +3435,11 @@ void TDeferredActions::SignalWaiters() { } } +template +void TDeferredActions::DestroyDecompressionInfos() { + for (const auto& info : DecompressionInfos) { + info->Cleanup(); + } +} + } // namespace NYdb::NTopic From c992db0c23cfbad3952c3413767742e07da89805 Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:49:39 +0000 Subject: [PATCH 38/93] Fixed SDK oss build (#34423) --- .github/last_commit.txt | 2 +- src/client/impl/internal/common/client_pid.cpp | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 1b44aa6987a..3bde50d5510 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -9c0a1bceb7dcc45e1e32aa9e33cfba9b55ec3b64 +caef9d7f7f97b623166c54561a1966ebba0006fc diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index cb5cc8eb0a5..4bc2136a99b 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,6 +3,8 @@ #include +#include + #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. From 447cb923e3edc87e8b2c61507bdc34d2ef339500 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Tue, 10 Mar 2026 07:49:45 +0000 Subject: [PATCH 39/93] LOGBROKER-7430 added unit test on reading with restarts (#34487) --- .github/last_commit.txt | 2 +- src/client/topic/impl/read_session_impl.h | 24 ++++-- src/client/topic/impl/read_session_impl.ipp | 76 ++++++++++++++----- src/client/topic/ut/basic_usage_ut.cpp | 47 ++++++++++-- .../ut/ut_utils/topic_sdk_test_setup.cpp | 38 ++++++++++ .../topic/utils/managed_executor.cpp | 6 +- 6 files changed, 159 insertions(+), 34 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 3bde50d5510..e01a1666431 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -caef9d7f7f97b623166c54561a1966ebba0006fc +9b39c77079f5426318342441c8e0150d1d62303a diff --git a/src/client/topic/impl/read_session_impl.h b/src/client/topic/impl/read_session_impl.h index a7809af3eb0..c60558caab0 100644 --- a/src/client/topic/impl/read_session_impl.h +++ b/src/client/topic/impl/read_session_impl.h @@ -275,6 +275,10 @@ class TDataDecompressionInfo : public std::enable_shared_from_this> GetReadyThreshold() const { size_t readyCount = 0; std::pair ret; @@ -405,8 +409,8 @@ class TDataDecompressionEvent { return true; } - Y_ASSERT(Ready); - return false; + // If Ready=false, decompression task already successfully cancelled + return !Ready; } void TakeData(TIntrusivePtr> partitionStream, @@ -415,6 +419,8 @@ class TDataDecompressionEvent { size_t& maxByteSize, size_t& dataSize) const; + size_t GetDataSize() const; + TDataDecompressionInfoPtr GetParent() const { return Parent; } @@ -578,6 +584,10 @@ class TRawPartitionStreamEventQueue { (NotReady.empty() ? Ready : NotReady).pop_back(); } + size_t size() const { + return NotReady.size() + Ready.size(); + } + void clear() noexcept { NotReady.clear(); Ready.clear(); @@ -595,7 +605,7 @@ class TRawPartitionStreamEventQueue { std::vector::TCompressedMessage>& compressedMessages, TUserRetrievedEventsInfoAccumulator& accumulator); - ui64 GetReadyEventsCount() const { + size_t GetReadyEventAmount() const { return Ready.size(); } @@ -846,8 +856,12 @@ class TPartitionStreamImpl : public TAPartitionStream { return Lock; } - ui64 GetReadyEventsCount() const { - return EventsQueue.GetReadyEventsCount(); + size_t GetEventAmount() const { + return EventsQueue.size(); + } + + size_t GetReadyEventAmount() const { + return EventsQueue.GetReadyEventAmount(); } private: diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index 6b659254767..563a44edf6b 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -209,29 +209,35 @@ template void TRawPartitionStreamEventQueue::DeleteNotReadyTail(TDeferredActions& deferred) { std::deque> ready; + std::vector> infos; + TUserRetrievedEventsInfoAccumulator accumulator; - auto i = NotReady.begin(); - for (; i != NotReady.end(); ++i) { - if (!i->IsReady()) { - // Cancel inflight decompression tasks if any - if (!i->IsDataEvent() || i->GetDataEvent().SetAbandoned()) { - break; + bool hasNonReadyEvents = false; + for (auto& event : NotReady) { + const bool isDataEvent = event.IsDataEvent(); + + if (event.IsReady() || + (isDataEvent && !event.GetDataEvent().SetAbandoned()) // Try to cancel inflight decompression tasks if any (returns true if message was decompressed and become ready) + ) { + if (!hasNonReadyEvents) { + // Continue ready events prefix + ready.push_back(std::move(event)); + } else if (isDataEvent) { + // We should release memory for this ready event here + accumulator.Add(event.GetDataEvent().GetParent(), event.GetDataEvent().GetDataSize()); } - // Message was decompressed and become ready - } - - ready.push_back(std::move(*i)); - } - - std::vector> infos; + } else { + hasNonReadyEvents = true; - for (; i != NotReady.end(); ++i) { - if (i->IsDataEvent()) { - infos.push_back(i->GetDataEvent().GetParent()); + if (isDataEvent) { + // We should release memory for non ready data events and cancel decompression tasks + infos.push_back(event.GetDataEvent().GetParent()); + } } } deferred.DeferDestroyDecompressionInfos(std::move(infos)); + accumulator.OnUserRetrievedEvent(); swap(ready, NotReady); } @@ -2184,7 +2190,7 @@ bool TSingleClusterReadSessionImpl::AllParentSessionsHasBe template void TSingleClusterReadSessionImpl::SelfCheck() { const auto delta = TInstant::Now() - LastActiveTime; - if (delta < TDuration::Minutes(1)) { + if (delta < TDuration::Seconds(30)) { // Session ok, we got at least one event from server since last 1 minute return; } @@ -2201,16 +2207,21 @@ void TSingleClusterReadSessionImpl::SelfCheck() { } ui64 readyEventsCount = 0; + ui64 amountEvents = 0; std::unordered_set handledPartitionStreams; for (const auto& [assignId, stream] : PartitionStreams) { handledPartitionStreams.emplace(assignId); - readyEventsCount += stream->GetReadyEventsCount(); + readyEventsCount += stream->GetReadyEventAmount(); + amountEvents += stream->GetEventAmount(); } // Some streams may have some ready events after finish + ui64 notStartedTasks = 0; for (const auto& decompressionItem : DecompressionQueue) { if (const auto& stream = *decompressionItem.PartitionStream; handledPartitionStreams.emplace(stream.GetAssignId()).second) { - readyEventsCount += stream.GetReadyEventsCount(); + readyEventsCount += stream.GetReadyEventAmount(); + amountEvents += stream.GetEventAmount(); + notStartedTasks += decompressionItem.BatchInfo->GetNotStartedTasksCunt(); } } @@ -2220,7 +2231,20 @@ void TSingleClusterReadSessionImpl::SelfCheck() { } // No read from server inflight, no ready events and no decompression tasks inflight => most likely hanging - LOG_LAZY(Log, TLOG_WARNING, GetLogPrefix() << "[SelfCheck] There is no ready events / inflight decompression since last: " << delta << ", most likely hanged after stop by back pressure"); + LOG_LAZY(Log, TLOG_WARNING, GetLogPrefix() + << "[SelfCheck] There is no ready events / inflight decompression since last: " << delta + << ", most likely hung after stop by back pressure. Session state: " + << "grpc connection alive: " << (Processor ? "yes" : "no") + << ", waiting grpc reconnect: " << (ConnectDelayContext && !ConnectDelayContext->IsCancelled() ? "yes" : "no") + << ", reconnect attempts done: " << ConnectionAttemptsDone + << ", pending decompression batches: " << DecompressionQueue.size() + << ", not ready data events: " << amountEvents + << ", alive partition streams: " << PartitionStreams.size() + << ", unhandled compressed data size: " << CompressedDataSize + << ", unhandled decompressed data size: " << DecompressedDataSize + << ", avg compression ratio: " << AverageCompressionRatio + << ", free memory to request: " << ReadSizeServerDelta + << ", not started tasks after reconnect: " << notStartedTasks); } template <> @@ -2821,11 +2845,16 @@ void TDataDecompressionInfo::Cleanup() { auto session = CbContext->LockShared(); Y_ASSERT(session); + ui64 sourceSize = 0; + ui64 messagesCount = 0; while (!Tasks.empty()) { const auto& task = Tasks.front(); - OnTaskCanceled(task.AddedDataSize(), task.AddedMessagesCount()); + sourceSize += task.AddedDataSize(); + messagesCount += task.AddedMessagesCount(); Tasks.pop_front(); } + + OnTaskCanceled(sourceSize, messagesCount); } template @@ -3074,6 +3103,11 @@ void TDataDecompressionEvent::TakeData(TIntrusivePtr +size_t TDataDecompressionEvent::GetDataSize() const { + return Parent->GetServerMessage().batches(Batch).message_data(Message).data().size(); +} + template void TDataDecompressionInfo::OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount) { diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index 35fdf76c816..d9d4b920702 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -164,7 +164,9 @@ void CreateTopicWithAutoPartitioning(TTopicClient& client) { client.CreateTopic(TEST_TOPIC, createSettings).Wait(); } -void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSessionSettings writeSettings, const std::string& message, std::uint32_t count, TTopicSdkTestSetup& setup, std::shared_ptr decompressor) { +void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSessionSettings writeSettings, const std::string& message, std::uint32_t count, + TTopicSdkTestSetup& setup, std::shared_ptr decompressor, ui32 restartPeriod = 7, ui32 maxRestartsCount = 10) +{ auto client = setup.MakeClient(); auto session = client.CreateSimpleBlockingWriteSession(writeSettings); @@ -179,9 +181,10 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess TTopicClient topicClient = setup.MakeClient(); - - auto WaitTasks = [&](auto f, size_t c) { + auto WaitTasks = [&, timeout = TInstant::Now() + TDuration::Seconds(60)](auto f, size_t c) { while (f() < c) { + UNIT_ASSERT(timeout > TInstant::Now()); + ReadSession->WaitEvent(); std::this_thread::sleep_for(100ms); }; }; @@ -206,7 +209,7 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess WaitPlannedTasks(e, n); size_t completed = e->GetExecutedCount(); - setup.GetServer().KillTopicPqrbTablet(setup.GetTopicPath()); + setup.GetServer().KillTopicPqrbTablet(JoinPath({TString(setup.MakeDriverConfig().GetDatabase()), TString(setup.GetTopicPath())})); std::this_thread::sleep_for(100ms); e->StartFuncs(tasks); @@ -228,9 +231,17 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << ">>> TEST: start reading" << Endl; + std::uint32_t i = 0; + ui32 restartCount = 0; while (AtomicGet(lastOffset) + 1 < count) { - RunTasks(decompressor, {i++}); + if (restartCount < maxRestartsCount && i % restartPeriod == 1) { + PlanTasksAndRestart(decompressor, {i++}); + restartCount++; + } else { + RunTasks(decompressor, {i++}); + } } ReadSession->Close(TDuration::MilliSeconds(10)); @@ -895,6 +906,32 @@ Y_UNIT_TEST_SUITE(BasicUsage) { WriteAndReadToEndWithRestarts(readSettings, writeSettings, message, count, setup, decompressor); } + Y_UNIT_TEST(ReadWithRestartsAndLargeData) { + TTopicSdkTestSetup setup(TEST_CASE_NAME); + auto compressor = std::make_shared(); + auto decompressor = CreateThreadPoolManagedExecutor(1); + + TReadSessionSettings readSettings; + readSettings + .ConsumerName(setup.GetConsumerName()) + .MaxMemoryUsageBytes(1_MB) + .DecompressionExecutor(decompressor) + .AppendTopics(setup.GetTopicPath()) + // .DirectRead(EnableDirectRead) + ; + + TWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath()).MessageGroupId(TEST_MESSAGE_GROUP_ID) + .Codec(ECodec::RAW) + .CompressionExecutor(compressor); + + std::uint32_t count = 3000; + std::string message(8'000, 'x'); + + WriteAndReadToEndWithRestarts(readSettings, writeSettings, message, count, setup, decompressor); + } + Y_UNIT_TEST(ConflictingWrites) { TTopicSdkTestSetup setup(TEST_CASE_NAME); diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp index 4272e5defa0..a402943c798 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp @@ -1,14 +1,52 @@ #include "topic_sdk_test_setup.h" +#include using namespace std::chrono_literals; namespace NYdb::inline V3::NTopic::NTests { +namespace { + +void TerminateHandler() { + Cerr << "======= terminate() call stack ========" << Endl; + FormatBackTrace(&Cerr); + if (const auto& backtrace = TBackTrace::FromCurrentException(); backtrace.size() > 0) { + Cerr << "======== exception call stack =========" << Endl; + backtrace.PrintTo(Cerr); + } + Cerr << "=======================================" << Endl; + + if (std::current_exception()) { + Cerr << "Uncaught exception: " << CurrentExceptionMessage() << Endl; + } else { + Cerr << "Terminate for unknown reason (no current exception)" << Endl; + } + + abort(); +} + +void BackTraceSignalHandler(int signal) { + Cerr << "======= Signal " << signal << " call stack ========" << Endl; + FormatBackTrace(&Cerr); + Cerr << "===============================================" << Endl; + + abort(); +} + +} // anonymous namespace + TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NKikimr::Tests::TServerSettings& settings, bool createTopic) : Database_("/Root") , Server_(settings, false) { + NKikimr::EnableYDBBacktraceFormat(); + + std::set_terminate(&TerminateHandler); + for (auto sig : {SIGFPE, SIGILL, SIGSEGV}) { + signal(sig, &BackTraceSignalHandler); + } + Log_.SetFormatter([testCaseName](ELogPriority priority, TStringBuf message) { return TStringBuilder() << TInstant::Now() << " :" << testCaseName << " " << priority << ": " << message << Endl; }); diff --git a/tests/integration/topic/utils/managed_executor.cpp b/tests/integration/topic/utils/managed_executor.cpp index 3e72dbc5a12..c96bc1d0716 100644 --- a/tests/integration/topic/utils/managed_executor.cpp +++ b/tests/integration/topic/utils/managed_executor.cpp @@ -55,10 +55,11 @@ void TManagedExecutor::StartFuncs(const std::vector& indicies) std::lock_guard lock(Mutex); for (auto index : indicies) { - Y_ABORT_UNLESS(index < Funcs.size()); - Y_ABORT_UNLESS(Funcs[index]); + Y_ABORT_UNLESS(index < Funcs.size()); + Y_ABORT_UNLESS(Funcs[index]); RunTask(std::move(Funcs[index])); + Funcs[index] = nullptr; } } @@ -91,6 +92,7 @@ void TManagedExecutor::RunAllTasks() for (auto& func : Funcs) { if (func) { RunTask(std::move(func)); + func = nullptr; } } } From 4d2b6051d0091e8a4d9e27aad7658a93274211ad Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:49:52 +0000 Subject: [PATCH 40/93] [C++ SDK] Fixed tsan fail in kqp/ut/effects (#34319) --- .github/last_commit.txt | 2 +- src/client/query/impl/client_session.cpp | 26 ++++++++++++------------ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e01a1666431..2bc4ef07206 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -9b39c77079f5426318342441c8e0150d1d62303a +84d802b285b54f878802fee1976f3c0bba2c0aad diff --git a/src/client/query/impl/client_session.cpp b/src/client/query/impl/client_session.cpp index c905ad82fcb..7492e0d13e3 100644 --- a/src/client/query/impl/client_session.cpp +++ b/src/client/query/impl/client_session.cpp @@ -12,22 +12,22 @@ namespace NYdb::inline V3::NQuery { // Custom lock primitive to protect session from destroying // during async read execution. // The problem is TSession::TImpl holds grpc stream processor by IntrusivePtr -// and this processor alredy refcounted by internal code. -// That mean during TSession::TImpl dtor no gurantee to grpc procerrot will be destroyed. -// StreamProcessor_->Cancel() doesn't help it just start async cancelation but we have no way -// to wait cancelation has done. -// So we need some way to protect access to row session impl pointer -// from async reader (processor callback). We can't use shared/weak ptr here because TSessionImpl -// stores as uniq ptr inside session pool and as shared ptr in the TSession +// and this processor already refcounted by internal code. +// That mean during TSession::TImpl dtor no guarantee to grpc processor will be destroyed. +// StreamProcessor_->Cancel() doesn't help it just start async cancellation but we have no way +// to wait cancellation has done. +// So we need some way to protect access to raw session impl pointer +// from async reader (processor callback). We can't use shared/weak ptr here because TSession::TImpl +// stores as unique_ptr inside session pool and as shared_ptr in the TSession // when user got session (see GetSmartDeleter related code). -// Why just not std::mutex? - Requirement do not destroy a mutex while it is locked +// Why just not std::mutex? - Requirement does not destroy a mutex while it is locked // makes it difficult to use here. Moreover we need to allow recursive lock. // Why recursive lock? - In happy path we destroy session from CloseFromServer call, // so the session dtor called from thread which already got the lock. -// TODO: Proably we can add sync version of Cancel method in to grpc stream procesor to make sure +// TODO: Probably we can add sync version of Cancel method in to grpc stream processor to make sure // no more callback will be called. class TSafeTSessionImplHolder { @@ -38,6 +38,7 @@ class TSafeTSessionImplHolder { TSafeTSessionImplHolder(TSession::TImpl* p) : Ptr(p) , Semaphore(0) + , OwnerThread(std::thread::id()) {} TSession::TImpl* TrySharedOwning() noexcept { @@ -62,10 +63,9 @@ class TSafeTSessionImplHolder { uint32_t cur = 0; uint32_t newVal = 1; - while (!Semaphore.compare_exchange_weak(cur, newVal, - std::memory_order_release, std::memory_order_relaxed)) { - std::this_thread::yield(); - cur = 0; + while (!Semaphore.compare_exchange_weak(cur, newVal)) { + std::this_thread::yield(); + cur = 0; } } }; From 691636332604fc9f0bd3eeedc0a27ef6627a9116 Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:49:59 +0000 Subject: [PATCH 41/93] [C++ SDK] Created CHANGELOG.md (#34599) --- .github/last_commit.txt | 2 +- CHANGELOG.md | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) create mode 100644 CHANGELOG.md diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 2bc4ef07206..87e72e25913 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -84d802b285b54f878802fee1976f3c0bba2c0aad +ffd8b94a68fe31d0696f0f05eeebd807cd313b2d diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1 @@ + From 84fe8144a754cba113a0bd0c269355698207bf60 Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:50:06 +0000 Subject: [PATCH 42/93] [C++ SDK] Added missing include in topics (#34591) --- .github/last_commit.txt | 2 +- src/client/topic/impl/write_session.cpp | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 87e72e25913..58224f252fc 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -ffd8b94a68fe31d0696f0f05eeebd807cd313b2d +bb617aeed0b3e1e75f7d5715eda1209dd8db5f89 diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 698c5122291..51ecaa233b9 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -1,15 +1,17 @@ -#include #include "write_session.h" -#include #include #include #include #include #include + #include #include +#include + +#include namespace NYdb::inline V3::NTopic { From 8a153446c7ac9ddba15db41ad9828e9537f71c54 Mon Sep 17 00:00:00 2001 From: Ivan Nikolaev Date: Tue, 10 Mar 2026 07:50:13 +0000 Subject: [PATCH 43/93] Support ALTER TABLE COMPACT in kqp and rpc_alter_table (#34516) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_table.proto | 27 +++++++++++++++++++ src/library/operation_id/operation_id.cpp | 3 +++ .../operation_id/protos/operation_id.proto | 1 + 4 files changed, 32 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 58224f252fc..f3b28d1e881 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -bb617aeed0b3e1e75f7d5715eda1209dd8db5f89 +52275aa6371548dfd8c1c52fc6801a9ffd4b6b54 diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 2b0deca192a..34e7c7607ee 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -938,6 +938,31 @@ message RenameIndexItem { bool replace_destination = 3; } +// State of compact operation +message CompactState { + enum State { + STATE_UNSPECIFIED = 0; + STATE_IN_PROGRESS = 1; + STATE_DONE = 2; + STATE_CANCELLED = 3; + } +} + +message CompactItem { + // Compact table with all indices + optional bool cascade = 1; + // Max shards in flight per operation + optional uint32 max_shards_in_flight = 2; +} + +message CompactMetadata { + optional string path = 1; + optional bool cascade = 2; + optional uint32 max_shards_in_flight = 3; + optional CompactState.State state = 4; + optional float progress = 5; +} + // Alter table with given path message AlterTableRequest { // Session identifier @@ -984,6 +1009,8 @@ message AlterTableRequest { // Rename existed index repeated RenameIndexItem rename_indexes = 21; reserved 22, 23; + // Start compaction for table + CompactItem compact = 24; } message AlterTableResponse { diff --git a/src/library/operation_id/operation_id.cpp b/src/library/operation_id/operation_id.cpp index 5fcaebd7d51..c4bea6f98f0 100644 --- a/src/library/operation_id/operation_id.cpp +++ b/src/library/operation_id/operation_id.cpp @@ -72,6 +72,9 @@ std::string ProtoToString(const Ydb::TOperationId& proto) { case Ydb::TOperationId::RESTORE: res << "ydb://restore"; break; + case Ydb::TOperationId::COMPACTION: + res << "ydb://compaction"; + break; default: Y_ABORT_UNLESS(false, "unexpected kind"); } diff --git a/src/library/operation_id/protos/operation_id.proto b/src/library/operation_id/protos/operation_id.proto index 936feb2d93f..c85d5645b94 100644 --- a/src/library/operation_id/protos/operation_id.proto +++ b/src/library/operation_id/protos/operation_id.proto @@ -17,6 +17,7 @@ message TOperationId { SS_BG_TASKS = 10; INCREMENTAL_BACKUP = 11; RESTORE = 12; + COMPACTION = 13; } message TData { From 66cd23ed31b5507e54054db014151aa5f5f37a84 Mon Sep 17 00:00:00 2001 From: Vitaliy Filippov Date: Tue, 10 Mar 2026 07:50:20 +0000 Subject: [PATCH 44/93] Allow to specify impl table proto settings for fulltext_relevance indexes (#34278) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/table/table.h | 8 ++++ src/api/protos/ydb_table.proto | 5 ++- src/client/table/table.cpp | 51 +++++++++++++++++------- 4 files changed, 50 insertions(+), 16 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index f3b28d1e881..d4338699259 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -52275aa6371548dfd8c1c52fc6801a9ffd4b6b54 +d212c86fbf1a58ce11f13ed4da978d3718e706e7 diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index f405fc654af..48c4cee69d9 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -230,6 +230,14 @@ class TReadReplicasSettings { }; struct TGlobalIndexSettings { + static constexpr const int VectorKMeansTreeLevelTablePosition = 0; + static constexpr const int VectorKMeansTreePostingTablePosition = 1; + static constexpr const int VectorKMeansTreePrefixTablePosition = 2; + static constexpr const int FulltextRelevanceDictTablePosition = 0; + static constexpr const int FulltextRelevanceDocsTablePosition = 1; + static constexpr const int FulltextRelevanceStatsTablePosition = 2; + static constexpr const int FulltextRelevancePostingTablePosition = 3; + using TUniformOrExplicitPartitions = std::variant; TPartitioningSettings PartitioningSettings; diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 34e7c7607ee..a70361e7c64 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -318,8 +318,11 @@ message GlobalFulltextPlainIndex { } message GlobalFulltextRelevanceIndex { - GlobalIndexSettings settings = 1; + GlobalIndexSettings posting_table_settings = 1; FulltextIndexSettings fulltext_settings = 2; + GlobalIndexSettings dict_table_settings = 3; + GlobalIndexSettings docs_table_settings = 4; + GlobalIndexSettings stats_table_settings = 5; } // Represent table index diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 8ed110c0147..e793841319a 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -2788,11 +2788,15 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { case TProto::kGlobalVectorKmeansTreeIndex: { type = EIndexType::GlobalVectorKMeansTree; const auto &vectorProto = proto.global_vector_kmeans_tree_index(); - globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(vectorProto.level_table_settings())); - globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(vectorProto.posting_table_settings())); const bool prefixVectorIndex = indexColumns.size() > 1; + globalIndexSettings.resize(prefixVectorIndex ? 3 : 2); + globalIndexSettings[TGlobalIndexSettings::VectorKMeansTreeLevelTablePosition] = + TGlobalIndexSettings::FromProto(vectorProto.level_table_settings()); + globalIndexSettings[TGlobalIndexSettings::VectorKMeansTreePostingTablePosition] = + TGlobalIndexSettings::FromProto(vectorProto.posting_table_settings()); if (prefixVectorIndex) { - globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(vectorProto.prefix_table_settings())); + globalIndexSettings[TGlobalIndexSettings::VectorKMeansTreePrefixTablePosition] = + TGlobalIndexSettings::FromProto(vectorProto.prefix_table_settings()); } specializedIndexSettings = TKMeansTreeSettings::FromProto(vectorProto.vector_settings()); break; @@ -2807,7 +2811,15 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { case TProto::kGlobalFulltextRelevanceIndex: { type = EIndexType::GlobalFulltextRelevance; const auto& fulltextProto = proto.global_fulltext_relevance_index(); - globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(fulltextProto.settings())); + globalIndexSettings.resize(4); + globalIndexSettings[TGlobalIndexSettings::FulltextRelevanceDictTablePosition] = + TGlobalIndexSettings::FromProto(fulltextProto.dict_table_settings()); + globalIndexSettings[TGlobalIndexSettings::FulltextRelevanceDocsTablePosition] = + TGlobalIndexSettings::FromProto(fulltextProto.docs_table_settings()); + globalIndexSettings[TGlobalIndexSettings::FulltextRelevanceStatsTablePosition] = + TGlobalIndexSettings::FromProto(fulltextProto.stats_table_settings()); + globalIndexSettings[TGlobalIndexSettings::FulltextRelevancePostingTablePosition] = + TGlobalIndexSettings::FromProto(fulltextProto.posting_table_settings()); specializedIndexSettings = TFulltextIndexSettings::FromProto(fulltextProto.fulltext_settings()); break; } @@ -2837,19 +2849,19 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { case EIndexType::GlobalSync: { auto& settings = *proto.mutable_global_index()->mutable_settings(); if (GlobalIndexSettings_.size() == 1) - GlobalIndexSettings_[0].SerializeTo(settings); + GlobalIndexSettings_.at(0).SerializeTo(settings); break; } case EIndexType::GlobalAsync: { auto& settings = *proto.mutable_global_async_index()->mutable_settings(); if (GlobalIndexSettings_.size() == 1) - GlobalIndexSettings_[0].SerializeTo(settings); + GlobalIndexSettings_.at(0).SerializeTo(settings); break; } case EIndexType::GlobalUnique: { auto& settings = *proto.mutable_global_unique_index()->mutable_settings(); if (GlobalIndexSettings_.size() == 1) - GlobalIndexSettings_[0].SerializeTo(settings); + GlobalIndexSettings_.at(0).SerializeTo(settings); break; } case EIndexType::GlobalVectorKMeansTree: { @@ -2857,9 +2869,14 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { auto& level_settings = *global_vector_kmeans_tree_index->mutable_level_table_settings(); auto& posting_settings = *global_vector_kmeans_tree_index->mutable_posting_table_settings(); auto& vector_settings = *global_vector_kmeans_tree_index->mutable_vector_settings(); - if (GlobalIndexSettings_.size() == 2) { - GlobalIndexSettings_[0].SerializeTo(level_settings); - GlobalIndexSettings_[1].SerializeTo(posting_settings); + const bool is_prefixed = IndexColumns_.size() > 1; + if (GlobalIndexSettings_.size() == (is_prefixed ? 3 : 2)) { + GlobalIndexSettings_.at(TGlobalIndexSettings::VectorKMeansTreeLevelTablePosition).SerializeTo(level_settings); + GlobalIndexSettings_.at(TGlobalIndexSettings::VectorKMeansTreePostingTablePosition).SerializeTo(posting_settings); + if (is_prefixed) { + auto& prefix_settings = *global_vector_kmeans_tree_index->mutable_prefix_table_settings(); + GlobalIndexSettings_.at(TGlobalIndexSettings::VectorKMeansTreePrefixTablePosition).SerializeTo(prefix_settings); + } } if (const auto* settings = std::get_if(&SpecializedIndexSettings_)) { settings->SerializeTo(vector_settings); @@ -2871,7 +2888,7 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { auto& settings = *global_fulltext_index->mutable_settings(); auto& fulltext_settings = *global_fulltext_index->mutable_fulltext_settings(); if (GlobalIndexSettings_.size() == 1) { - GlobalIndexSettings_[0].SerializeTo(settings); + GlobalIndexSettings_.at(0).SerializeTo(settings); } if (const auto* ftSettings = std::get_if(&SpecializedIndexSettings_)) { ftSettings->SerializeTo(fulltext_settings); @@ -2880,10 +2897,16 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { } case EIndexType::GlobalFulltextRelevance: { auto* global_fulltext_index = proto.mutable_global_fulltext_relevance_index(); - auto& settings = *global_fulltext_index->mutable_settings(); + auto& dict_settings = *global_fulltext_index->mutable_dict_table_settings(); + auto& docs_settings = *global_fulltext_index->mutable_docs_table_settings(); + auto& stats_settings = *global_fulltext_index->mutable_stats_table_settings(); + auto& posting_settings = *global_fulltext_index->mutable_posting_table_settings(); auto& fulltext_settings = *global_fulltext_index->mutable_fulltext_settings(); - if (GlobalIndexSettings_.size() == 1) { - GlobalIndexSettings_[0].SerializeTo(settings); + if (GlobalIndexSettings_.size() == 4) { + GlobalIndexSettings_.at(TGlobalIndexSettings::FulltextRelevanceDictTablePosition).SerializeTo(dict_settings); + GlobalIndexSettings_.at(TGlobalIndexSettings::FulltextRelevanceDocsTablePosition).SerializeTo(docs_settings); + GlobalIndexSettings_.at(TGlobalIndexSettings::FulltextRelevanceStatsTablePosition).SerializeTo(stats_settings); + GlobalIndexSettings_.at(TGlobalIndexSettings::FulltextRelevancePostingTablePosition).SerializeTo(posting_settings); } if (const auto* ftSettings = std::get_if(&SpecializedIndexSettings_)) { ftSettings->SerializeTo(fulltext_settings); From 433eff0e5ba5bf89c8037aa514bb442bdae5da26 Mon Sep 17 00:00:00 2001 From: xyliganSereja Date: Tue, 10 Mar 2026 07:50:26 +0000 Subject: [PATCH 45/93] Columnshard bloom skip index support (#34458) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_table.proto | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index d4338699259..b120b1d5296 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -d212c86fbf1a58ce11f13ed4da978d3718e706e7 +20464fe3faa28991f670f1b39266c9cf2a7de9f1 diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index a70361e7c64..874f9a60dbd 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -325,6 +325,18 @@ message GlobalFulltextRelevanceIndex { GlobalIndexSettings stats_table_settings = 5; } +message LocalBloomFilterIndex { + optional double false_positive_probability = 1; +} + +message LocalBloomNgramFilterIndex { + uint32 ngram_size = 1; + uint32 hashes_count = 2; + uint32 filter_size_bytes = 3; + uint32 records_count = 4; + optional bool case_sensitive = 5; +} + // Represent table index message TableIndex { // Name of index @@ -339,6 +351,8 @@ message TableIndex { GlobalVectorKMeansTreeIndex global_vector_kmeans_tree_index = 7; GlobalFulltextPlainIndex global_fulltext_plain_index = 8; GlobalFulltextRelevanceIndex global_fulltext_relevance_index = 9; + LocalBloomFilterIndex local_bloom_filter_index = 10; + LocalBloomNgramFilterIndex local_bloom_ngram_filter_index = 11; } // list of columns content to be copied in to index table repeated string data_columns = 5; @@ -365,6 +379,8 @@ message TableIndexDescription { GlobalVectorKMeansTreeIndex global_vector_kmeans_tree_index = 9; GlobalFulltextPlainIndex global_fulltext_plain_index = 10; GlobalFulltextRelevanceIndex global_fulltext_relevance_index = 11; + LocalBloomFilterIndex local_bloom_filter_index = 12; + LocalBloomNgramFilterIndex local_bloom_ngram_filter_index = 13; } Status status = 4; // list of columns content to be copied in to index table From bc2d75aaa129bee4c0e025ebadf1e46aeff9a2e6 Mon Sep 17 00:00:00 2001 From: kseleznyov Date: Tue, 10 Mar 2026 07:50:33 +0000 Subject: [PATCH 46/93] Specify user SID in change (CDC) records (#33262) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/table/table.h | 4 ++++ src/api/protos/ydb_table.proto | 4 ++++ src/client/table/table.cpp | 13 +++++++++++++ 4 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index b120b1d5296..2bf40cad798 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -20464fe3faa28991f670f1b39266c9cf2a7de9f1 +006a6f26eb6b0a2ae72f4ff7c5c6a8f39ddf7c18 diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 48c4cee69d9..fb103f0d343 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -466,6 +466,8 @@ class TChangefeedDescription { TChangefeedDescription& WithRetentionPeriod(const TDuration& value); // Initial scan will output the current state of the table first TChangefeedDescription& WithInitialScan(); + // Enable UserSIDs + TChangefeedDescription& WithUserSIDs(); // Attributes TChangefeedDescription& AddAttribute(const std::string& key, const std::string& value); TChangefeedDescription& SetAttributes(const std::unordered_map& attrs); @@ -481,6 +483,7 @@ class TChangefeedDescription { bool GetSchemaChanges() const; const std::optional& GetResolvedTimestamps() const; bool GetInitialScan() const; + bool GetUserSIDs() const; const std::unordered_map& GetAttributes() const; const std::string& GetAwsRegion() const; const std::optional& GetInitialScanProgress() const; @@ -510,6 +513,7 @@ class TChangefeedDescription { std::optional ResolvedTimestamps_; std::optional RetentionPeriod_; bool InitialScan_ = false; + bool UserSIDs_ = false; std::unordered_map Attributes_; std::string AwsRegion_; std::optional InitialScanProgress_; diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 874f9a60dbd..3237af482d4 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -467,6 +467,8 @@ message Changefeed { Topic.PartitioningSettings topic_partitioning_settings = 10; // Emit schema change events or not bool schema_changes = 11; + // Emit user security identifier (SID) or not. + bool user_sids = 12; } message ChangefeedDescription { @@ -510,6 +512,8 @@ message ChangefeedDescription { InitialScanProgress initial_scan_progress = 9; // State of emitting of schema change events bool schema_changes = 10; + // State of emitting of user security identifier (SID) + bool user_sids = 11; } message StoragePool { diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index e793841319a..6e9f1421157 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -3040,6 +3040,11 @@ TChangefeedDescription& TChangefeedDescription::WithInitialScan() { return *this; } +TChangefeedDescription& TChangefeedDescription::WithUserSIDs() { + UserSIDs_ = true; + return *this; +} + TChangefeedDescription& TChangefeedDescription::AddAttribute(const std::string& key, const std::string& value) { Attributes_[key] = value; return *this; @@ -3092,6 +3097,10 @@ bool TChangefeedDescription::GetInitialScan() const { return InitialScan_; } +bool TChangefeedDescription::GetUserSIDs() const { + return UserSIDs_; +} + const std::unordered_map& TChangefeedDescription::GetAttributes() const { return Attributes_; } @@ -3151,6 +3160,9 @@ TChangefeedDescription TChangefeedDescription::FromProto(const TProto& proto) { if (proto.schema_changes()) { ret.WithSchemaChanges(); } + if (proto.user_sids()) { + ret.WithUserSIDs(); + } if (proto.has_resolved_timestamps_interval()) { ret.WithResolvedTimestamps(TDuration::MilliSeconds( ::google::protobuf::util::TimeUtil::DurationToMilliseconds(proto.resolved_timestamps_interval()))); @@ -3196,6 +3208,7 @@ void TChangefeedDescription::SerializeCommonFields(TProto& proto) const { proto.set_virtual_timestamps(VirtualTimestamps_); proto.set_schema_changes(SchemaChanges_); proto.set_aws_region(TStringType{AwsRegion_}); + proto.set_user_sids(UserSIDs_); switch (Mode_) { case EChangefeedMode::KeysOnly: From bdeb6f6bbfddaa76e178205559db6c787866a3fd Mon Sep 17 00:00:00 2001 From: stanislav_shchetinin Date: Tue, 10 Mar 2026 07:50:40 +0000 Subject: [PATCH 47/93] Support include_index_data in ExportFs (#34641) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_export.proto | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 2bf40cad798..ce02ac0e478 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -006a6f26eb6b0a2ae72f4ff7c5c6a8f39ddf7c18 +72acad87a9603dcf39447997fc5f634904c070ff diff --git a/src/api/protos/ydb_export.proto b/src/api/protos/ydb_export.proto index c334944f1c4..1d6b327b7b9 100644 --- a/src/api/protos/ydb_export.proto +++ b/src/api/protos/ydb_export.proto @@ -247,6 +247,12 @@ message ExportToFsSettings { // - Patterns are matched against the object path relative to the export's source_path. // - Object is excluded from export operation if it matches any of the specified exclude regexps. repeated string exclude_regexps = 8; + + // Include index data or not. + // By default, only index metadata is uploaded and indexes are built during import — it saves space + // and reduces export time, but it can potentially increase the import time. + // Index data can be uploaded and downloaded back during import. + bool include_index_data = 9; } message ExportToFsResult { From 3982cf101c6193ab71edd2c26e1d8771fc1c6a3c Mon Sep 17 00:00:00 2001 From: Maksim Date: Tue, 10 Mar 2026 07:50:47 +0000 Subject: [PATCH 48/93] [NBS-6905] Make nbs2 volume and partition tablets (#34315) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_nbs.proto | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index ce02ac0e478..f2268abd1af 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -72acad87a9603dcf39447997fc5f634904c070ff +df2ba2d668d27c3eb8fb7de2f69afda12ce536da diff --git a/src/api/protos/draft/ydb_nbs.proto b/src/api/protos/draft/ydb_nbs.proto index f323b291b7d..2e120f32908 100644 --- a/src/api/protos/draft/ydb_nbs.proto +++ b/src/api/protos/draft/ydb_nbs.proto @@ -33,6 +33,9 @@ message CreatePartitionRequest { // Storage media type StorageMediaKind StorageMedia = 5; + + // Disk id. + string DiskId = 6; } message CreatePartitionResponse { From 449218c1796ffe0d00e3fb46cc81bac6ea2bd180 Mon Sep 17 00:00:00 2001 From: Ivan Nikolaev Date: Tue, 10 Mar 2026 07:50:54 +0000 Subject: [PATCH 49/93] Support cancel/forget/list forced compaction methods in RPC (#34676) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/library/operation_id/operation_id.h | 1 + src/library/operation_id/operation_id.cpp | 4 ++++ 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index f2268abd1af..c57b3dbf27d 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -df2ba2d668d27c3eb8fb7de2f69afda12ce536da +8b3f8599459202834d0eff07f39655a75c532ba8 diff --git a/include/ydb-cpp-sdk/library/operation_id/operation_id.h b/include/ydb-cpp-sdk/library/operation_id/operation_id.h index 0524ceccd8b..5727a9bf178 100644 --- a/include/ydb-cpp-sdk/library/operation_id/operation_id.h +++ b/include/ydb-cpp-sdk/library/operation_id/operation_id.h @@ -28,6 +28,7 @@ class TOperationId { SS_BG_TASKS = 10, INCREMENTAL_BACKUP = 11, RESTORE = 12, + COMPACTION = 13, }; struct TData { diff --git a/src/library/operation_id/operation_id.cpp b/src/library/operation_id/operation_id.cpp index c4bea6f98f0..b9c0f56d484 100644 --- a/src/library/operation_id/operation_id.cpp +++ b/src/library/operation_id/operation_id.cpp @@ -326,6 +326,10 @@ TOperationId::EKind ParseKind(const std::string_view value) { return TOperationId::RESTORE; } + if (value.starts_with("compaction")) { + return TOperationId::COMPACTION; + } + return TOperationId::UNUSED; } From 3dcfb17a0f2f506803e23c38801200b4e5619a29 Mon Sep 17 00:00:00 2001 From: stanislav_shchetinin Date: Tue, 10 Mar 2026 07:51:01 +0000 Subject: [PATCH 50/93] Support index_population_mode in ImportFs (#34832) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/import/import.h | 1 + src/api/protos/ydb_import.proto | 4 ++++ src/client/import/import.cpp | 3 +++ 4 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index c57b3dbf27d..13fe57257aa 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -8b3f8599459202834d0eff07f39655a75c532ba8 +dbdfe6cfa63a715847209056fa1a67ce9e2dd75e diff --git a/include/ydb-cpp-sdk/client/import/import.h b/include/ydb-cpp-sdk/client/import/import.h index aa88e24b1c2..742627595b6 100644 --- a/include/ydb-cpp-sdk/client/import/import.h +++ b/include/ydb-cpp-sdk/client/import/import.h @@ -169,6 +169,7 @@ struct TImportFromFsSettings : public TOperationRequestSettings Date: Tue, 10 Mar 2026 07:51:08 +0000 Subject: [PATCH 51/93] [NBS2]: introduce batch sync and erase (#34613) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_nbs.proto | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 13fe57257aa..3d32b5be739 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -dbdfe6cfa63a715847209056fa1a67ce9e2dd75e +8caed9c233bdde57dcad3c9b9f3baba83a8dbeee diff --git a/src/api/protos/draft/ydb_nbs.proto b/src/api/protos/draft/ydb_nbs.proto index 2e120f32908..411c6567689 100644 --- a/src/api/protos/draft/ydb_nbs.proto +++ b/src/api/protos/draft/ydb_nbs.proto @@ -36,6 +36,8 @@ message CreatePartitionRequest { // Disk id. string DiskId = 6; + + uint32 SyncRequestsBatchSize = 7; } message CreatePartitionResponse { From 6572e7632ecaaf043514de61a34ad0498a5d3bfd Mon Sep 17 00:00:00 2001 From: Vitaliy Filippov Date: Tue, 10 Mar 2026 07:51:15 +0000 Subject: [PATCH 52/93] Remove Layout (FLAT/FLAT_RELEVANCE) from fulltext settings (#34802) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/table/table.h | 13 +- src/api/protos/ydb_table.proto | 155 +++++++++++------------ src/client/table/out.cpp | 18 +-- src/client/table/table.cpp | 55 ++------ 5 files changed, 94 insertions(+), 149 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 3d32b5be739..673e82ce4a7 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -8caed9c233bdde57dcad3c9b9f3baba83a8dbeee +324cd64ade2da6b14076e5bcdbdd520aa1f11a1c diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index fb103f0d343..ab3466100df 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -313,12 +313,6 @@ struct TKMeansTreeSettings { struct TFulltextIndexSettings { public: - enum class ELayout { - Unspecified = 0, - Flat, - FlatRelevance, - }; - enum class ETokenizer { Unspecified = 0, Whitespace, @@ -345,7 +339,6 @@ struct TFulltextIndexSettings { std::optional Analyzers; }; - std::optional Layout; std::vector Columns; static TFulltextIndexSettings FromProto(const Ydb::Table::FulltextIndexSettings& proto); @@ -820,8 +813,8 @@ class TTableDescription { void AddVectorKMeansTreeIndex(const std::string& indexName, const std::vector& indexColumns, const TKMeansTreeSettings& indexSettings); void AddVectorKMeansTreeIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TKMeansTreeSettings& indexSettings); // fulltext - void AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings); - void AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings); + void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings); + void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings); // default void AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); @@ -1066,6 +1059,8 @@ class TTableBuilder { // fulltext TTableBuilder& AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings); TTableBuilder& AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings); + TTableBuilder& AddFulltextRelevanceIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings); + TTableBuilder& AddFulltextRelevanceIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings); // default TTableBuilder& AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 3237af482d4..9477c6b5856 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -118,86 +118,6 @@ message GlobalVectorKMeansTreeIndex { } message FulltextIndexSettings { - // Specifies the layout strategy for storing and updating the full-text index - enum Layout { - LAYOUT_UNSPECIFIED = 0; - - // Simplest possible layout without support for relevance. - // Uses a single flat inverted index table (indexImplTable). - // Supports a single column only. - // Example source table: - // ┌────┬────────────────────────────┐ - // │ id │ text │ - // ├────┼────────────────────────────┤ - // │ 1 │ "The quick brown fox" │ - // │ 2 │ "The quick blue hare" │ - // └────┴────────────────────────────┘ - // Example inverted index table (indexImplTable): - // ┌──────────────┬────┐ - // │ __ydb_token │ id │ - // ├──────────────┼────┤ - // │ "The" │ 1 │ - // │ "The" │ 2 │ - // │ "blue" │ 2 │ - // │ "brown" │ 1 │ - // │ "fox" │ 1 │ - // │ "hare" │ 2 │ - // │ "quick" │ 1 │ - // │ "quick" │ 2 │ - // └──────────────┴────┘ - FLAT = 1; - - // Simple layout like FLAT, but with support for relevance. - // Uses multiple index tables. Supports a single column only. - // Example source table: - // ┌────┬───────────────────────────────────────┐ - // │ id │ text │ - // ├────┼───────────────────────────────────────┤ - // │ 1 │ "A quick fox catches quick hare" │ - // │ 2 │ "The quick blue hare" │ - // └────┴───────────────────────────────────────┘ - // Example inverted index table with term frequencies (indexImplTable): - // ┌──────────────┬────┬────────────┐ - // │ __ydb_token │ id │ __ydb_freq │ - // ├──────────────┼────┼────────────┤ - // │ "A" │ 1 │ 1 │ - // │ "The" │ 2 │ 1 │ - // │ "blue" │ 2 │ 1 │ - // │ "catches" │ 1 │ 1 │ - // │ "fox" │ 1 │ 1 │ - // │ "hare" │ 1 │ 1 │ - // │ "hare" │ 2 │ 1 │ - // │ "quick" │ 1 │ 2 │ - // │ "quick" │ 2 │ 1 │ - // └──────────────┴────┴────────────┘ - // Term dictionary table (indexImplDictTable): - // ┌──────────────┬────────────┐ - // │ __ydb_token │ __ydb_freq │ - // ├──────────────┼────────────┤ - // │ "A" │ 1 │ - // │ "The" │ 1 │ - // │ "blue" │ 1 │ - // │ "catches" │ 1 │ - // │ "fox" │ 1 │ - // │ "hare" │ 2 │ - // │ "quick" │ 2 │ - // └──────────────┴────────────┘ - // Document statistics table (indexImplDocsTable): - // ┌────┬──────────────┐ - // │ id │ __ydb_length │ - // ├────┼──────────────┤ - // │ 1 │ 6 │ - // │ 2 │ 4 │ - // └────┴──────────────┘ - // Global statistics table (indexImplStatsTable): - // ┌──────────┬─────────────────┬────────────────────┐ - // │ __ydb_id │ __ydb_doc_count │ __ydb_total_length │ - // ├──────────┼─────────────────┼────────────────────┤ - // │ 1 │ 2 │ 10 │ - // └──────────┴─────────────────┴────────────────────┘ - FLAT_RELEVANCE = 2; - } - // Specifies how text is tokenized during indexing enum Tokenizer { TOKENIZER_UNSPECIFIED = 0; @@ -302,8 +222,8 @@ message FulltextIndexSettings { Analyzers analyzers = 2; } - // See Layout enum - optional Layout layout = 1; + // Removed Layout field + reserved 1; // List of columns and their fulltext settings // Currently, this list should contain a single entry with specified analyzers @@ -312,11 +232,82 @@ message FulltextIndexSettings { repeated ColumnAnalyzers columns = 2; } +// Simplest possible layout without support for relevance. +// Uses a single flat inverted index table (indexImplTable). +// Supports a single column only. +// Example source table: +// ┌────┬────────────────────────────┐ +// │ id │ text │ +// ├────┼────────────────────────────┤ +// │ 1 │ "The quick brown fox" │ +// │ 2 │ "The quick blue hare" │ +// └────┴────────────────────────────┘ +// Example inverted index table (indexImplTable): +// ┌──────────────┬────┐ +// │ __ydb_token │ id │ +// ├──────────────┼────┤ +// │ "The" │ 1 │ +// │ "The" │ 2 │ +// │ "blue" │ 2 │ +// │ "brown" │ 1 │ +// │ "fox" │ 1 │ +// │ "hare" │ 2 │ +// │ "quick" │ 1 │ +// │ "quick" │ 2 │ +// └──────────────┴────┘ message GlobalFulltextPlainIndex { GlobalIndexSettings settings = 1; FulltextIndexSettings fulltext_settings = 2; } +// Simple layout like Plain, but with support for relevance. +// Uses multiple index tables. Supports a single column only. +// Example source table: +// ┌────┬───────────────────────────────────────┐ +// │ id │ text │ +// ├────┼───────────────────────────────────────┤ +// │ 1 │ "A quick fox catches quick hare" │ +// │ 2 │ "The quick blue hare" │ +// └────┴───────────────────────────────────────┘ +// Example inverted index table with term frequencies (indexImplTable): +// ┌──────────────┬────┬────────────┐ +// │ __ydb_token │ id │ __ydb_freq │ +// ├──────────────┼────┼────────────┤ +// │ "A" │ 1 │ 1 │ +// │ "The" │ 2 │ 1 │ +// │ "blue" │ 2 │ 1 │ +// │ "catches" │ 1 │ 1 │ +// │ "fox" │ 1 │ 1 │ +// │ "hare" │ 1 │ 1 │ +// │ "hare" │ 2 │ 1 │ +// │ "quick" │ 1 │ 2 │ +// │ "quick" │ 2 │ 1 │ +// └──────────────┴────┴────────────┘ +// Term dictionary table (indexImplDictTable): +// ┌──────────────┬────────────┐ +// │ __ydb_token │ __ydb_freq │ +// ├──────────────┼────────────┤ +// │ "A" │ 1 │ +// │ "The" │ 1 │ +// │ "blue" │ 1 │ +// │ "catches" │ 1 │ +// │ "fox" │ 1 │ +// │ "hare" │ 2 │ +// │ "quick" │ 2 │ +// └──────────────┴────────────┘ +// Document statistics table (indexImplDocsTable): +// ┌────┬──────────────┐ +// │ id │ __ydb_length │ +// ├────┼──────────────┤ +// │ 1 │ 6 │ +// │ 2 │ 4 │ +// └────┴──────────────┘ +// Global statistics table (indexImplStatsTable): +// ┌──────────┬─────────────────┬────────────────────┐ +// │ __ydb_id │ __ydb_doc_count │ __ydb_total_length │ +// ├──────────┼─────────────────┼────────────────────┤ +// │ 1 │ 2 │ 10 │ +// └──────────┴─────────────────┴────────────────────┘ message GlobalFulltextRelevanceIndex { GlobalIndexSettings posting_table_settings = 1; FulltextIndexSettings fulltext_settings = 2; diff --git a/src/client/table/out.cpp b/src/client/table/out.cpp index 7e90331c995..7d29674d4f2 100644 --- a/src/client/table/out.cpp +++ b/src/client/table/out.cpp @@ -86,20 +86,6 @@ Y_DECLARE_OUT_SPEC(, NYdb::NTable::TKMeansTreeSettings, stream, value) { " }"; } -Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::ELayout, stream, value) { - switch (value) { - case NYdb::NTable::TFulltextIndexSettings::ELayout::Flat: - stream << "flat"; - break; - case NYdb::NTable::TFulltextIndexSettings::ELayout::FlatRelevance: - stream << "flat_relevance"; - break; - case NYdb::NTable::TFulltextIndexSettings::ELayout::Unspecified: - stream << "unspecified"; - break; - } -} - Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::ETokenizer, stream, value) { switch (value) { case NYdb::NTable::TFulltextIndexSettings::ETokenizer::Whitespace: @@ -164,9 +150,9 @@ Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::TColumnAnalyzers, str } Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings, stream, value) { - stream << "{ layout: " << value.Layout.value_or(NYdb::NTable::TFulltextIndexSettings::ELayout::Unspecified); + stream << "{"; if (!value.Columns.empty()) { - stream << ", columns: ["; + stream << " columns: ["; for (size_t i = 0; i < value.Columns.size(); ++i) { if (i > 0) stream << ", "; stream << value.Columns[i]; diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 6e9f1421157..63bc32c6ce8 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -792,19 +792,11 @@ void TTableDescription::AddVectorKMeansTreeIndex(const std::string& indexName, c Impl_->AddVectorKMeansTreeIndex(indexName, EIndexType::GlobalVectorKMeansTree, indexColumns, dataColumns, indexSettings); } -void TTableDescription::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { - EIndexType indexType = EIndexType::GlobalFulltextPlain; - if (indexSettings.Layout.has_value() && indexSettings.Layout.value() == TFulltextIndexSettings::ELayout::FlatRelevance) { - indexType = EIndexType::GlobalFulltextRelevance; - } +void TTableDescription::AddFulltextIndex(const std::string& indexName, EIndexType indexType, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { Impl_->AddFulltextIndex(indexName, indexType, indexColumns, indexSettings); } -void TTableDescription::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { - EIndexType indexType = EIndexType::GlobalFulltextPlain; - if (indexSettings.Layout.has_value() && indexSettings.Layout.value() == TFulltextIndexSettings::ELayout::FlatRelevance) { - indexType = EIndexType::GlobalFulltextRelevance; - } +void TTableDescription::AddFulltextIndex(const std::string& indexName, EIndexType indexType, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { Impl_->AddFulltextIndex(indexName, indexType, indexColumns, dataColumns, indexSettings); } @@ -1298,12 +1290,22 @@ TTableBuilder& TTableBuilder::AddVectorKMeansTreeIndex(const std::string& indexN } TTableBuilder& TTableBuilder::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { - TableDescription_.AddFulltextIndex(indexName, indexColumns, dataColumns, indexSettings); + TableDescription_.AddFulltextIndex(indexName, EIndexType::GlobalFulltextPlain, indexColumns, dataColumns, indexSettings); return *this; } TTableBuilder& TTableBuilder::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { - TableDescription_.AddFulltextIndex(indexName, indexColumns, indexSettings); + TableDescription_.AddFulltextIndex(indexName, EIndexType::GlobalFulltextPlain, indexColumns, indexSettings); + return *this; +} + +TTableBuilder& TTableBuilder::AddFulltextRelevanceIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { + TableDescription_.AddFulltextIndex(indexName, EIndexType::GlobalFulltextRelevance, indexColumns, dataColumns, indexSettings); + return *this; +} + +TTableBuilder& TTableBuilder::AddFulltextRelevanceIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { + TableDescription_.AddFulltextIndex(indexName, EIndexType::GlobalFulltextRelevance, indexColumns, indexSettings); return *this; } @@ -2715,43 +2717,14 @@ Ydb::Table::FulltextIndexSettings::ColumnAnalyzers ToProto(const TFulltextIndexS } TFulltextIndexSettings TFulltextIndexSettings::FromProto(const Ydb::Table::FulltextIndexSettings& proto) { - auto convertLayout = [&] { - switch (proto.layout()) { - case Ydb::Table::FulltextIndexSettings::FLAT: - return ELayout::Flat; - case Ydb::Table::FulltextIndexSettings::FLAT_RELEVANCE: - return ELayout::FlatRelevance; - default: - return ELayout::Unspecified; - } - }; - TFulltextIndexSettings result; - result.Layout = convertLayout(); for (const auto& columnProto : proto.columns()) { result.Columns.push_back(NTable::FromProto(columnProto)); } - return result; } void TFulltextIndexSettings::SerializeTo(Ydb::Table::FulltextIndexSettings& settings) const { - auto convertLayout = [&] { - switch (*Layout) { - case ELayout::Flat: - return Ydb::Table::FulltextIndexSettings::FLAT; - case ELayout::FlatRelevance: - return Ydb::Table::FulltextIndexSettings::FLAT_RELEVANCE; - case ELayout::Unspecified: - return Ydb::Table::FulltextIndexSettings::LAYOUT_UNSPECIFIED; - } - return Ydb::Table::FulltextIndexSettings::LAYOUT_UNSPECIFIED; - }; - - if (Layout.has_value()) { - settings.set_layout(convertLayout()); - } - for (const auto& column : Columns) { *settings.add_columns() = ToProto(column); } From 50943b0a74f554f2a410fee8539c0690d5b8977f Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Tue, 10 Mar 2026 07:51:22 +0000 Subject: [PATCH 53/93] YQ-5091 rename PQ partitions balancer and fix registration (#34667) --- .github/last_commit.txt | 2 +- .../ut/ut_utils/topic_sdk_test_setup.cpp | 38 +------------------ 2 files changed, 3 insertions(+), 37 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 673e82ce4a7..aa32890972c 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -324cd64ade2da6b14076e5bcdbdd520aa1f11a1c +f72db541a5e9b41206574f528106e2a8d0ea1d1a diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp index a402943c798..6f16928a608 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp @@ -1,51 +1,17 @@ #include "topic_sdk_test_setup.h" #include +#include using namespace std::chrono_literals; namespace NYdb::inline V3::NTopic::NTests { -namespace { - -void TerminateHandler() { - Cerr << "======= terminate() call stack ========" << Endl; - FormatBackTrace(&Cerr); - if (const auto& backtrace = TBackTrace::FromCurrentException(); backtrace.size() > 0) { - Cerr << "======== exception call stack =========" << Endl; - backtrace.PrintTo(Cerr); - } - Cerr << "=======================================" << Endl; - - if (std::current_exception()) { - Cerr << "Uncaught exception: " << CurrentExceptionMessage() << Endl; - } else { - Cerr << "Terminate for unknown reason (no current exception)" << Endl; - } - - abort(); -} - -void BackTraceSignalHandler(int signal) { - Cerr << "======= Signal " << signal << " call stack ========" << Endl; - FormatBackTrace(&Cerr); - Cerr << "===============================================" << Endl; - - abort(); -} - -} // anonymous namespace - TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NKikimr::Tests::TServerSettings& settings, bool createTopic) : Database_("/Root") , Server_(settings, false) { - NKikimr::EnableYDBBacktraceFormat(); - - std::set_terminate(&TerminateHandler); - for (auto sig : {SIGFPE, SIGILL, SIGSEGV}) { - signal(sig, &BackTraceSignalHandler); - } + NTestUtils::SetupSignalHandlers(); Log_.SetFormatter([testCaseName](ELogPriority priority, TStringBuf message) { return TStringBuilder() << TInstant::Now() << " :" << testCaseName << " " << priority << ": " << message << Endl; From b75187a50c08accab23a3aa52906ea6bd12c44d8 Mon Sep 17 00:00:00 2001 From: Maxim Yurchuk Date: Tue, 10 Mar 2026 07:51:28 +0000 Subject: [PATCH 54/93] Revert "YQ-5091 rename PQ partitions balancer and fix registration" (#34982) --- .github/last_commit.txt | 2 +- .../ut/ut_utils/topic_sdk_test_setup.cpp | 38 ++++++++++++++++++- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index aa32890972c..0d4594327ff 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -f72db541a5e9b41206574f528106e2a8d0ea1d1a +c461e6aaf71af5e48653282f9c8407d331d21dc9 diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp index 6f16928a608..a402943c798 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp @@ -1,17 +1,51 @@ #include "topic_sdk_test_setup.h" #include -#include using namespace std::chrono_literals; namespace NYdb::inline V3::NTopic::NTests { +namespace { + +void TerminateHandler() { + Cerr << "======= terminate() call stack ========" << Endl; + FormatBackTrace(&Cerr); + if (const auto& backtrace = TBackTrace::FromCurrentException(); backtrace.size() > 0) { + Cerr << "======== exception call stack =========" << Endl; + backtrace.PrintTo(Cerr); + } + Cerr << "=======================================" << Endl; + + if (std::current_exception()) { + Cerr << "Uncaught exception: " << CurrentExceptionMessage() << Endl; + } else { + Cerr << "Terminate for unknown reason (no current exception)" << Endl; + } + + abort(); +} + +void BackTraceSignalHandler(int signal) { + Cerr << "======= Signal " << signal << " call stack ========" << Endl; + FormatBackTrace(&Cerr); + Cerr << "===============================================" << Endl; + + abort(); +} + +} // anonymous namespace + TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NKikimr::Tests::TServerSettings& settings, bool createTopic) : Database_("/Root") , Server_(settings, false) { - NTestUtils::SetupSignalHandlers(); + NKikimr::EnableYDBBacktraceFormat(); + + std::set_terminate(&TerminateHandler); + for (auto sig : {SIGFPE, SIGILL, SIGSEGV}) { + signal(sig, &BackTraceSignalHandler); + } Log_.SetFormatter([testCaseName](ELogPriority priority, TStringBuf message) { return TStringBuilder() << TInstant::Now() << " :" << testCaseName << " " << priority << ": " << message << Endl; From ba8cddd57b1a5b884cbe3358e708b024b2d86a88 Mon Sep 17 00:00:00 2001 From: Vitaliy Filippov Date: Tue, 10 Mar 2026 07:51:35 +0000 Subject: [PATCH 55/93] Add a note about removed full text Layout field (#34998) --- .github/last_commit.txt | 2 +- CHANGELOG.md | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 0d4594327ff..9dd094f5322 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -c461e6aaf71af5e48653282f9c8407d331d21dc9 +a03ac2c3c300df29715b7f638f7bc67326e26b5f diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b137891791..1d624f72a4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1 +1,3 @@ - +* Removed the `layout` field from `FulltextIndexSettings` and replaced it with separate index types in `TableIndexDescription`. + `layout` was a preliminary version of API, actual YDB release 26-1 uses separate index types, so please note that creating + full text indexes via gRPC won't work with previous versions of SDK. From 65a6e7e78c5b2e1970db657a729692a78824e189 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Tue, 10 Mar 2026 07:51:42 +0000 Subject: [PATCH 56/93] Revert "Revert "YQ-5091 rename PQ partitions balancer and fix registration"" (#34989) --- .github/last_commit.txt | 2 +- .../ut/ut_utils/topic_sdk_test_setup.cpp | 38 +------------------ 2 files changed, 3 insertions(+), 37 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 9dd094f5322..8b091c524bb 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -a03ac2c3c300df29715b7f638f7bc67326e26b5f +9fefd18aab36174963665dc3b75df460995a09c1 diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp index a402943c798..6f16928a608 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp @@ -1,51 +1,17 @@ #include "topic_sdk_test_setup.h" #include +#include using namespace std::chrono_literals; namespace NYdb::inline V3::NTopic::NTests { -namespace { - -void TerminateHandler() { - Cerr << "======= terminate() call stack ========" << Endl; - FormatBackTrace(&Cerr); - if (const auto& backtrace = TBackTrace::FromCurrentException(); backtrace.size() > 0) { - Cerr << "======== exception call stack =========" << Endl; - backtrace.PrintTo(Cerr); - } - Cerr << "=======================================" << Endl; - - if (std::current_exception()) { - Cerr << "Uncaught exception: " << CurrentExceptionMessage() << Endl; - } else { - Cerr << "Terminate for unknown reason (no current exception)" << Endl; - } - - abort(); -} - -void BackTraceSignalHandler(int signal) { - Cerr << "======= Signal " << signal << " call stack ========" << Endl; - FormatBackTrace(&Cerr); - Cerr << "===============================================" << Endl; - - abort(); -} - -} // anonymous namespace - TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NKikimr::Tests::TServerSettings& settings, bool createTopic) : Database_("/Root") , Server_(settings, false) { - NKikimr::EnableYDBBacktraceFormat(); - - std::set_terminate(&TerminateHandler); - for (auto sig : {SIGFPE, SIGILL, SIGSEGV}) { - signal(sig, &BackTraceSignalHandler); - } + NTestUtils::SetupSignalHandlers(); Log_.SetFormatter([testCaseName](ELogPriority priority, TStringBuf message) { return TStringBuilder() << TInstant::Now() << " :" << testCaseName << " " << priority << ": " << message << Endl; From 66410622fab5bb02f001d3b920f323e277a08f6b Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Tue, 10 Mar 2026 07:51:49 +0000 Subject: [PATCH 57/93] Fixed TSAN error in tests (#35091) --- .github/last_commit.txt | 2 +- .../ut/ut_utils/topic_sdk_test_setup.cpp | 43 ++++++++++++++----- .../topic/ut/ut_utils/topic_sdk_test_setup.h | 7 +++ 3 files changed, 40 insertions(+), 12 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 8b091c524bb..8bbda038e9d 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -9fefd18aab36174963665dc3b75df460995a09c1 +06f0aa0e59f6769b5719ba22f26ea54efb89d625 diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp index 6f16928a608..71608ae13ae 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp @@ -18,6 +18,7 @@ TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NK }); Server_.StartServer(true, GetDatabase()); + Driver = std::make_unique(MakeDriverConfig()); Log_ << "TTopicSdkTestSetup started"; @@ -27,6 +28,13 @@ TTopicSdkTestSetup::TTopicSdkTestSetup(const std::string& testCaseName, const NK } } +TTopicSdkTestSetup::~TTopicSdkTestSetup() { + if (Driver) { + Driver->Stop(true); + Driver = nullptr; + } +} + void TTopicSdkTestSetup::CreateTopicWithAutoscale(const std::string& name, const std::string& consumer, std::size_t partitionCount, @@ -45,7 +53,7 @@ void TTopicSdkTestSetup::CreateTopic(const std::string& name, const std::string& TConsumerDescription TTopicSdkTestSetup::DescribeConsumer(const std::string& name, const std::string& consumer) { - TTopicClient client(MakeDriver()); + TTopicClient client(*Driver); TDescribeConsumerSettings settings; settings.IncludeStats(true); @@ -54,7 +62,7 @@ TConsumerDescription TTopicSdkTestSetup::DescribeConsumer(const std::string& nam auto status = client.DescribeConsumer(GetTopicPath(name), GetConsumerName(consumer), settings).GetValueSync(); UNIT_ASSERT(status.IsSuccess()); - return status.GetConsumerDescription(); + return std::move(status.GetConsumerDescription()); } void TTopicSdkTestSetup::Write(const std::string& message, std::uint32_t partitionId, @@ -65,8 +73,8 @@ void TTopicSdkTestSetup::Write(const std::string& message, std::uint32_t partiti void TTopicSdkTestSetup::Write(const std::string& topic, const std::string& message, std::uint32_t partitionId, const std::optional producer, - std::optional seqNo) { - TTopicClient client(MakeDriver()); + std::optional seqNo) { + TTopicClient client(*Driver); TWriteSessionSettings settings; settings.Path(topic); @@ -80,13 +88,27 @@ void TTopicSdkTestSetup::Write(const std::string& topic, const std::string& mess UNIT_ASSERT(session->Write(message, seqNo)); - session->Close(TDuration::Seconds(5)); + session->Close(); +} + +TTopicSdkTestSetup::TReadResult::TReadResult(TDriver& driver) + : Client(driver) +{ +} + +TTopicSdkTestSetup::TReadResult::~TReadResult() +{ + if (Reader) { + Reader->Close(); + Reader.reset(); + } } TTopicSdkTestSetup::TReadResult TTopicSdkTestSetup::Read(const std::string& topic, const std::string& consumer, std::function handler, std::optional partition, const TDuration timeout, bool autoPartitioningSupport) { - TTopicClient client(MakeDriver()); + + TReadResult result(*Driver); auto topicSettings = TTopicReadSettings(topic); if (partition) { @@ -98,11 +120,10 @@ TTopicSdkTestSetup::TReadResult TTopicSdkTestSetup::Read(const std::string& topi .AppendTopics(topicSettings) .ConsumerName(consumer); - auto reader = client.CreateReadSession(settings); + auto reader = result.Client.CreateReadSession(settings); TInstant deadlineTime = TInstant::Now() + timeout; - TReadResult result; result.Reader = reader; bool continueFlag = true; @@ -146,7 +167,7 @@ TTopicSdkTestSetup::TReadResult TTopicSdkTestSetup::Read(const std::string& topi } TStatus TTopicSdkTestSetup::Commit(const std::string& path, const std::string& consumerName, size_t partitionId, size_t offset, std::optional sessionId) { - TTopicClient client(MakeDriver()); + TTopicClient client(*Driver); TCommitOffsetSettings commitSettings {.ReadSessionId_ = sessionId}; return client.CommitOffset(path, partitionId, consumerName, offset, commitSettings).GetValueSync(); @@ -224,12 +245,12 @@ NKikimr::Tests::TServerSettings TTopicSdkTestSetup::MakeServerSettings() TTopicClient TTopicSdkTestSetup::MakeClient() const { - return TTopicClient(MakeDriver()); + return TTopicClient(*Driver); } NYdb::NTable::TTableClient TTopicSdkTestSetup::MakeTableClient() const { - return NYdb::NTable::TTableClient(MakeDriver(), NYdb::NTable::TClientSettings() + return NYdb::NTable::TTableClient(*Driver, NYdb::NTable::TClientSettings() .UseQueryCache(false)); } diff --git a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.h b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.h index e3fe5d48a0b..e4381554e10 100644 --- a/src/client/topic/ut/ut_utils/topic_sdk_test_setup.h +++ b/src/client/topic/ut/ut_utils/topic_sdk_test_setup.h @@ -13,6 +13,8 @@ namespace NYdb::inline V3::NTopic::NTests { class TTopicSdkTestSetup : public ITopicTestSetup { public: explicit TTopicSdkTestSetup(const std::string& testCaseName, const NKikimr::Tests::TServerSettings& settings = MakeServerSettings(), bool createTopic = true); + TTopicSdkTestSetup(TTopicSdkTestSetup&& other) noexcept = default; + ~TTopicSdkTestSetup(); void CreateTopic(const std::string& name = TEST_TOPIC, const std::string& consumer = TEST_CONSUMER, @@ -37,6 +39,10 @@ class TTopicSdkTestSetup : public ITopicTestSetup { std::optional seqNo = std::nullopt); struct TReadResult { + TReadResult(TDriver& driver); + ~TReadResult(); + + TTopicClient Client; std::shared_ptr Reader; bool Timeout; @@ -76,6 +82,7 @@ class TTopicSdkTestSetup : public ITopicTestSetup { std::string Database_; ::NPersQueue::TTestServer Server_; + std::unique_ptr Driver; TLog Log_ = CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG); }; From 7abd00a0d6d67d914e34d48087cb3038ff3c07bb Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:51:56 +0000 Subject: [PATCH 58/93] [C++ SDK] Added grpc load balancing policy option (#35137) --- .github/last_commit.txt | 2 +- include/ydb-cpp-sdk/client/driver/driver.h | 6 ++++++ src/client/driver/driver.cpp | 8 ++++++++ .../impl/internal/grpc_connections/grpc_connections.cpp | 1 + .../impl/internal/grpc_connections/grpc_connections.h | 3 ++- src/client/impl/internal/grpc_connections/params.h | 1 + 6 files changed, 19 insertions(+), 2 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 8bbda038e9d..272d96af807 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -06f0aa0e59f6769b5719ba22f26ea54efb89d625 +fc76da9b9b22ab7364236471559de2e2cbcdacf4 diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index b5a4df7faf2..199459c58f7 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -124,6 +124,7 @@ class TDriverConfig { //! Params is a optionally field to set policy settings //! default: EBalancingPolicy::UsePreferableLocation TDriverConfig& SetBalancingPolicy(EBalancingPolicy policy, const std::string& params = ""); + //! Set grpc level keep alive. If keepalive ping was delayed more than given timeout //! internal grpc routine fails request with TRANSIENT_FAILURE or TRANSPORT_UNAVAILABLE error //! Note: this timeout should not be too small to prevent fail due to @@ -133,6 +134,11 @@ class TDriverConfig { TDriverConfig& SetGRpcKeepAliveTimeout(TDuration timeout); TDriverConfig& SetGRpcKeepAlivePermitWithoutCalls(bool permitWithoutCalls); + //! Set grpc load balancing policy + //! policy - name of the load balancing policy, see grpc documentation for available policies + //! default: "round_robin" + TDriverConfig& SetGRpcLoadBalancingPolicy(const std::string& policy); + //! Set inactive socket timeout. //! Used to close connections, that were inactive for given time. //! Closes unused connections every 1/10 of timeout, so deletion time is approximate. diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 18ac064e480..9bfc5f80560 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -45,6 +45,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { TBalancingPolicy::TImpl GetBalancingSettings() const override { return BalancingSettings; } TDuration GetGRpcKeepAliveTimeout() const override { return GRpcKeepAliveTimeout; } bool GetGRpcKeepAlivePermitWithoutCalls() const override { return GRpcKeepAlivePermitWithoutCalls; } + std::string GetGRpcLoadBalancingPolicy() const override { return GRpcLoadBalancingPolicy; } TDuration GetSocketIdleTimeout() const override { return SocketIdleTimeout; } uint64_t GetMemoryQuota() const override { return MemoryQuota; } uint64_t GetMaxInboundMessageSize() const override { return MaxInboundMessageSize; } @@ -77,6 +78,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { TBalancingPolicy::TImpl BalancingSettings = TBalancingPolicy::TImpl::UsePreferableLocation(std::nullopt); TDuration GRpcKeepAliveTimeout = TDuration::Seconds(10); bool GRpcKeepAlivePermitWithoutCalls = true; + std::string GRpcLoadBalancingPolicy = "round_robin"; TDuration SocketIdleTimeout = TDuration::Minutes(6); uint64_t MemoryQuota = 0; uint64_t MaxInboundMessageSize = 0; @@ -210,6 +212,11 @@ TDriverConfig& TDriverConfig::SetGRpcKeepAlivePermitWithoutCalls(bool permitWith return *this; } +TDriverConfig& TDriverConfig::SetGRpcLoadBalancingPolicy(const std::string& policy) { + Impl_->GRpcLoadBalancingPolicy = policy; + return *this; +} + TDriverConfig& TDriverConfig::SetSocketIdleTimeout(TDuration timeout) { Impl_->SocketIdleTimeout = timeout; return *this; @@ -297,6 +304,7 @@ TDriverConfig TDriver::GetConfig() const { config.SetBalancingPolicy(std::make_unique(Impl_->BalancingSettings_)); config.SetGRpcKeepAliveTimeout(std::chrono::duration_cast(Impl_->GRpcKeepAliveTimeout_)); config.SetGRpcKeepAlivePermitWithoutCalls(Impl_->GRpcKeepAlivePermitWithoutCalls_); + config.SetGRpcLoadBalancingPolicy(Impl_->GRpcLoadBalancingPolicy_); config.SetSocketIdleTimeout(std::chrono::duration_cast(Impl_->SocketIdleTimeout_)); config.SetMaxInboundMessageSize(Impl_->MaxInboundMessageSize_); config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index d17cc078195..625358f0023 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -157,6 +157,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , BalancingSettings_(params->GetBalancingSettings()) , GRpcKeepAliveTimeout_(TDeadline::SafeDurationCast(params->GetGRpcKeepAliveTimeout())) , GRpcKeepAlivePermitWithoutCalls_(params->GetGRpcKeepAlivePermitWithoutCalls()) + , GRpcLoadBalancingPolicy_(params->GetGRpcLoadBalancingPolicy()) , MemoryQuota_(params->GetMemoryQuota()) , MaxInboundMessageSize_(params->GetMaxInboundMessageSize()) , MaxOutboundMessageSize_(params->GetMaxOutboundMessageSize()) diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 945cb97bcd4..99555ffc24c 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -105,7 +105,7 @@ class TGRpcConnectionsImpl clientConfig.MaxOutboundMessageSize = MaxOutboundMessageSize_; } - clientConfig.LoadBalancingPolicy = "round_robin"; + clientConfig.LoadBalancingPolicy = GRpcLoadBalancingPolicy_; if (dbState->DiscoveryMode != EDiscoveryMode::Off) { if (std::is_same() @@ -706,6 +706,7 @@ class TGRpcConnectionsImpl const TBalancingPolicy::TImpl BalancingSettings_; const TDeadline::Duration GRpcKeepAliveTimeout_; const bool GRpcKeepAlivePermitWithoutCalls_; + const std::string GRpcLoadBalancingPolicy_; const std::uint64_t MemoryQuota_; const std::uint64_t MaxInboundMessageSize_; const std::uint64_t MaxOutboundMessageSize_; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 35c0186421d..7f23bcceffb 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -35,6 +35,7 @@ class IConnectionsParams { virtual TBalancingPolicy::TImpl GetBalancingSettings() const = 0; virtual TDuration GetGRpcKeepAliveTimeout() const = 0; virtual bool GetGRpcKeepAlivePermitWithoutCalls() const = 0; + virtual std::string GetGRpcLoadBalancingPolicy() const = 0; virtual TDuration GetSocketIdleTimeout() const = 0; virtual const TLog& GetLog() const = 0; virtual uint64_t GetMemoryQuota() const = 0; From 53a6fecc9ed65b4424ff7a6af26276112e464593 Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:52:03 +0000 Subject: [PATCH 59/93] Added changelog entry for grpc load balancing policy (#35150) --- .github/last_commit.txt | 2 +- CHANGELOG.md | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 272d96af807..a5f6b65e109 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -fc76da9b9b22ab7364236471559de2e2cbcdacf4 +ef5dbcea3bbe1f36b2d4bda12261d2bd0e5b80a9 diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d624f72a4d..abc1c1de6c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +* Added gRPC load balancing policy option for `TDriver`. Default policy: `round_robin`. + * Removed the `layout` field from `FulltextIndexSettings` and replaced it with separate index types in `TableIndexDescription`. `layout` was a preliminary version of API, actual YDB release 26-1 uses separate index types, so please note that creating full text indexes via gRPC won't work with previous versions of SDK. From fdaf0ef5b5d0d5873c2505b9bcdfc683e2b5adc9 Mon Sep 17 00:00:00 2001 From: Bulat Date: Tue, 10 Mar 2026 07:52:10 +0000 Subject: [PATCH 60/93] [C++ SDK] Moved private executors to adapters (#35197) --- .github/last_commit.txt | 2 +- examples/executor/main.cpp | 11 ++--------- .../ydb-cpp-sdk/client/types/executor/executor.h | 14 -------------- src/client/persqueue_public/ut/read_session_ut.cpp | 3 ++- src/client/types/executor/executor.cpp | 10 ---------- 5 files changed, 5 insertions(+), 35 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index a5f6b65e109..079ce64de38 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -ef5dbcea3bbe1f36b2d4bda12261d2bd0e5b80a9 +5a0805bde3b091c5e0c37075049be2d4ba351e26 diff --git a/examples/executor/main.cpp b/examples/executor/main.cpp index 2d314403fe8..5be6cb85cb7 100644 --- a/examples/executor/main.cpp +++ b/examples/executor/main.cpp @@ -4,20 +4,13 @@ #include -#include - #include void ExecutorExample(const std::string& endpoint, const std::string& database) { auto driverConfig = NYdb::CreateFromEnvironment(endpoint + "/?database=" + database) - .SetExecutor(NYdb::CreateThreadPoolExecutorAdapter( - std::make_shared(TThreadPool::TParams() - .SetBlocking(true) - .SetCatching(false) - .SetForkAware(false)), - std::thread::hardware_concurrency()) - ); + .SetExecutor(NYdb::CreateThreadPoolExecutor(std::thread::hardware_concurrency(), 0) + ); NYdb::TDriver driver(driverConfig); NYdb::NQuery::TQueryClient client(driver); diff --git a/include/ydb-cpp-sdk/client/types/executor/executor.h b/include/ydb-cpp-sdk/client/types/executor/executor.h index 977680869c6..f74d52afaba 100644 --- a/include/ydb-cpp-sdk/client/types/executor/executor.h +++ b/include/ydb-cpp-sdk/client/types/executor/executor.h @@ -1,9 +1,5 @@ #pragma once -#ifndef YDB_SDK_OSS -#include -#endif - #include #include #include @@ -45,14 +41,4 @@ class IExecutor { // Create default executor for thread pool. IExecutor::TPtr CreateThreadPoolExecutor(std::size_t threadCount, std::size_t maxQueueSize = 0); -#ifndef YDB_SDK_OSS -// Create executor adapter for util thread pool. -// Thread pool is started and stopped by SDK. -IExecutor::TPtr CreateThreadPoolExecutorAdapter(std::shared_ptr threadPool, std::size_t threadCount, std::size_t maxQueueSize = 0); - -// Create executor adapter for util thread pool. -// Thread pool is expected to have been started and stopped by user. -IExecutor::TPtr CreateExternalThreadPoolExecutorAdapter(std::shared_ptr threadPool); -#endif - } // namespace NYdb diff --git a/src/client/persqueue_public/ut/read_session_ut.cpp b/src/client/persqueue_public/ut/read_session_ut.cpp index 071ca522a9d..63f60b3909b 100644 --- a/src/client/persqueue_public/ut/read_session_ut.cpp +++ b/src/client/persqueue_public/ut/read_session_ut.cpp @@ -4,6 +4,7 @@ #include #undef INCLUDE_YDB_INTERNAL_H +#include #include #include @@ -599,7 +600,7 @@ ::IExecutor::TPtr TReadSessionImplTestSetup::GetDefaultExecutor() { if (!DefaultExecutor) { ThreadPool = std::make_shared(); ThreadPool->Start(1); - DefaultExecutor = CreateExternalThreadPoolExecutorAdapter(ThreadPool); + DefaultExecutor = NAdapters::CreateExternalThreadPoolExecutorAdapter(ThreadPool); } return DefaultExecutor; } diff --git a/src/client/types/executor/executor.cpp b/src/client/types/executor/executor.cpp index a5e0f8f85cc..75f0d85f0a7 100644 --- a/src/client/types/executor/executor.cpp +++ b/src/client/types/executor/executor.cpp @@ -12,14 +12,4 @@ std::shared_ptr CreateThreadPoolExecutor(std::size_t threadCount, std return std::make_shared(CreateThreadPool(threadCount), threadCount, maxQueueSize); } -#ifndef YDB_SDK_OSS -std::shared_ptr CreateThreadPoolExecutorAdapter(std::shared_ptr threadPool, std::size_t threadCount, std::size_t maxQueueSize) { - return std::make_shared(threadPool, threadCount, maxQueueSize); -} - -std::shared_ptr CreateExternalThreadPoolExecutorAdapter(std::shared_ptr threadPool) { - return std::make_shared(threadPool); -} -#endif - } From 8bae80607a7c09889a89ee91b69995b4e31d2f14 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Tue, 10 Mar 2026 07:52:17 +0000 Subject: [PATCH 61/93] LOGBROKER-10206 New interface (#34369) --- .github/last_commit.txt | 2 +- CHANGELOG.md | 3 + .../producer/basic_write/main.cpp | 123 ++ include/ydb-cpp-sdk/client/topic/client.h | 12 +- include/ydb-cpp-sdk/client/topic/producer.h | 202 ++ .../ydb-cpp-sdk/client/topic/write_session.h | 203 +- src/client/topic/impl/producer.cpp | 1882 +++++++++++++++++ src/client/topic/impl/producer.h | 449 ++++ src/client/topic/impl/topic.cpp | 10 +- src/client/topic/impl/topic_impl.cpp | 35 +- src/client/topic/impl/topic_impl.h | 11 +- src/client/topic/impl/write_session.cpp | 1818 ---------------- src/client/topic/impl/write_session.h | 452 ---- src/client/topic/ut/basic_usage_ut.cpp | 1015 ++++----- src/client/topic/ut/ut_utils/event_loop.cpp | 93 - src/client/topic/ut/ut_utils/event_loop.h | 37 - tests/integration/topic/basic_usage_it.cpp | 100 +- 17 files changed, 3341 insertions(+), 3106 deletions(-) create mode 100644 examples/topic_writer/producer/basic_write/main.cpp create mode 100644 include/ydb-cpp-sdk/client/topic/producer.h create mode 100644 src/client/topic/impl/producer.cpp create mode 100644 src/client/topic/impl/producer.h delete mode 100644 src/client/topic/ut/ut_utils/event_loop.cpp delete mode 100644 src/client/topic/ut/ut_utils/event_loop.h diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 079ce64de38..f0e7add639d 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -5a0805bde3b091c5e0c37075049be2d4ba351e26 +395280d887bc0d1c34863836ee9136d8ffdfa81d diff --git a/CHANGELOG.md b/CHANGELOG.md index abc1c1de6c0..599b3e6c047 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,6 @@ +* EXPERIMENTAL! Added `IProducer` interface to the SDK. This interface is used to write messages to a topic. +Each message can be associated with a partitioning key, which is used to determine the partition to which the message will be written. + * Added gRPC load balancing policy option for `TDriver`. Default policy: `round_robin`. * Removed the `layout` field from `FulltextIndexSettings` and replaced it with separate index types in `TableIndexDescription`. diff --git a/examples/topic_writer/producer/basic_write/main.cpp b/examples/topic_writer/producer/basic_write/main.cpp new file mode 100644 index 00000000000..83af470330b --- /dev/null +++ b/examples/topic_writer/producer/basic_write/main.cpp @@ -0,0 +1,123 @@ +#include +#include +#include + +#include + +std::shared_ptr CreateProducer(const std::string& topic, NYdb::NTopic::TTopicClient& topicClient) { + NYdb::NTopic::TProducerSettings producerSettings; + producerSettings.Path(topic); + producerSettings.Codec(NYdb::NTopic::ECodec::ZSTD); + producerSettings.ProducerIdPrefix("producer_basic"); + producerSettings.PartitionChooserStrategy(NYdb::NTopic::TProducerSettings::EPartitionChooserStrategy::Bound); + producerSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + producerSettings.MaxBlock(TDuration::Seconds(30)); + producerSettings.MaxMemoryUsage(100_MB); + return topicClient.CreateProducer(producerSettings); +} + +std::string GetResultStatus(const NYdb::NTopic::TWriteResult& result) { + switch (result.Status) { + case NYdb::NTopic::EWriteStatus::Queued: + return "Queued"; + case NYdb::NTopic::EWriteStatus::Error: + return "Error"; + case NYdb::NTopic::EWriteStatus::Timeout: + return "Timeout"; + } +} + +std::string GetResultStatus(const NYdb::NTopic::TFlushResult& result) { + switch (result.Status) { + case NYdb::NTopic::EFlushStatus::Success: + return "Success"; + case NYdb::NTopic::EFlushStatus::ProducerClosed: + return "ProducerClosed"; + } +} + +std::string GetErrorMessage(const NYdb::NTopic::TFlushResult& result) { + std::string errorMessage = "error occurred while flushing messages"; + errorMessage += ", flush status: " + GetResultStatus(result); + errorMessage += ", last written sequence number: " + ToString(result.LastWrittenSeqNo); + if (result.ClosedDescription) { + errorMessage += ", producer is closed: "; + errorMessage += result.ClosedDescription.value().DebugString(); + } + return errorMessage; +} + +std::string GetErrorMessage(const NYdb::NTopic::TWriteResult& result) { + std::string errorMessage = "error occurred while writing message"; + errorMessage += ", write status: " + GetResultStatus(result); + if (result.ErrorMessage) { + errorMessage += ", reason: "; + errorMessage += result.ErrorMessage.value(); + } + if (result.ClosedDescription) { + errorMessage += ", producer is closed: "; + errorMessage += result.ClosedDescription.value().DebugString(); + } + return errorMessage; +} + +void WriteWithHandlingResult(std::shared_ptr producer, NYdb::NTopic::TWriteMessage&& writeMessage) { + static constexpr size_t MAX_RETRIES = 10; + + for (size_t retries = 0; retries < MAX_RETRIES; retries++) { + auto writeResult = producer->Write(std::move(writeMessage)); + if (writeResult.IsSuccess()) { + // if write was successful, we can continue writing messages + continue; + } + + if (writeResult.IsError()) { + // this means that some non retryable error occurred, for example, producer was closed due to user error + // in this case we need to stop retrying and see the close description (to simplify the example, we just print it to standard error) + std::cerr << GetErrorMessage(writeResult) << std::endl; + return; + } + + if (writeResult.IsTimeout()) { + // when timeout occurs this means that producer's buffer is overloaded by memory (see MaxMemoryUsage setting) + // so we need to wait for some time and try to write again later + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + continue; + } + } + + auto flushResult = producer->Flush().GetValueSync(); + if (flushResult.IsSuccess()) { + // if flush was successful, we can return, because all messages were written to the server + return; + } + + if (flushResult.IsClosed()) { + // if flush was not successful, this means that producer was closed due to non retryable error + // in this case we should see the close description (to simplify the example, we just print it to standard error) + std::cerr << GetErrorMessage(flushResult) << std::endl; + } +} + +int main() { + const std::string ENDPOINT = "HOST:PORT"; + const std::string DATABASE = "DATABASE"; + const std::string TOPIC = "PATH/TO/TOPIC"; + + NYdb::TDriverConfig config; + config.SetEndpoint(ENDPOINT); + config.SetDatabase(DATABASE); + NYdb::TDriver driver(config); + + NYdb::NTopic::TTopicClient topicClient(driver); + + auto producer = CreateProducer(TOPIC, topicClient); + auto messageData = std::string(1_KB, 'a'); + + for (int i = 0; i < 10; i++) { + NYdb::NTopic::TWriteMessage writeMessage(messageData); + writeMessage.Key("key" + ToString(i)); + WriteWithHandlingResult(producer, std::move(writeMessage)); + } + return 0; +} diff --git a/include/ydb-cpp-sdk/client/topic/client.h b/include/ydb-cpp-sdk/client/topic/client.h index 88d50cca684..eb6ddb0c673 100644 --- a/include/ydb-cpp-sdk/client/topic/client.h +++ b/include/ydb-cpp-sdk/client/topic/client.h @@ -3,6 +3,7 @@ #include "control_plane.h" #include "read_session.h" #include "write_session.h" +#include "producer.h" namespace NYdb::inline V3::NTopic { @@ -49,11 +50,14 @@ class TTopicClient { //! Create write session. std::shared_ptr CreateSimpleBlockingWriteSession(const TWriteSessionSettings& settings); - //! Create simple blocking keyed write session. Experimental feature. DO NOT USE IN PRODUCTION. - std::shared_ptr CreateSimpleBlockingKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + //! Create producer. Experimental feature. DO NOT USE IN PRODUCTION. + std::shared_ptr CreateProducer(const TProducerSettings& settings); - //! Create keyed write session. Experimental feature. DO NOT USE IN PRODUCTION. - std::shared_ptr CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + //! Create typed producer. + template + std::shared_ptr> CreateTypedProducer(const TProducerSettings& settings) { + return std::make_shared>(CreateProducer(settings)); + } //! Create write session. std::shared_ptr CreateWriteSession(const TWriteSessionSettings& settings); diff --git a/include/ydb-cpp-sdk/client/topic/producer.h b/include/ydb-cpp-sdk/client/topic/producer.h new file mode 100644 index 00000000000..de48becd82e --- /dev/null +++ b/include/ydb-cpp-sdk/client/topic/producer.h @@ -0,0 +1,202 @@ +#pragma once + +#include "write_session.h" + +#include +#include + +namespace NYdb::NTopic { + +struct TProducerSettings : public TWriteSessionSettings { + using TSelf = TProducerSettings; + + enum class EPartitionChooserStrategy { + Bound, + Hash, + }; + + TProducerSettings() = default; + TProducerSettings(const TProducerSettings&) = default; + TProducerSettings(TProducerSettings&&) = default; + + TProducerSettings& operator=(const TProducerSettings&) = default; + TProducerSettings& operator=(TProducerSettings&&) = default; + + //! Session lifetime. + FLUENT_SETTING_DEFAULT(TDuration, SubSessionIdleTimeout, TDuration::Seconds(30)); + + //! Partition chooser strategy. + FLUENT_SETTING_DEFAULT(EPartitionChooserStrategy, PartitionChooserStrategy, EPartitionChooserStrategy::Bound); + + //! Hasher function. + FLUENT_SETTING_DEFAULT(std::function, PartitioningKeyHasher, DefaultPartitioningKeyHasher); + + //! Default partitioning key hasher. + //! Uses MurmurHash. + static std::string DefaultPartitioningKeyHasher(const std::string_view key); + + //! ProducerId prefix to use. + //! ProducerId is generated as ProducerIdPrefix + partition id. + FLUENT_SETTING(std::string, ProducerIdPrefix); + + //! SessionID to use. + FLUENT_SETTING_DEFAULT(std::string, SessionId, ""); + + //! Maximum block time for write. If set, write will block for up to MaxBlockMs when the buffer is overloaded. + FLUENT_SETTING_DEFAULT(TDuration, MaxBlock, TDuration::Zero()); + + //! Key producer function. + FLUENT_SETTING_OPTIONAL(std::function, KeyProducer); + +private: + using TWriteSessionSettings::ProducerId; +}; + +//! Write status. +//! If write was successfully added to buffer, returns Queued. +//! If write was not successful because of closed session, returns Closed. +//! If write was not successful because of timeout, returns Timeout. Usually it means that the producer's buffer is overloaded. You can try to increase MaxBlock setting or try to write later. +//! If write was not successful because of error, returns Error. +enum class EWriteStatus { + Queued = 0, + Timeout = 1, + Error = 2, +}; + +//! Flush status. +//! If flush was successful, returns Success. This status means that all messages in buffer were persistently written to the server. +//! If flush was not successful because of closed session, returns Closed. +//! If flush was not successful because of timeout, returns Timeout. +enum class EFlushStatus { + Success = 0, + ProducerClosed = 1, +}; + +//! Result of write operation. +//! If write was successful, returns Queued. This status means that write was successfully added to buffer. +//! If write was not successful due to overloaded buffer, returns Overloaded. +//! If write was not successful because of closed session, returns Closed. +//! If write was not successful because of timeout, returns Timeout. Usually it means that the producer's buffer is overloaded. You can try to increase MaxBlock setting or try to write later. +//! If write was not successful because of error, returns Error. +struct TWriteResult { + //! Status of write operation. + EWriteStatus Status; + //! Error message. + //! Value is empty if the write was successful. + std::optional ErrorMessage = std::nullopt; + //! Description why session was closed. + //! Value is std::nullopt if the session is not closed. + std::optional ClosedDescription; + + bool IsSuccess() const { + return Status == EWriteStatus::Queued; + } + + bool IsTimeout() const { + return Status == EWriteStatus::Timeout; + } + + bool IsError() const { + return Status == EWriteStatus::Error; + } +}; + +//! Result of flush operation. +//! If flush was successful, returns Success. This status means that all messages in buffer were persistently written to the server. +//! If flush was not successful because of closed session, returns Closed. +//! If flush was not successful because of timeout, returns Timeout. +struct TFlushResult { + //! Status of flush operation. + EFlushStatus Status; + //! Last written sequence number. + std::uint64_t LastWrittenSeqNo; + //! Description why session was closed. + //! Value is std::nullopt if the session is not closed. + std::optional ClosedDescription; + + bool IsSuccess() const { + return Status == EFlushStatus::Success; + } + + bool IsClosed() const { + return Status == EFlushStatus::ProducerClosed; + } +}; + +//! Statistics of write operations. +struct TWriteStats { + //! Last written sequence number. If messages do not have sequence numbers, returns std::nullopt. + std::optional LastWrittenSeqNo; + //! Number of messages written. + std::uint64_t MessagesWritten; +}; + +//! Producer is an abstraction that can write messages to the topic. +//! It has three versions of Write method: +//! - Write without key and partition (partition is chosen by key generated by KeyProducer function in ProducerSettings or by ProducerIdPrefix interpreted as key) +//! - Write with key (partition is chosen based on the key) +//! - Write with partition +//! EXPERIMENTAL SDK, DO NOT USE IN PRODUCTION. +class IProducer { +public: + //! Write single message to partition. + //! Returns write result. + //! If write was successful, returns Queued. + //! If write was not successful due to overloaded buffer, returns Overloaded. + //! If write was not successful because of closed session, returns ProducerClosed. + //! DO NOT IGNORE THE RETURN VALUE. + [[nodiscard]] virtual TWriteResult Write(TWriteMessage&& message) = 0; + + //! Flush all messages to the server. + //! Returns future that is set when flush is complete. + //! If flush was successful, returns TFlushResult with Status Success and LastWrittenSeqNo set to the last written sequence number. + //! If flush was not successful because of closed session, returns TFlushResult with Status ProducerClosed and ClosedDescription set to the description why session was closed. + //! If flush was not successful because of timeout, returns TFlushResult with Status Timeout. + [[nodiscard]] virtual NThreading::TFuture Flush() = 0; + + //! Close the producer. + //! Returns close result. + //! If close was successful, returns Success. This status means that all messages in buffer were persistently written to the server. + //! If close was not successful because of timeout, returns Timeout. + //! If close was not successful because of error, returns Error. + //! DO NOT IGNORE THE RETURN VALUE. + [[nodiscard]] virtual TCloseResult Close(TDuration closeTimeout = TDuration::Max()) = 0; + + //! Get statistics of write operations. + //! Returns statistics of write operations. + virtual TWriteStats GetWriteStats() = 0; + + virtual ~IProducer() = default; +}; + +//! Typed producer — inherits from IProducer; adds only Write(T&&). +//! Type T is implicitly converted to TWriteMessage (e.g. via TWriteMessage's template constructor). +template +class TTypedProducer { +public: + explicit TTypedProducer(std::shared_ptr impl) + : Impl_(impl) + {} + + //! Write single message. T converts to TWriteMessage, then IProducer::Write is called. + [[nodiscard]] TWriteResult Write(T&& message) { + return Impl_->Write(TWriteMessage(std::forward(message))); + } + + [[nodiscard]] NThreading::TFuture Flush() { + return Impl_->Flush(); + } + + [[nodiscard]] TCloseResult Close(TDuration closeTimeout = TDuration::Max()) { + return Impl_->Close(closeTimeout); + } + + TWriteStats GetWriteStats() { + return Impl_->GetWriteStats(); + } + +private: + std::shared_ptr Impl_; +}; + +} // namespace NYdb::NTopic diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index ceef9701fe9..89752c6feb6 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -14,6 +14,44 @@ namespace NYdb::inline V3::NTopic { +//! Result of close operation. +//! If close was successful, returns Success. This status means that all messages in buffer were persistently written to the server. +//! If close was not successful because of timeout, returns Timeout. +//! If close was not successful because of error, returns Error. +enum class ECloseStatus { + Success = 0, + Timeout = 1, + Error = 2, + AlreadyClosed = 3, +}; + +//! Description why session was closed. +struct TCloseDescription : public TSessionClosedEvent {}; + +//! Result of close operation. +struct TCloseResult { + //! Status of close operation. + ECloseStatus Status; + //! Description why session was closed. + std::optional ClosedDescription; + + bool IsSuccess() const { + return Status == ECloseStatus::Success; + } + + bool IsTimeout() const { + return Status == ECloseStatus::Timeout; + } + + bool IsError() const { + return Status == ECloseStatus::Error; + } + + bool IsAlreadyClosed() const { + return Status == ECloseStatus::AlreadyClosed; + } +}; + //! Settings for write session. struct TWriteSessionSettings : public TRequestSettings { using TSelf = TWriteSessionSettings; @@ -110,6 +148,7 @@ struct TWriteSessionSettings : public TRequestSettings { //! Function to handle ReadyToAccept event. //! If this handler is set, write these events will be handled by handler, //! otherwise sent to TWriteSession::GetEvent(). + //! NOTE: DO NOT USE THIS HANDLER IN IProducer INTERFACE. FLUENT_SETTING(TReadyToAcceptHandler, ReadyToAcceptHandler); //! Function to handle close session events. @@ -144,53 +183,94 @@ struct TWriteSessionSettings : public TRequestSettings { FLUENT_SETTING_DEFAULT(bool, ValidateSeqNo, true); }; -struct TKeyedWriteSessionSettings : public TWriteSessionSettings { - using TSelf = TKeyedWriteSessionSettings; - - enum class EPartitionChooserStrategy { - Bound, - Hash, +template +concept Serializable = + requires(const T& t) { + { Serialize(t) } -> std::convertible_to; }; - TKeyedWriteSessionSettings() = default; - TKeyedWriteSessionSettings(const TKeyedWriteSessionSettings&) = default; - TKeyedWriteSessionSettings(TKeyedWriteSessionSettings&&) = default; +//! Contains the message to write and all the options. +struct TWriteMessage { + using TSelf = TWriteMessage; + using TMessageMeta = std::vector>; +private: + //! This field is used to store serialized data. + std::optional DataHolder; - TKeyedWriteSessionSettings& operator=(const TKeyedWriteSessionSettings&) = default; - TKeyedWriteSessionSettings& operator=(TKeyedWriteSessionSettings&&) = default; +public: + TWriteMessage() = delete; + TWriteMessage(std::string_view data) + : Data(data) + {} - //! Session lifetime. - FLUENT_SETTING_DEFAULT(TDuration, SubSessionIdleTimeout, TDuration::Seconds(30)); + TWriteMessage(const TWriteMessage& other) + : DataHolder(other.DataHolder) + , Data(other.DataHolder ? std::string_view(*DataHolder) : other.Data) + , Codec(other.Codec) + , OriginalSize(other.OriginalSize) + , SeqNo_(other.SeqNo_) + , CreateTimestamp_(other.CreateTimestamp_) + , MessageMeta_(other.MessageMeta_) + , Key_(other.Key_) + , Partition_(other.Partition_) + , Tx_(other.Tx_) + {} - //! Partition chooser strategy. - FLUENT_SETTING_DEFAULT(EPartitionChooserStrategy, PartitionChooserStrategy, EPartitionChooserStrategy::Bound); + TWriteMessage(TWriteMessage&& other) noexcept + : DataHolder(std::move(other.DataHolder)) + , Data(DataHolder ? std::string_view(*DataHolder) : other.Data) + , Codec(std::move(other.Codec)) + , OriginalSize(other.OriginalSize) + , SeqNo_(std::move(other.SeqNo_)) + , CreateTimestamp_(std::move(other.CreateTimestamp_)) + , MessageMeta_(std::move(other.MessageMeta_)) + , Key_(std::move(other.Key_)) + , Partition_(std::move(other.Partition_)) + , Tx_(std::move(other.Tx_)) + {} - //! Hasher function. - FLUENT_SETTING_DEFAULT(std::function, PartitioningKeyHasher, DefaultPartitioningKeyHasher); + TWriteMessage& operator=(const TWriteMessage& other) { + if (this == &other) { + return *this; + } - //! Default partitioning key hasher. - //! Uses MurmurHash. - static std::string DefaultPartitioningKeyHasher(const std::string_view key); + DataHolder = other.DataHolder; + Data = DataHolder ? std::string_view(*DataHolder) : other.Data; + Codec = other.Codec; + OriginalSize = other.OriginalSize; + SeqNo_ = other.SeqNo_; + CreateTimestamp_ = other.CreateTimestamp_; + MessageMeta_ = other.MessageMeta_; + Key_ = other.Key_; + Partition_ = other.Partition_; + Tx_ = other.Tx_; - //! ProducerId prefix to use. - //! ProducerId is generated as ProducerIdPrefix + partition id. - FLUENT_SETTING(std::string, ProducerIdPrefix); + return *this; + } - //! SessionID to use. - FLUENT_SETTING_DEFAULT(std::string, SessionId, ""); + TWriteMessage& operator=(TWriteMessage&& other) noexcept { + if (this == &other) { + return *this; + } -private: - using TWriteSessionSettings::ProducerId; -}; + DataHolder = std::move(other.DataHolder); + Data = DataHolder ? std::string_view(*DataHolder) : other.Data; + Codec = std::move(other.Codec); + OriginalSize = other.OriginalSize; + SeqNo_ = std::move(other.SeqNo_); + CreateTimestamp_ = std::move(other.CreateTimestamp_); + MessageMeta_ = std::move(other.MessageMeta_); + Key_ = std::move(other.Key_); + Partition_ = std::move(other.Partition_); + Tx_ = std::move(other.Tx_); -//! Contains the message to write and all the options. -struct TWriteMessage { - using TSelf = TWriteMessage; - using TMessageMeta = std::vector>; -public: - TWriteMessage() = delete; - TWriteMessage(std::string_view data) - : Data(data) + return *this; + } + + template + TWriteMessage(const T& data) + : DataHolder(Serialize(data)) + , Data(*DataHolder) {} //! A message that is already compressed by codec. Codec from WriteSessionSettings does not apply to this message. @@ -205,7 +285,6 @@ struct TWriteMessage { bool Compressed() const { return Codec.has_value(); } - //! Message body. std::string_view Data; @@ -225,6 +304,12 @@ struct TWriteMessage { //! Message metadata. Limited to 4096 characters overall (all keys and values combined). FLUENT_SETTING(TMessageMeta, MessageMeta); + //! Message key. It will be used to route message to the partition. + FLUENT_SETTING_OPTIONAL(std::string, Key); + + //! Partition to write to. It is not recommended to use this option, use Key instead. + FLUENT_SETTING_OPTIONAL(std::uint32_t, Partition); + //! Transaction id FLUENT_SETTING_OPTIONAL(std::reference_wrapper, Tx); @@ -315,48 +400,4 @@ class IWriteSession { virtual ~IWriteSession() = default; }; -//! Keyed write session. Experimental SDK. DO NOT USE IN PRODUCTION. -class IKeyedWriteSession { -public: - //! Write single message. - //! continuationToken - a token earlier provided to client with ReadyToAccept event. - virtual void Write(TContinuationToken&& continuationToken, const std::string& key, TWriteMessage&& message, - TTransactionBase* tx = nullptr) = 0; - - //! Future that is set when next event is available. - virtual NThreading::TFuture WaitEvent() = 0; - - //! Wait and return next event. Use WaitEvent() for non-blocking wait. - virtual std::optional GetEvent(bool block = false) = 0; - - //! Get several events in one call. - //! If blocking = false, instantly returns up to maxEventsCount available events. - //! If blocking = true, blocks till maxEventsCount events are available. - //! If maxEventsCount is unset, write session decides the count to return itself. - virtual std::vector GetEvents(bool block = false, std::optional maxEventsCount = std::nullopt) = 0; - - virtual bool Close(TDuration closeTimeout = TDuration::Max()) = 0; - virtual TWriterCounters::TPtr GetCounters() = 0; - virtual ~IKeyedWriteSession() = default; -}; - -//! Simple blocking keyed write session. Experimental SDK. DO NOT USE IN PRODUCTION. -class ISimpleBlockingKeyedWriteSession { -public: - //! Write single message. - //! continuationToken - a token earlier provided to client with ReadyToAccept event. - virtual bool Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx = nullptr, - TDuration blockTimeout = TDuration::Max()) = 0; - - //! Wait for all writes to complete (no more that closeTimeout()), then close. - //! Return true if all writes were completed and acked, false if timeout was reached and some writes were aborted. - virtual bool Close(TDuration closeTimeout = TDuration::Max()) = 0; - - //! Writer counters with different stats (see TWriterConuters). - virtual TWriterCounters::TPtr GetCounters() = 0; - - //! Close() with timeout = 0 and destroy everything instantly. - virtual ~ISimpleBlockingKeyedWriteSession() = default; -}; - } // namespace NYdb::NTopic diff --git a/src/client/topic/impl/producer.cpp b/src/client/topic/impl/producer.cpp new file mode 100644 index 00000000000..8057902e0fa --- /dev/null +++ b/src/client/topic/impl/producer.cpp @@ -0,0 +1,1882 @@ +#include +#include +#include +#include +#include + +namespace NYdb::inline V3::NTopic { + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer + +// TProducerSettings + +std::string TProducerSettings::DefaultPartitioningKeyHasher(const std::string_view key) { + const std::uint64_t lo = MurmurHash(key.data(), key.size(), std::uint64_t{0}); + const std::uint64_t loBe = InetToHost(lo); + + std::string out; + out.resize(8); + memcpy(out.data(), &loBe, 8); + return out; // 8 bytes +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TPartitionInfo + +bool TProducer::TPartitionInfo::InRange(const std::string_view key) const { + if (FromBound_ > key) { + return false; + } + if (ToBound_.has_value() && *ToBound_ <= key) { + return false; + } + return true; +} + +bool TProducer::TPartitionInfo::IsSplitted() const { + return !Children_.empty(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TMessageInfo + +TProducer::TMessageInfo::TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition) + : Key(key) + , Data(message.Data) + , Codec(message.Codec) + , OriginalSize(message.OriginalSize) + , SeqNo(message.SeqNo_) + , CreateTimestamp(message.CreateTimestamp_) + , Tx(message.Tx_) + , Partition(partition) +{ + for (const auto& [key, value] : message.MessageMeta_) { + MessageMeta.Fields.emplace_back(key, value); + } +} + +TWriteMessage TProducer::TMessageInfo::BuildMessage() const { + TWriteMessage message(Data); + message.Codec = Codec; + message.OriginalSize = OriginalSize; + message.SeqNo(SeqNo); + message.CreateTimestamp(CreateTimestamp); + for (const auto& [key, value] : MessageMeta.Fields) { + message.MessageMeta_.emplace_back(key, value); + } + message.Tx(Tx); + return message; +} +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TWriteSessionWrapper + +TProducer::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition) + : Session(std::move(session)) + , Partition(partition) +{} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TIdleSession + +bool TProducer::TIdleSession::Less(const std::shared_ptr& other) const { + if (EmptySince == other->EmptySince) { + return Session->Partition < other->Session->Partition; + } + + return EmptySince < other->EmptySince; +} + +bool TProducer::TIdleSession::Comparator::operator()( + const std::shared_ptr& first, + const std::shared_ptr& second) const { + return first->Less(second); +} + +bool TProducer::TIdleSession::IsExpired() const { + return TInstant::Now() - EmptySince > IdleTimeout; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TSplittedPartitionWorker + +TProducer::TSplittedPartitionWorker::TSplittedPartitionWorker(TProducer* producer, std::uint32_t partitionId) + : Producer(producer) + , PartitionId(partitionId) +{ + LOG_LAZY(Producer->DbDriverState->Log, TLOG_INFO, Producer->LogPrefix() << "Creating splitted partition worker for partition " << PartitionId); +} + +std::string TProducer::TSplittedPartitionWorker::GetStateName() const { + switch (State) { + case EState::Init: + return "Init"; + case EState::PendingDescribe: + return "PendingDescribe"; + case EState::GotDescribe: + return "GotDescribe"; + case EState::PendingMaxSeqNo: + return "PendingMaxSeqNo"; + case EState::Done: + return "Done"; + case EState::GotMaxSeqNo: + return "GotMaxSeqNo"; + } +} + +void TProducer::TSplittedPartitionWorker::DoWork() { + std::unique_lock lock(Lock); + std::weak_ptr producer = Producer->shared_from_this(); + switch (State) { + case EState::Init: + DescribeTopicFuture = Producer->Client->DescribeTopic(Producer->Settings.Path_, TDescribeTopicSettings()); + lock.unlock(); + DescribeTopicFuture.Subscribe([this, producer](const NThreading::TFuture&) { + auto producerPtr = producer.lock(); + if (!producerPtr) { + return; + } + + { + std::lock_guard lock(Lock); + MoveTo(EState::GotDescribe); + } + + producerPtr->RunMainWorker(static_cast(PartitionId)); + }); + lock.lock(); + if (State == EState::Init) { + MoveTo(EState::PendingDescribe); + } + break; + case EState::GotDescribe: + HandleDescribeResult(); + if (State != EState::GotDescribe) { + break; + } + + LaunchGetMaxSeqNoFutures(lock); + if (State == EState::GotDescribe) { + MoveTo(EState::PendingMaxSeqNo); + } + break; + case EState::PendingDescribe: + case EState::PendingMaxSeqNo: + case EState::Done: + break; + case EState::GotMaxSeqNo: + Producer->MessagesWorker->RebuildPendingMessagesIndex(PartitionId); + Producer->MessagesWorker->ScheduleResendMessages(PartitionId, MaxSeqNo); + for (const auto& child : Producer->Partitions[PartitionId].Children_) { + Producer->Partitions[child].Locked(false); + } + Producer->Partitions[PartitionId].Locked_ = false; + MoveTo(EState::Done); + break; + } +} + +void TProducer::TSplittedPartitionWorker::MoveTo(EState state) { + State = state; + LOG_LAZY(Producer->DbDriverState->Log, TLOG_INFO, Producer->LogPrefix() << "Moving splitted partition worker for partition " << PartitionId << " to state " << GetStateName()); +} + +void TProducer::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t maxSeqNo) { + MaxSeqNo = std::max(MaxSeqNo, maxSeqNo); +} + +bool TProducer::TSplittedPartitionWorker::IsDone() { + std::lock_guard lock(Lock); + DoneAt = TInstant::Now(); + return State == EState::Done; +} + +bool TProducer::TSplittedPartitionWorker::IsInit() { + std::lock_guard lock(Lock); + return State == EState::Init; +} + +void TProducer::TSplittedPartitionWorker::HandleDescribeResult() { + std::vector newPartitionsIds; + const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); + for (const auto& partition : partitions) { + if (partition.GetPartitionId() != PartitionId) { + continue; + } + + LOG_LAZY(Producer->DbDriverState->Log, TLOG_DEBUG, Producer->LogPrefix() << "Found partition " << partition.GetPartitionId() << " for partition " << PartitionId << " children: " << partition.GetChildPartitionIds().size()); + for (const auto& childPartitionId : partition.GetChildPartitionIds()) { + newPartitionsIds.push_back(childPartitionId); + } + break; + } + + if (newPartitionsIds.empty()) { + // describe response is incomplete, we need to resend describe request + MoveTo(EState::Init); + Y_ABORT_UNLESS(++Retries < 40, "Too many retries for partition %u", PartitionId); + LOG_LAZY(Producer->DbDriverState->Log, TLOG_ERR, Producer->LogPrefix() << "Describe response is incomplete, we need to resend describe request for partition " << PartitionId); + return; + } + + std::vector children; + const auto& splittedPartition = Producer->Partitions[PartitionId]; + Producer->PartitionsIndex.erase(splittedPartition.FromBound_); + + for (const auto& newPartitionId : newPartitionsIds) { + auto partitionDescribeInfo = std::find_if(partitions.begin(), partitions.end(), [newPartitionId](const auto& partition) { + return partition.GetPartitionId() == newPartitionId; + }); + Y_ABORT_UNLESS(partitionDescribeInfo != partitions.end(), "Partition describe info not found"); + Producer->PartitionsIndex[partitionDescribeInfo->GetFromBound().value_or("")] = newPartitionId; + Producer->Partitions[newPartitionId] = TPartitionInfo() + .PartitionId(newPartitionId) + .FromBound(partitionDescribeInfo->GetFromBound().value_or("")) + .ToBound(partitionDescribeInfo->GetToBound()) + .Locked(true); + children.push_back(newPartitionId); + } + + Producer->Partitions[PartitionId].Children(children); +} + +void TProducer::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_lock& lock) { + Y_ABORT_UNLESS(DescribeTopicFuture.IsReady(), "DescribeTopicFuture is not ready yet"); + + std::unordered_map partitionIdToParentId; + const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); + for (const auto& partition : partitions) { + auto parentPartitions = partition.GetParentPartitionIds(); + if (parentPartitions.empty()) { + continue; + } + + // we consider here that each partition has only one parent partition + partitionIdToParentId[partition.GetPartitionId()] = parentPartitions.front(); + } + + std::vector ancestors; + std::uint32_t currentPartitionId = PartitionId; + while (true) { + ancestors.push_back(currentPartitionId); + + auto parentPartitionId = partitionIdToParentId.find(currentPartitionId); + if (parentPartitionId == partitionIdToParentId.end()) { + break; + } + currentPartitionId = parentPartitionId->second; + } + + NotReadyFutures = ancestors.size(); + for (const auto& ancestor : ancestors) { + auto wrappedSession = Producer->SessionsWorker->GetWriteSession(ancestor, false); + Y_ABORT_UNLESS(wrappedSession, "Write session not found"); + WriteSessions.push_back(wrappedSession); + + auto future = wrappedSession->Session->GetInitSeqNo(); + std::weak_ptr producer = Producer->shared_from_this(); + lock.unlock(); + future.Subscribe([this, producer, wrappedSession, ancestor](const NThreading::TFuture& result) { + auto producerPtr = producer.lock(); + if (!producerPtr) { + return; + } + + if (IsDone()) { + return; + } + + bool gotMaxSeqNo = false; + { + std::lock_guard lock(Lock); + if (result.HasException()) { + LOG_LAZY(producerPtr->DbDriverState->Log, TLOG_ERR, producerPtr->LogPrefix() << "Failed to get max seq no for partition " << ancestor << " for splitted partition " << PartitionId); + TSessionClosedEvent sessionClosedEvent(EStatus::INTERNAL_ERROR, {}); + producerPtr->GetSessionClosedEventAndDie(wrappedSession, std::move(sessionClosedEvent)); + MoveTo(EState::Done); + return; + } + + UpdateMaxSeqNo(result.GetValue()); + if (--NotReadyFutures == 0) { + MoveTo(EState::GotMaxSeqNo); + gotMaxSeqNo = true; + } + } + + if (gotMaxSeqNo) { + producerPtr->RunMainWorker(static_cast(PartitionId)); + } + }); + lock.lock(); + GetMaxSeqNoFutures.push_back(future); + } + + if (ancestors.empty()) { + LOG_LAZY(Producer->DbDriverState->Log, TLOG_INFO, Producer->LogPrefix() << "No ancestors found for partition " << PartitionId); + MoveTo(EState::Init); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TEventsWorkerWrapper + +TProducer::TEventsWorker::TEventsWorker(TProducer* producer) + : Producer(producer) +{ + EventsPromise = NThreading::NewPromise(); + EventsFuture = EventsPromise.GetFuture(); +} + +void TProducer::TEventsWorker::HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); + queueIt->second.push_back(TWriteSessionEvent::TEvent(std::move(event))); +} + +void TProducer::TEventsWorker::HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event) { + Producer->MessagesWorker->HandleContinuationToken(partition, std::move(event.ContinuationToken)); +} + +void TProducer::TEventsWorker::HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition) { + if (event.IsSuccess()) { + return; + } + + Producer->Partitions[partition].Locked_ = true; + + if (event.GetStatus() == EStatus::OVERLOADED) { + Producer->HandleAutoPartitioning(partition); + return; + } + + if (!CloseEvent.has_value()) { + CloseEvent = std::move(event); + } + Producer->NonBlockingClose(); +} + +bool TProducer::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition) { + while (true) { + auto event = wrappedSession->Session->GetEvent(false); + if (!event) { + break; + } + + if (auto sessionClosedEvent = std::get_if(&*event); sessionClosedEvent) { + HandleSessionClosedEvent(std::move(*sessionClosedEvent), partition); + return true; + } + + if (auto readyToAcceptEvent = std::get_if(&*event)) { + HandleReadyToAcceptEvent(partition, std::move(*readyToAcceptEvent)); + continue; + } + + if (auto acksEvent = std::get_if(&*event)) { + // Producer->SessionsWorker->OnReadFromSession(wrappedSession, acksEvent->Acks.size()); + HandleAcksEvent(partition, std::move(*acksEvent)); + continue; + } + } + + return false; +} + +std::optional> TProducer::TEventsWorker::DoWork() { + std::unique_lock lock(Lock); + + while (!ReadyFutures.empty()) { + auto idx = *ReadyFutures.begin(); + ReadyFutures.erase(idx); + lock.unlock(); + // RunEventLoop without Lock: sub-session's WaitEvent() completion may run the Subscribe + // callback (ReadyFutures.insert) synchronously; that callback takes Lock -> same-thread deadlock. + auto isSessionClosed = RunEventLoop(Producer->SessionsWorker->GetWriteSession(idx), idx); + if (!isSessionClosed) { + SubscribeToPartition(idx); + } else { + UnsubscribeFromPartition(idx); + } + lock.lock(); + } + + if (!Producer->Done.load() && TransferEventsToOutputQueue()) { + return EventsPromise; + } + + return std::nullopt; +} + +void TProducer::TEventsWorker::SubscribeToPartition(std::uint32_t partition) { + if (Producer->Partitions[partition].IsSplitted() || Producer->SplittedPartitionWorkers.contains(partition)) { + Producer->Partitions[partition].Future(NThreading::MakeFuture()); + return; + } + + auto wrappedSession = Producer->SessionsWorker->GetWriteSession(partition); + auto newFuture = wrappedSession->Session->WaitEvent(); + std::weak_ptr producer = Producer->shared_from_this(); + std::weak_ptr self = shared_from_this(); + + newFuture.Subscribe([self, producer, partition](const NThreading::TFuture&) { + auto producerPtr = producer.lock(); + if (!producerPtr) { + return; + } + + auto selfPtr = self.lock(); + if (!selfPtr) { + return; + } + + { + std::lock_guard lock(selfPtr->Lock); + selfPtr->ReadyFutures.insert(partition); + } + producerPtr->RunMainWorker(static_cast(partition)); + }); + Producer->Partitions[partition].Future(newFuture); +} + +std::optional TProducer::TEventsWorker::GetSessionClosedEvent() { + std::unique_lock lock(Lock); + if (CloseEvent.has_value()) { + return CloseEvent; + } + return std::nullopt; +} + +std::optional> TProducer::TEventsWorker::HandleNewMessage() { + std::lock_guard lock(Lock); + if (Producer->MessagesWorker->IsMemoryUsageOK()) { + AddContinuationToken(); + return EventsPromise; + } + + Producer->Metrics.IncBufferFull(); + return std::nullopt; +} + +void TProducer::TEventsWorker::AddContinuationToken() { + auto continuationToken = IssueContinuationToken(); + TokensQueue.push_back(std::move(continuationToken)); + Producer->Metrics.IncContinuationTokensSent(); +} + +bool TProducer::TEventsWorker::AddSessionClosedIfNeeded() { + if (!Producer->Closed.load()) { + return false; + } + + if (!CloseEvent.has_value()) { + CloseEvent = TSessionClosedEvent(EStatus::SUCCESS, {}); + } + + if (EventsOutputQueue.empty() && (Producer->MessagesWorker->IsQueueEmpty() || Producer->Done.load())) { + EventsOutputQueue.push_back(*CloseEvent); + return true; + } + + return false; +} + +bool TProducer::TEventsWorker::TransferEventsToOutputQueue() { + bool eventsTransferred = false; + bool shouldAddContinuationToken = false; + std::unordered_map> acks; + + auto messagesWorker = Producer->MessagesWorker; + auto buildOutputAckEvent = [&](std::deque& acksQueue, std::uint64_t partition, std::optional expectedSeqNo) -> TWriteSessionEvent::TAcksEvent { + TWriteSessionEvent::TAcksEvent ackEvent; + + if (expectedSeqNo.has_value()) { + Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << Producer->Partitions[partition].PartitionId_); + } + + auto ack = std::move(acksQueue.front()); + ackEvent.Acks.push_back(std::move(ack)); + acksQueue.pop_front(); + return ackEvent; + }; + auto finishWithAck = [this, messagesWorker, &shouldAddContinuationToken](std::uint64_t seqNo) { + Producer->LastWrittenSeqNo = std::max(Producer->LastWrittenSeqNo, seqNo); + Producer->MessagesWritten++; + bool wasMemoryUsageOk = messagesWorker->IsMemoryUsageOK(); + messagesWorker->HandleAck(); + if (messagesWorker->IsMemoryUsageOK() && !wasMemoryUsageOk) { + shouldAddContinuationToken = true; + } + }; + + while (messagesWorker->HasInFlightMessages()) { + const auto& head = messagesWorker->GetFrontInFlightMessage(); + + auto remainingAcks = acks.find(head.Partition); + if (remainingAcks != acks.end() && remainingAcks->second.size() > 0) { + auto seqNo = remainingAcks->second.front().SeqNo; + EventsOutputQueue.push_back(buildOutputAckEvent(remainingAcks->second, head.Partition, head.SeqNo)); + finishWithAck(seqNo); + continue; + } + + auto eventsQueueIt = PartitionsEventQueues.find(head.Partition); + if (eventsQueueIt == PartitionsEventQueues.end() || eventsQueueIt->second.empty()) { + // No events for this message yet, stop processing (preserve order) + break; + } + + auto event = std::move(eventsQueueIt->second.front()); + auto acksEvent = std::get_if(&event); + Y_ABORT_UNLESS(acksEvent, "Expected AcksEvent only in PartitionsEventQueues"); + + std::deque acksQueue; + std::copy(acksEvent->Acks.begin(), acksEvent->Acks.end(), std::back_inserter(acksQueue)); + auto seqNo = acksEvent->Acks.front().SeqNo; + EventsOutputQueue.push_back(buildOutputAckEvent(acksQueue, head.Partition, head.SeqNo)); + acks[head.Partition] = std::move(acksQueue); + eventsQueueIt->second.pop_front(); + eventsTransferred = true; + + finishWithAck(seqNo); + } + + // this case handles situation: + // 1st message is written to partition 0 + // 2nd message is written to partition 1 + // 3rd message is written to partition 0 + // 4th message is written to partition 1 + // but AcksEvent for partition 0 looks like: + // [ack1, ack3] + // In this case we can not just forget about ack3, because 3rd message is in-flight + // so we will push 'AcksEvent' back to the queue for partition 0 + for (auto& [partition, acksQueue] : acks) { + if (acksQueue.size() > 0) { + TWriteSessionEvent::TAcksEvent ackEvent; + std::copy(acksQueue.begin(), acksQueue.end(), std::back_inserter(ackEvent.Acks)); + PartitionsEventQueues[partition].push_front(std::move(ackEvent)); + } + } + + if (shouldAddContinuationToken) { + AddContinuationToken(); + } + + return eventsTransferred; +} + +std::optional TProducer::TEventsWorker::GetContinuationToken() { + std::lock_guard lock(Lock); + if (TokensQueue.empty()) { + return std::nullopt; + } + + auto continuationToken = std::move(TokensQueue.front()); + TokensQueue.pop_front(); + return std::move(continuationToken); +} + +std::list::iterator TProducer::TEventsWorker::AckQueueBegin(std::uint32_t partition) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); + return queueIt->second.begin(); +} + +std::list::iterator TProducer::TEventsWorker::AckQueueEnd(std::uint32_t partition) { + auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); + return queueIt->second.end(); +} + +TProducer::TEventsWorker::EEventType TProducer::TEventsWorker::GetEventType(const TWriteSessionEvent::TEvent& event) { + if (std::holds_alternative(event)) { + return EEventType::SessionClosed; + } else if (std::holds_alternative(event)) { + return EEventType::ReadyToAccept; + } else if (std::holds_alternative(event)) { + return EEventType::Ack; + } + + Y_ABORT_UNLESS(false, "Unexpected event type"); +} + +std::optional TProducer::TEventsWorker::GetEventImpl(bool block, const std::vector& eventTypes) { + std::unique_lock lock(Lock); + if (EventsOutputQueue.empty() && block) { + lock.unlock(); + WaitEvent().Wait(); + lock.lock(); + } + + if (!EventsOutputQueue.empty()) { + if (!eventTypes.empty() && std::find(eventTypes.begin(), eventTypes.end(), GetEventType(EventsOutputQueue.front())) == eventTypes.end()) { + return std::nullopt; + } + + auto event = std::move(EventsOutputQueue.front()); + EventsOutputQueue.pop_front(); + return event; + } + + return std::nullopt; +} + +std::optional TProducer::TEventsWorker::GetEvent(bool block, const std::vector& eventTypes) { + { + std::unique_lock lock(Lock); + AddSessionClosedIfNeeded(); + } + auto event = GetEventImpl(block, eventTypes); + + return event; +} + +std::vector TProducer::TEventsWorker::GetEvents(bool block, std::optional maxEventsCount, const std::vector& eventTypes) { + if (maxEventsCount.has_value() && maxEventsCount.value() == 0) { + return {}; + } + + { + std::unique_lock lock(Lock); + AddSessionClosedIfNeeded(); + } + + std::vector events; + while (true) { + auto event = GetEventImpl(block, eventTypes); + if (!event) { + break; + } + + events.push_back(std::move(*event)); + if (maxEventsCount.has_value() && events.size() >= maxEventsCount.value()) { + break; + } + } + + return events; +} + +NThreading::TFuture TProducer::TEventsWorker::WaitEvent() { + std::unique_lock lock(Lock); + + AddSessionClosedIfNeeded(); + if (!EventsOutputQueue.empty()) { + return NThreading::MakeFuture(); + } + + if (EventsFuture.IsReady() && !Producer->Closed.load()) { + EventsPromise = NThreading::NewPromise(); + EventsFuture = EventsPromise.GetFuture(); + } + + return EventsFuture; +} + +void TProducer::TEventsWorker::UnsubscribeFromPartition(std::uint32_t partition) { + ReadyFutures.erase(partition); + Producer->Partitions[partition].Future(NThreading::MakeFuture()); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TSessionsWorker + +TProducer::TSessionsWorker::TSessionsWorker(TProducer* producer) + : Producer(producer) +{} + +TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::GetWriteSession(std::uint32_t partition, bool directToPartition) { + auto sessionIter = SessionsIndex.find(partition); + if (sessionIter == SessionsIndex.end() || !directToPartition) { + return CreateWriteSession(partition, directToPartition); + } + + return sessionIter->second; +} + +std::string TProducer::TSessionsWorker::GetProducerId(std::uint32_t partitionId) { + return std::format("{}_{}", Producer->Settings.ProducerIdPrefix_, partitionId); +} + +TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::CreateWriteSession(std::uint32_t partition, bool directToPartition) { + auto partitionId = Producer->Partitions[partition].PartitionId_; + auto producerId = GetProducerId(partitionId); + TWriteSessionSettings alteredSettings = Producer->Settings; + + alteredSettings + .ProducerId(producerId) + .MessageGroupId(producerId) + .MaxMemoryUsage(std::numeric_limits::max()) + .RetryPolicy(Producer->RetryPolicy) + .EventHandlers(TWriteSessionSettings::TEventHandlers() + .ReadyToAcceptHandler({}) + .AcksHandler({}) + .SessionClosedHandler({})); + + if (directToPartition) { + alteredSettings.DirectWriteToPartition(true); + alteredSettings.PartitionId(partitionId); + } + auto writeSession = std::make_shared( + Producer->Client->CreateWriteSession(alteredSettings), + partition); + + if (directToPartition) { + SessionsIndex.emplace(partition, writeSession); + Producer->EventsWorker->SubscribeToPartition(partition); + } + return writeSession; +} + +void TProducer::TSessionsWorker::DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout) { + if (it == SessionsIndex.end() || !it->second) { + return; + } + + it->second->Session->Close(closeTimeout); + const auto partition = it->second->Partition; + if (static_cast(partition) == Producer->MainWorkerOwner) { + SessionsToRemove.push_back(it->second); + } + it = SessionsIndex.erase(it); + Producer->EventsWorker->UnsubscribeFromPartition(partition); +} + +size_t TProducer::TSessionsWorker::GetSessionsCount() const { + return SessionsIndex.size(); +} + +size_t TProducer::TSessionsWorker::GetIdleSessionsCount() const { + return IdlerSessions.size(); +} + +void TProducer::TSessionsWorker::AddIdleSession(std::uint32_t partition) { + auto wrappedSession = SessionsIndex.find(partition); + if (wrappedSession == SessionsIndex.end()) { + return; + } + + if (wrappedSession->second->IdleSession) { + return; + } + + auto idleSessionPtr = std::make_shared(wrappedSession->second.get(), TInstant::Now(), Producer->Settings.SubSessionIdleTimeout_); + auto [itIdle, inserted] = IdlerSessions.insert(idleSessionPtr); + Y_ABORT_UNLESS(inserted, "Duplicate idle session for partition"); + IdlerSessionsIndex[partition] = itIdle; + wrappedSession->second->IdleSession = idleSessionPtr; +} + +void TProducer::TSessionsWorker::RemoveIdleSession(std::uint32_t partition) { + auto itIdle = IdlerSessionsIndex.find(partition); + if (itIdle == IdlerSessionsIndex.end()) { + return; + } + + auto wrappedSession = SessionsIndex.find(partition); + if (wrappedSession == SessionsIndex.end()) { + return; + } + + IdlerSessions.erase(itIdle->second); + IdlerSessionsIndex.erase(itIdle); + wrappedSession->second->IdleSession.reset(); +} + +void TProducer::TSessionsWorker::DoWork() { + while (!SessionsToRemove.empty()) { + if (static_cast(SessionsToRemove.front()->Partition) == Producer->MainWorkerOwner) { + break; + } + + SessionsToRemove.pop_front(); + } + + while (!IdlerSessions.empty()) { + auto it = IdlerSessions.begin(); + if (!(*it)->IsExpired()) { + break; + } + + LOG_LAZY(Producer->DbDriverState->Log, TLOG_DEBUG, TStringBuilder() << Producer->LogPrefix() << "Removing idle session for partition " << (*it)->Session->Partition); + + const auto partition = (*it)->Session->Partition; + if (Producer->Partitions[partition].Locked_) { + continue; + } + + // Remove idle tracking first to keep containers consistent even if the session + // is already absent from SessionsIndex. + IdlerSessions.erase(it); + IdlerSessionsIndex.erase(partition); + + auto sessionIter = SessionsIndex.find(partition); + if (sessionIter != SessionsIndex.end()) { + sessionIter->second->IdleSession.reset(); + DestroyWriteSession(sessionIter, TDuration::Zero()); + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TMessagesWorker + +TProducer::TMessagesWorker::TMessagesWorker(TProducer* producer) + : Producer(producer) +{ +} + +void TProducer::TMessagesWorker::RechoosePartitionIfNeeded(MessageIter message) { + const auto& partitionInfo = Producer->Partitions[message->Partition]; + if (partitionInfo.Children_.empty()) { + return; + } + + // this case means that partition was split, so we need to rechoose the partition for the message + auto newPartition = Producer->PartitionChooser->ChoosePartition(message->Key); + message->Partition = newPartition; +} + +void TProducer::TMessagesWorker::DoWork() { + auto sessionsWorker = Producer->SessionsWorker; + + auto iterateMessagesIndex = [&](std::unordered_map>& messagesIndex, auto stopCondition) { + std::vector partitionsProcessed; + for (auto& [partition, messages] : messagesIndex) { + while (!messages.empty()) { + auto head = messages.front(); + if (stopCondition(head)) { + break; + } + + auto wrappedSession = sessionsWorker->GetWriteSession(head->Partition); + if (!SendMessage(wrappedSession, *head)) { + break; + } + + Producer->Metrics.AddWriteLag((TInstant::Now() - head->CreateTimestamp.value_or(TInstant::Now())).MilliSeconds()); + head->Sent = true; + // sessionsWorker->OnWriteToSession(wrappedSession); + messages.pop_front(); + } + + if (messages.empty()) { + partitionsProcessed.push_back(partition); + } + } + + for (const auto& partition : partitionsProcessed) { + messagesIndex.erase(partition); + } + }; + + iterateMessagesIndex( + MessagesToResendIndex, + [](MessageIter) { + return false; + } + ); + + iterateMessagesIndex( + PendingMessagesIndex, + [this](MessageIter head) { + return Producer->Partitions[head->Partition].Locked_ || + MessagesToResendIndex.contains(head->Partition); + } + ); +} + +bool TProducer::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message) { + if (!wrappedSession) { + return false; + } + + auto continuationToken = GetContinuationToken(message.Partition); + if (!continuationToken) { + return false; + } + + Producer->Metrics.IncOutgoingMessages(); + auto builtMessage = message.BuildMessage(); + wrappedSession->Session->Write(std::move(*continuationToken), std::move(builtMessage)); + return true; +} + +void TProducer::TMessagesWorker::PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message) { + auto iter = InFlightMessages.insert(InFlightMessages.end(), std::move(message)); + auto [inFlightMessagesIndexIt, wasInserted] = InFlightMessagesIndex.try_emplace(partition); + inFlightMessagesIndexIt->second.push_back(iter); + + auto [pendingMessagesIndexIt, __] = PendingMessagesIndex.try_emplace(partition); + pendingMessagesIndexIt->second.push_back(iter); + + if (wasInserted) { + Producer->SessionsWorker->RemoveIdleSession(partition); + } +} + +void TProducer::TMessagesWorker::HandleAck() { + PopInFlightMessage(); +} + +void TProducer::TMessagesWorker::PopInFlightMessage() { + Y_ABORT_UNLESS(!InFlightMessages.empty()); + const std::uint64_t partition = InFlightMessages.front().Partition; + const auto it = InFlightMessages.begin(); + + auto mapIt = InFlightMessagesIndex.find(partition); + if (mapIt != InFlightMessagesIndex.end()) { + auto& list = mapIt->second; + for (auto listIt = list.begin(); listIt != list.end(); ++listIt) { + if (*listIt == it) { + list.erase(listIt); + break; + } + } + if (list.empty()) { + InFlightMessagesIndex.erase(mapIt); + Producer->SessionsWorker->AddIdleSession(partition); + } + } + + Y_ABORT_UNLESS(it->Data.size() <= MemoryUsage, "MemoryUsage is less than the size of the message"); + MemoryUsage -= it->Data.size(); + + if (it->FlushPromise.Initialized()) { + Producer->FlushPromises.push_back(std::make_pair(it->FlushPromise, TFlushResult{ + .Status = EFlushStatus::Success, + .LastWrittenSeqNo = Producer->LastWrittenSeqNo, + .ClosedDescription = std::nullopt, + })); + } + InFlightMessages.pop_front(); +} + +bool TProducer::TMessagesWorker::IsMemoryUsageOK() const { + return MemoryUsage <= Producer->Settings.MaxMemoryUsage_ / 2; +} + +void TProducer::TMessagesWorker::AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition) { + MemoryUsage += message.Data.size(); + PushInFlightMessage(partition, TMessageInfo(key, std::move(message), partition)); +} + +std::optional TProducer::TMessagesWorker::GetContinuationToken(std::uint32_t partition) { + auto it = ContinuationTokens.find(partition); + if (it != ContinuationTokens.end() && !it->second.empty()) { + auto token = std::move(it->second.front()); + it->second.pop_front(); + if (it->second.empty()) { + ContinuationTokens.erase(it); + } + return token; + } + + return std::nullopt; +} + +void TProducer::TMessagesWorker::HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken) { + auto [it, _] = ContinuationTokens.try_emplace(partition); + it->second.push_back(std::move(continuationToken)); +} + +bool TProducer::TMessagesWorker::IsQueueEmpty() const { + return InFlightMessages.empty(); +} + +const TProducer::TMessageInfo& TProducer::TMessagesWorker::GetFrontInFlightMessage() const { + Y_ABORT_UNLESS(!InFlightMessages.empty()); + return InFlightMessages.front(); +} + +bool TProducer::TMessagesWorker::HasInFlightMessages() const { + return !InFlightMessages.empty(); +} + +void TProducer::TMessagesWorker::SetClosedStatusToFlushPromises(std::optional closedDescription) { + for (auto& inFlightMessage : InFlightMessages) { + if (inFlightMessage.FlushPromise.Initialized()) { + inFlightMessage.FlushPromise.TrySetValue(TFlushResult{ + .Status = EFlushStatus::ProducerClosed, + .LastWrittenSeqNo = Producer->LastWrittenSeqNo, + .ClosedDescription = closedDescription, + }); + } + } +} + +void TProducer::TMessagesWorker::ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo) { + auto it = InFlightMessagesIndex.find(partition); + if (it == InFlightMessagesIndex.end()) { + return; + } + + auto& list = it->second; + auto resendIt = list.begin(); + auto ackQueueIt = Producer->EventsWorker->AckQueueBegin(partition); + size_t ackIdx = 0; + auto ackQueueEnd = Producer->EventsWorker->AckQueueEnd(partition); + std::vector acksToSend; + + while (resendIt != list.end()) { + if (!(*resendIt)->SeqNo.has_value() || (*resendIt)->SeqNo.value() > afterSeqNo) { + break; + } + + auto seqNo = (*resendIt)->SeqNo.value(); + if (ackQueueIt == ackQueueEnd) { + // this case can happen if the message was sent, but session was closed before the ack was received + TWriteSessionEvent::TWriteAck ack; + ack.SeqNo = seqNo; + acksToSend.push_back(std::move(ack)); + } else { + auto acksEvent = std::get_if(&*ackQueueIt); + if (ackIdx == acksEvent->Acks.size()) { + ++ackQueueIt; + ackIdx = 0; + continue; + } + + if (acksEvent->Acks[ackIdx].SeqNo > seqNo) { + // this case can happen if the message was sent, but session was closed before the ack was received + TWriteSessionEvent::TWriteAck ack; + ack.SeqNo = seqNo; + acksEvent->Acks.insert(acksEvent->Acks.begin() + ackIdx, std::move(ack)); + } + ++ackIdx; + } + ++resendIt; + } + + if (!acksToSend.empty()) { + TWriteSessionEvent::TAcksEvent event; + event.Acks = std::move(acksToSend); + Producer->EventsWorker->HandleAcksEvent(partition, std::move(event)); + } + + // IMPORTANT: do not mutate InFlightMessagesIndex while holding references/iterators to its elements. + // try_emplace()/rehash may invalidate 'it' and 'list' -> use-after-free and segfaults. + std::vector> messagesFromOldPartition; + messagesFromOldPartition.reserve(std::distance(resendIt, list.end())); + auto currentSeqNo = resendIt != list.end() ? (*resendIt)->SeqNo.value_or(0) : 0; + for (auto iter = resendIt; iter != list.end(); ++iter) { + if (iter != resendIt && currentSeqNo != 0) { + Y_ABORT_UNLESS((*iter)->SeqNo.value_or(0) > currentSeqNo, "SeqNo is not increasing for partition %d", partition); + } + + auto newPartition = Producer->PartitionChooser->ChoosePartition((*iter)->Key); + (*iter)->Partition = newPartition; + messagesFromOldPartition.emplace_back(newPartition, *iter); + + currentSeqNo = (*iter)->SeqNo.value_or(0); + } + + list.erase(resendIt, list.end()); + for (const auto& [newPartition, msgIt] : messagesFromOldPartition) { + auto [inFlightMessagesIndexChainIt, _] = InFlightMessagesIndex.try_emplace(newPartition); + inFlightMessagesIndexChainIt->second.push_back(msgIt); + + if (msgIt->Sent) { + auto [messagesToResendChainIt, __] = MessagesToResendIndex.try_emplace(newPartition); + messagesToResendChainIt->second.push_back(msgIt); + } + } + + InFlightMessagesIndex.erase(partition); + Producer->SessionsWorker->AddIdleSession(partition); +} + +void TProducer::TMessagesWorker::RebuildPendingMessagesIndex(std::uint32_t partition) { + auto [oldPendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(partition); + std::unordered_map> pendingMessagesForNewPartitions; + for (auto it = oldPendingMessagesIndexChainIt->second.begin(); it != oldPendingMessagesIndexChainIt->second.end(); ++it) { + auto newPartition = Producer->PartitionChooser->ChoosePartition((*it)->Key); + auto [pendingMessagesForNewPartitionsIt, __] = pendingMessagesForNewPartitions.try_emplace(newPartition); + pendingMessagesForNewPartitionsIt->second.push_back(*it); + } + + for (const auto& [newPartition, pendingMessagesForNewPartition] : pendingMessagesForNewPartitions) { + auto [pendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(newPartition); + for (auto reverseIt = pendingMessagesForNewPartition.rbegin(); reverseIt != pendingMessagesForNewPartition.rend(); ++reverseIt) { + pendingMessagesIndexChainIt->second.push_front(*reverseIt); + } + } + + PendingMessagesIndex.erase(partition); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::TProducerRetryPolicy + +TProducer::TProducerRetryPolicy::TProducerRetryPolicy(TProducer* producer) + : Producer(producer) +{} + +typename TProducer::TProducerRetryPolicy::IRetryState::TPtr TProducer::TProducerRetryPolicy::CreateRetryState() const { + struct TRetryState : public IRetryState { + TRetryState(TProducer* producer) + : Producer(producer) + {} + ~TRetryState() = default; + TMaybe GetNextRetryDelay(EStatus status) override { + if (status == EStatus::OVERLOADED) { + return Nothing(); + } + + if (!UserRetryState) { + auto policy = Producer->Settings.RetryPolicy_ ? Producer->Settings.RetryPolicy_ : NYdb::NTopic::IRetryPolicy::GetDefaultPolicy(); + UserRetryState = policy->CreateRetryState(); + } + + return UserRetryState->GetNextRetryDelay(status); + } + + private: + TProducer* Producer; + IRetryState::TPtr UserRetryState; + }; + + return std::make_unique(Producer); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer::Metrics + +void TProducer::TMetricGauge::Add(std::uint64_t value) { + Sum += value; + Max = std::max(Max, value); + MetricCount++; +} + +void TProducer::TMetricGauge::Clear() { + Sum = 0; + MetricCount = 0; + Max = 0; +} + +long double TProducer::TMetricGauge::Average() { + if (MetricCount == 0) { + return 0; + } + + return (long double)Sum / (long double)MetricCount; +} + +std::uint64_t TProducer::TMetricGauge::GetMax() const { + return Max; +} + +std::uint64_t TProducer::TMetricGauge::GetSum() const { + return Sum; +} + +TProducer::TMetrics::TMetrics(TProducer* producer): Producer(producer) {} + +void TProducer::TMetrics::AddMainWorkerTime(std::uint64_t ms) { + std::lock_guard lock(Lock); + MainWorkerTimeMs.Add(ms); +} + +void TProducer::TMetrics::AddCycleTime(std::uint64_t ms) { + std::lock_guard lock(Lock); + CycleTimeMs.Add(ms); +} + +void TProducer::TMetrics::AddWriteLag(std::uint64_t lagMs) { + std::lock_guard lock(Lock); + WriteLagMs.Add(lagMs); +} + +void TProducer::TMetrics::IncContinuationTokensSent() { + std::lock_guard lock(Lock); + ContinuationTokensSent.Add(1); +} + +void TProducer::TMetrics::IncBufferFull() { + std::lock_guard lock(Lock); + BufferFull.Add(1); +} + +void TProducer::TMetrics::IncIncomingMessages() { + std::lock_guard lock(Lock); + IncomingMessages.Add(1); +} + +void TProducer::TMetrics::IncOutgoingMessages() { + std::lock_guard lock(Lock); + OutgoingMessages.Add(1); +} + +void TProducer::TMetrics::PrintMetrics() { + std::lock_guard lock(Lock); + LOG_LAZY( + Producer->DbDriverState->Log, + TLOG_DEBUG, + Producer->LogPrefix() + << "METRICS: average MainWorkerTimeMs: " << MainWorkerTimeMs.Average() + << " ms, average CycleTimeMs: " << CycleTimeMs.Average() + << " ms, average WriteLagMs: " << WriteLagMs.Average() << " ms, " + << "max MainWorkerTimeMs: " << MainWorkerTimeMs.GetMax() << " ms, " + << "max CycleTimeMs: " << CycleTimeMs.GetMax() << " ms, " + << "max WriteLagMs: " << WriteLagMs.GetMax() << " ms, " + << "ContinuationTokensSent: " << ContinuationTokensSent.GetSum() << " tokens, " + << "BufferFull: " << BufferFull.GetSum() << " times, " + << "IncomingMessages: " << IncomingMessages.GetSum() << " messages, " + << "OutgoingMessages: " << OutgoingMessages.GetSum() << " messages"); + MainWorkerTimeMs.Clear(); + CycleTimeMs.Clear(); + WriteLagMs.Clear(); + ContinuationTokensSent.Clear(); + BufferFull.Clear(); + IncomingMessages.Clear(); + OutgoingMessages.Clear(); +} + + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer + +TProducer::TProducer( + const TProducerSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState) + : Connections(connections), + Client(client), + DbDriverState(dbDriverState), + Metrics(this), + Settings(settings) +{ + if (settings.ProducerIdPrefix_.empty()) { + ythrow TContractViolation("ProducerIdPrefix is required for KeyedWriteSession"); + } + + if (!settings.ProducerId_.empty()) { + ythrow TContractViolation("ProducerId should be empty for KeyedWriteSession, use ProducerIdPrefix instead"); + } + + if (!settings.MessageGroupId_.empty()) { + ythrow TContractViolation("MessageGroupId should be empty for KeyedWriteSession"); + } + + TDescribeTopicSettings describeTopicSettings; + auto topicConfig = client->DescribeTopic(settings.Path_, describeTopicSettings).GetValueSync(); + auto partitions = topicConfig.GetTopicDescription().GetPartitions(); + std::sort(partitions.begin(), partitions.end(), [](const auto& a, const auto& b) -> bool { + return a.GetPartitionId() < b.GetPartitionId(); + }); + + PartitionChooserStrategy = settings.PartitionChooserStrategy_; + auto strategy = topicConfig.GetTopicDescription().GetPartitioningSettings().GetAutoPartitioningSettings().GetStrategy(); + auto autoPartitioningEnabled = (strategy != EAutoPartitioningStrategy::Disabled && + strategy != EAutoPartitioningStrategy::Unspecified); + + for (const auto& partition : partitions) { + auto partitionId = partition.GetPartitionId(); + auto fromBound = partition.GetFromBound().value_or(""); + auto toBound = partition.GetToBound(); + LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "Adding partition " << partitionId << " from bound " << fromBound << " to bound " << (toBound.has_value() ? toBound.value() : "null")); + Partitions[partitionId] = TPartitionInfo() + .PartitionId(partitionId) + .FromBound(fromBound) + .ToBound(toBound); + } + + for (const auto& partition : partitions) { + auto children = partition.GetChildPartitionIds(); + + std::vector childrenIndices; + childrenIndices.reserve(children.size()); + for (auto child : children) { + childrenIndices.push_back(child); + } + Partitions[partition.GetPartitionId()].Children(childrenIndices); + } + + if (Settings.EventHandlers_.CommonHandler_) { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); + } else { + if (Settings.EventHandlers_.SessionClosedHandler_) { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); + } + if (Settings.EventHandlers_.AcksHandler_) { + EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); + } + } + + switch (PartitionChooserStrategy) { + case TProducerSettings::EPartitionChooserStrategy::Bound: + PartitioningKeyHasher = settings.PartitioningKeyHasher_; + PartitionChooser = std::make_unique(this); + for (size_t i = 0; i < Partitions.size(); ++i) { + if (i > 0 && Partitions[i].FromBound_.empty() && !Partitions[i].ToBound_.has_value()) { + ythrow TContractViolation("Unbounded partition is not supported for Bound partition chooser strategy"); + } + + if (!Partitions[i].Children_.empty()) { + continue; + } + + PartitionsIndex[Partitions[i].FromBound_] = Partitions[i].PartitionId_; + } + break; + case TProducerSettings::EPartitionChooserStrategy::Hash: + if (autoPartitioningEnabled) { + throw TContractViolation("Hash partition chooser strategy is not supported with auto partitioning enabled"); + } + + std::vector partitionsIds; + partitionsIds.reserve(partitions.size()); + for (const auto& partition : partitions) { + partitionsIds.push_back(partition.GetPartitionId()); + } + + PartitionChooser = std::make_unique(std::move(partitionsIds)); + break; + } + + ClosePromise = NThreading::NewPromise(); + CloseFuture = ClosePromise.GetFuture(); + ShutdownPromise = NThreading::NewPromise(); + ShutdownFuture = ShutdownPromise.GetFuture(); + + SessionsWorker = std::make_shared(this); + MessagesWorker = std::make_shared(this); + EventsWorker = std::make_shared(this); + RetryPolicy = std::make_shared(this); + + EventsWorker->AddContinuationToken(); + + // Start handlers executor for user callbacks (Acks/ReadyToAccept/SessionClosed/Common). + Settings.EventHandlers_.HandlersExecutor_->Start(); + + CloseFuture.Subscribe([this](const NThreading::TFuture&) { + RunMainWorker(-1); + }); + + RunMainWorker(-1); + + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Keyed write session created"); +} + +std::vector TProducer::GetPartitions() const { + std::vector partitions; + partitions.reserve(Partitions.size()); + for (const auto& [partitionId, partitionInfo] : Partitions) { + partitions.push_back(partitionInfo); + } + return partitions; +} + +std::unordered_map TProducer::GetPartitionsMap() const { + return Partitions; +} + +std::map TProducer::GetPartitionsIndex() const { + return PartitionsIndex; +} + +size_t TProducer::GetSessionsCount() { + std::lock_guard lock(GlobalLock); + return SessionsWorker->GetSessionsCount(); +} + +size_t TProducer::GetIdleSessionsCount() { + std::lock_guard lock(GlobalLock); + return SessionsWorker->GetIdleSessionsCount(); +} + +TCloseResult TProducer::Close(TDuration closeTimeout) { + if (Closed.exchange(true)) { + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + return TCloseResult{ + .Status = ECloseStatus::AlreadyClosed, + .ClosedDescription = sessionClosedEvent ? std::make_optional(*sessionClosedEvent) : std::nullopt + }; + } + + SetCloseDeadline(closeTimeout); + ClosePromise.TrySetValue(); + + Flush().Wait(CloseDeadline); + ShutdownFuture.Wait(CloseDeadline); + RunUserEventLoop(); + Done.store(true); + + if (MessagesWorker->IsQueueEmpty()) { + return TCloseResult{ .Status = ECloseStatus::Success }; + } + + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + if (sessionClosedEvent && sessionClosedEvent->GetStatus() != EStatus::SUCCESS) { + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + return TCloseResult{ + .Status = ECloseStatus::Error, + .ClosedDescription = sessionClosedEvent ? std::make_optional(*sessionClosedEvent) : std::nullopt + }; + } + + return TCloseResult{ .Status = ECloseStatus::Timeout }; +} + +void TProducer::NonBlockingClose() { + Closed.store(true); + Done.store(true); +} + +void TProducer::SetCloseDeadline(const TDuration& closeTimeout) { + std::lock_guard lock(GlobalLock); + CloseDeadline = TInstant::Now() + closeTimeout; +} + +TProducer::~TProducer() { + auto _ = Close(TDuration::Zero()); // Ignore the result, because we are destroying the producer + Settings.EventHandlers_.HandlersExecutor_->Stop(); + ShutdownFuture.Wait(); +} + +NThreading::TFuture TProducer::WaitEvent() { + return EventsWorker->WaitEvent(); +} + +std::optional TProducer::GetEvent(bool block) { + if (Settings.EventHandlers_.CommonHandler_) { + return std::nullopt; + } + + return EventsWorker->GetEvent(block); +} + +std::vector TProducer::GetEvents(bool block, std::optional maxEventsCount) { + if (Settings.EventHandlers_.CommonHandler_) { + return {}; + } + + return EventsWorker->GetEvents(block, maxEventsCount); +} + +TDuration TProducer::GetCloseTimeout() { + std::lock_guard lock(GlobalLock); + auto now = TInstant::Now(); + if (CloseDeadline <= now) { + return TDuration::Zero(); + } + return CloseDeadline - now; +} + +bool TProducer::RunSplittedPartitionWorkers() { + if (SplittedPartitionWorkers.empty() && ReadySplittedPartitionWorkers.empty()) { + return false; + } + + bool needRerun = false; + std::unordered_map> readySplittedPartitionWorkers; + for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { + if (splittedPartitionWorker->IsDone()) { + readySplittedPartitionWorkers[partition] = splittedPartitionWorker; + continue; + } + + splittedPartitionWorker->DoWork(); + needRerun = needRerun || splittedPartitionWorker->IsInit(); + needRerun = needRerun || splittedPartitionWorker->IsDone(); + } + + for (const auto& [partition, splittedPartitionWorker] : readySplittedPartitionWorkers) { + ReadySplittedPartitionWorkers[partition] = splittedPartitionWorker; + SplittedPartitionWorkers.erase(partition); + } + + std::vector partitionsToRemove; + for (const auto& [partition, splittedPartitionWorker] : ReadySplittedPartitionWorkers) { + if (MainWorkerOwner != partition) { + partitionsToRemove.push_back(partition); + } + } + + for (const auto& partition : partitionsToRemove) { + ReadySplittedPartitionWorkers.erase(partition); + } + + return needRerun; +} + +void TProducer::RunUserEventLoop() { + if (!Settings.EventHandlers_.AcksHandler_ && + !Settings.EventHandlers_.ReadyToAcceptHandler_ && + !Settings.EventHandlers_.SessionClosedHandler_ && + !Settings.EventHandlers_.CommonHandler_) { + return; + } + + auto handlersExecutor = Settings.EventHandlers_.HandlersExecutor_; + if (!handlersExecutor) { + return; + } + + while (true) { + auto event = EventsWorker->GetEvent(false, EventTypesWithHandlers); + if (!event) { + break; + } + + if (auto* readyToAcceptEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.ReadyToAcceptHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*readyToAcceptEvent)]() mutable { + Settings.EventHandlers_.ReadyToAcceptHandler_(ev); + }); + } else if (Settings.EventHandlers_.CommonHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + continue; + } + + if (auto* acksEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.AcksHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*acksEvent)]() mutable { + Settings.EventHandlers_.AcksHandler_(ev); + }); + } else if (Settings.EventHandlers_.CommonHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + continue; + } + + if (auto* sessionClosedEvent = std::get_if(&*event)) { + if (Settings.EventHandlers_.SessionClosedHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*sessionClosedEvent)]() mutable { + Settings.EventHandlers_.SessionClosedHandler_(ev); + }); + } else if (Settings.EventHandlers_.CommonHandler_) { + handlersExecutor->Post( + [this, ev = std::move(*event)]() mutable { + Settings.EventHandlers_.CommonHandler_(ev); + }); + } + break; + } + } +} + +void TProducer::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent) { + std::optional receivedSessionClosedEvent; + while (true) { + auto event = wrappedSession->Session->GetEvent(false); + if (!event) { + break; + } + + if (auto* closedEvent = std::get_if(&*event)) { + receivedSessionClosedEvent = std::move(*closedEvent); + break; + } + } + + if (!receivedSessionClosedEvent || receivedSessionClosedEvent->GetStatus() == EStatus::SUCCESS || receivedSessionClosedEvent->GetStatus() == EStatus::OVERLOADED) { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Failed to get session closed event"); + EventsWorker->HandleSessionClosedEvent(std::move(*sessionClosedEvent), wrappedSession->Partition); + } else { + EventsWorker->HandleSessionClosedEvent(std::move(*receivedSessionClosedEvent), wrappedSession->Partition); + } +} + +TStringBuilder TProducer::LogPrefix() { + return TStringBuilder() << " SessionId: " << Settings.SessionId_ << " Epoch: " << Epoch.load() << " "; +} + +void TProducer::NextEpoch() { + auto maxEpoch = MAX_EPOCH - 1; + if (Epoch.compare_exchange_weak(maxEpoch, 0)) { + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Epoch overflow, resetting to 0"); + return; + } + + Epoch.fetch_add(1); +} + +void TProducer::RunMainWorker(std::int64_t owner) { + // This function is both "request to run" and the runner itself. + // We must handle two properties: + // - TFuture::Subscribe may call back synchronously when future is already ready. + // - A callback may race with the runner trying to go idle (avoid lost wakeups). + enum : std::uint8_t { + Idle = 0, + Running = 1, + Rerun = 2, + }; + + // Try to become the runner. If already running, just request a rerun. + std::uint8_t state = MainWorkerState.load(std::memory_order_acquire); + for (;;) { + if (state & Running) { + if (MainWorkerState.compare_exchange_weak(state, std::uint8_t(state | Rerun), + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return; + } + continue; + } else { + if (MainWorkerState.compare_exchange_weak(state, Running, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + break; // we are the runner now + } + continue; + } + } + + MainWorkerOwner = owner; + NextEpoch(); + + auto startWorkerTime = TInstant::Now(); + // Runner loop: process, arm subscription, then either go idle or loop again. + for (;;) { + auto startIter = TInstant::Now(); + // Clear rerun request for this iteration. + MainWorkerState.fetch_and(std::uint8_t(~Rerun), std::memory_order_acq_rel); + bool needRerun = false; + std::optional> eventsPromise; + + { + std::unique_lock lock(GlobalLock); + eventsPromise = EventsWorker->DoWork(); + RunUserEventLoop(); + needRerun = RunSplittedPartitionWorkers(); + if (!Done.load()) { + SessionsWorker->DoWork(); + MessagesWorker->DoWork(); + } + } + + if (eventsPromise) { + eventsPromise->TrySetValue(); + } + + while (!FlushPromises.empty()) { + auto& [promise, flushResult] = FlushPromises.front(); + promise.TrySetValue(flushResult); + FlushPromises.pop_front(); + } + + const auto isClosed = Closed.load(); + const auto closeTimeout = GetCloseTimeout(); + if (isClosed && (Done.load() || MessagesWorker->IsQueueEmpty() || closeTimeout == TDuration::Zero())) { + EventsWorker->EventsPromise.TrySetValue(); + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + MessagesWorker->SetClosedStatusToFlushPromises( + sessionClosedEvent ? + std::make_optional(TCloseDescription(*sessionClosedEvent)) : + std::nullopt); + MainWorkerState.store(Idle, std::memory_order_release); + ShutdownPromise.TrySetValue(); + return; + } + + if (needRerun) { + // we need this case to start resending messages if there are any + Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); + continue; + } + + // Try to go idle. If someone requested rerun concurrently, keep running. + std::uint8_t cur = MainWorkerState.load(std::memory_order_acquire); + for (;;) { + if (cur & Rerun) { + Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); + break; // continue outer loop + } + if (MainWorkerState.compare_exchange_weak(cur, Idle, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + auto workerFinished = TInstant::Now(); + Metrics.AddCycleTime((workerFinished - startIter).MilliSeconds()); + Metrics.AddMainWorkerTime((workerFinished - startWorkerTime).MilliSeconds()); + return; // successfully went idle + } + } + // Rerun was requested; continue the loop without recursion. + } +} + +TWriteResult TProducer::WriteInternal(TContinuationToken&&, TWriteMessage&& message) { + std::optional> eventsPromise; + { + std::lock_guard lock(GlobalLock); + Metrics.IncIncomingMessages(); + if (Closed.load()) { + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + return TWriteResult{ + .Status = EWriteStatus::Error, + .ErrorMessage = "producer is closed", + .ClosedDescription = sessionClosedEvent ? std::make_optional(TCloseDescription(*sessionClosedEvent)) : std::nullopt, + }; + } + + if ((message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithoutSeqNo) + || (!message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithSeqNo)) { + ythrow TContractViolation("Can not mix messages with and without seqNo"); + } + + if (SeqNoStrategy == ESeqNoStrategy::NotInitialized) { + SeqNoStrategy = message.SeqNo_.has_value() ? ESeqNoStrategy::WithSeqNo : ESeqNoStrategy::WithoutSeqNo; + } + + std::uint32_t chosenPartition; + if (message.Partition_.has_value()) { + if (!Partitions[message.Partition_.value()].Children_.empty()) { + return TWriteResult{ + .Status = EWriteStatus::Error, + .ErrorMessage = "Partition was split", + }; + } + + chosenPartition = message.Partition_.value(); + } else if (!message.Key_.has_value()) { + std::string key; + if (Settings.KeyProducer_) { + key = (*Settings.KeyProducer_)(message); + } else { + key = Settings.ProducerIdPrefix_; + } + message.Key(key); + chosenPartition = PartitionChooser->ChoosePartition(key); + } else { + chosenPartition = PartitionChooser->ChoosePartition(*message.Key_); + } + + MessagesWorker->AddMessage(message.Key_.value_or(""), std::move(message), chosenPartition); + eventsPromise = EventsWorker->HandleNewMessage(); + RunUserEventLoop(); + } + + RunMainWorker(-1); + if (eventsPromise) { + eventsPromise->TrySetValue(); + } + + return TWriteResult{ + .Status = EWriteStatus::Queued, + }; +} + +TWriteResult TProducer::Write(TWriteMessage&& message) { + auto remainingTimeout = Settings.MaxBlock_; + auto sleepTimeMs = DEFAULT_START_BLOCK_TIMEOUT; + for (;;) { + if (Closed.load()) { + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + return TWriteResult{ + .Status = EWriteStatus::Error, + .ErrorMessage = "producer is closed", + .ClosedDescription = sessionClosedEvent ? std::make_optional(TCloseDescription(*sessionClosedEvent)) : std::nullopt, + }; + } + + auto continuationToken = EventsWorker->GetContinuationToken(); + if (!continuationToken) { + if (remainingTimeout > TDuration::Zero()) { + auto toSleep = Min(sleepTimeMs, remainingTimeout); + Sleep(toSleep); + sleepTimeMs *= 2; + if (remainingTimeout > toSleep) { + remainingTimeout -= toSleep; + continue; + } + + return TWriteResult{ + .Status = EWriteStatus::Timeout, + }; + } + + return TWriteResult{ + .Status = EWriteStatus::Timeout, + }; + } + + return WriteInternal(std::move(*continuationToken), std::move(message)); + } +} + +void TProducer::Write(TContinuationToken&& continuationToken, TWriteMessage&& message) { + WriteInternal(std::move(continuationToken), std::move(message)); +} + +TWriteStats TProducer::GetWriteStats() { + std::lock_guard lock(GlobalLock); + return TWriteStats{ + .LastWrittenSeqNo = LastWrittenSeqNo, + .MessagesWritten = MessagesWritten, + }; +} + +NThreading::TFuture TProducer::Flush() { + std::unique_lock lock(GlobalLock); + if (Closed.load() || MessagesWorker->InFlightMessages.empty()) { + auto sessionClosedEvent = EventsWorker->GetSessionClosedEvent(); + return NThreading::MakeFuture(TFlushResult{ + .Status = EFlushStatus::Success, + .LastWrittenSeqNo = LastWrittenSeqNo, + .ClosedDescription = sessionClosedEvent ? std::make_optional(TCloseDescription(*sessionClosedEvent)) : std::nullopt, + }); + } + + auto lastInFlightMessage = std::prev(MessagesWorker->InFlightMessages.end()); + if (!lastInFlightMessage->FlushPromise.Initialized()) { + lastInFlightMessage->FlushPromise = NThreading::NewPromise(); + } + + return lastInFlightMessage->FlushPromise.GetFuture(); +} + +TInstant TProducer::GetCloseDeadline() { + std::lock_guard lock(GlobalLock); + return CloseDeadline; +} + +void TProducer::HandleAutoPartitioning(std::uint32_t partition) { + LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "HandleAutoPartitioning: " << partition); + auto splittedPartitionWorker = std::make_shared(this, partition); + SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); +} + +std::string TProducer::GetProducerId(std::uint32_t partition) { + return std::format("{}_{}", Settings.ProducerIdPrefix_, partition); +} + +TWriterCounters::TPtr TProducer::GetCounters() { + return nullptr; +} + +TProducer::TBoundPartitionChooser::TBoundPartitionChooser(TProducer* producer) + : Producer(producer) +{} + +std::uint32_t TProducer::TBoundPartitionChooser::ChoosePartition(const std::string_view key) { + auto hashedKey = Producer->PartitioningKeyHasher(key); + + auto lowerBound = Producer->PartitionsIndex.lower_bound(hashedKey); + if (lowerBound != Producer->PartitionsIndex.end() && lowerBound->first == hashedKey) { + return lowerBound->second; + } + + Y_ABORT_IF(lowerBound == Producer->PartitionsIndex.begin(), "Lower bound is the first element"); + return std::prev(lowerBound)->second; +} + +TProducer::THashPartitionChooser::THashPartitionChooser(std::vector&& partitions) + : Partitions(std::move(partitions)) +{} + +std::uint32_t TProducer::THashPartitionChooser::ChoosePartition(const std::string_view key) { + auto hash = MurmurHash(key.data(), key.size()); + return Partitions[hash % Partitions.size()]; +} + +} // namespace NYdb::inline V3::NTopic diff --git a/src/client/topic/impl/producer.h b/src/client/topic/impl/producer.h new file mode 100644 index 00000000000..ff0c997132a --- /dev/null +++ b/src/client/topic/impl/producer.h @@ -0,0 +1,449 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace NYdb::inline V3::NTopic { + +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TProducer + +class TProducer : public IProducer, + public TContinuationTokenIssuer, + public std::enable_shared_from_this { +private: + static constexpr size_t MAX_EPOCH = 1'000'000'000; + static constexpr TDuration DEFAULT_START_BLOCK_TIMEOUT = TDuration::MilliSeconds(1); + + using WriteSessionPtr = std::shared_ptr; + + struct TPartitionInfo { + using TSelf = TPartitionInfo; + + bool InRange(const std::string_view key) const; + bool IsSplitted() const; + + FLUENT_SETTING(std::string, FromBound); + FLUENT_SETTING(std::optional, ToBound); + FLUENT_SETTING(std::uint32_t, PartitionId); + FLUENT_SETTING(std::vector, Children); + FLUENT_SETTING_DEFAULT(bool, Locked, false); + FLUENT_SETTING_DEFAULT(NThreading::TFuture, Future, NThreading::MakeFuture()); + }; + + struct TMessageInfo { + TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition); + + std::string Key; + std::string Data; + std::optional Codec; + uint32_t OriginalSize = 0; + std::optional SeqNo; + std::optional CreateTimestamp; + TMessageMeta MessageMeta; + std::optional> Tx; + std::uint32_t Partition; + bool Sent = false; + NThreading::TPromise FlushPromise; + + TWriteMessage BuildMessage() const; + }; + + struct TIdleSession; + + struct TWriteSessionWrapper { + WriteSessionPtr Session; + const std::uint32_t Partition; + std::shared_ptr IdleSession = nullptr; + + TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition); + }; + + using WrappedWriteSessionPtr = std::shared_ptr; + + struct TIdleSession { + TIdleSession(TWriteSessionWrapper* session, TInstant emptySince, TDuration idleTimeout) + : Session(session) + , EmptySince(emptySince) + , IdleTimeout(idleTimeout) + {} + + const TWriteSessionWrapper* Session; + const TInstant EmptySince; + const TDuration IdleTimeout; + + bool Less(const std::shared_ptr& other) const; + bool IsExpired() const; + + struct Comparator { + bool operator()(const std::shared_ptr& first, const std::shared_ptr& second) const; + }; + }; + + using IdleSessionPtr = std::shared_ptr; + + enum class ESeqNoStrategy { + NotInitialized, + WithoutSeqNo, + WithSeqNo, + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Custom retry policy + + struct TProducerRetryPolicy : public ::IRetryPolicy { + using TSelf = TProducerRetryPolicy; + using TPtr = std::shared_ptr; + + TProducerRetryPolicy(TProducer* producer); + ~TProducerRetryPolicy() = default; + typename IRetryState::TPtr CreateRetryState() const override; + + private: + TProducer* Producer; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Workers + + struct TEventsWorker; + + struct TSessionsWorker { + TSessionsWorker(TProducer* producer); + WrappedWriteSessionPtr GetWriteSession(std::uint32_t partition, bool directToPartition = true); + void AddIdleSession(std::uint32_t partition); + void RemoveIdleSession(std::uint32_t partition); + void DoWork(); + size_t GetSessionsCount() const; + size_t GetIdleSessionsCount() const; + + private: + WrappedWriteSessionPtr CreateWriteSession(std::uint32_t partition, bool directToPartition = true); + + using TSessionsIndexIterator = std::unordered_map::iterator; + void DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout); + + std::string GetProducerId(std::uint32_t partitionId); + + TProducer* Producer; + std::set IdlerSessions; + using IdlerSessionsIterator = std::set::iterator; + std::unordered_map IdlerSessionsIndex; + std::unordered_map SessionsIndex; + std::deque SessionsToRemove; + + static constexpr TDuration SESSION_REMOVE_DELAY = TDuration::Seconds(5); + }; + + struct TMessagesWorker { + TMessagesWorker(TProducer* producer); + + void DoWork(); + + void AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition); + void ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo); + void RebuildPendingMessagesIndex(std::uint32_t partition); + void HandleAck(); + void HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken); + bool IsMemoryUsageOK() const; + bool IsQueueEmpty() const; + bool HasInFlightMessages() const; + const TMessageInfo& GetFrontInFlightMessage() const; + void SetClosedStatusToFlushPromises(std::optional closedDescription); + + private: + using MessageIter = std::list::iterator; + + void PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message); + void PopInFlightMessage(); + bool SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message); + std::optional GetContinuationToken(std::uint32_t partition); + void RechoosePartitionIfNeeded(MessageIter message); + + TProducer* Producer; + + std::list InFlightMessages; + std::unordered_map> InFlightMessagesIndex; + std::unordered_map> PendingMessagesIndex; + std::unordered_map> MessagesToResendIndex; + std::unordered_map> ContinuationTokens; + + std::uint64_t MemoryUsage = 0; + + friend class TProducer; + }; + + struct TSplittedPartitionWorker : public std::enable_shared_from_this { + private: + enum class EState { + Init = 0, + PendingDescribe = 1, + GotDescribe = 2, + PendingMaxSeqNo = 3, + GotMaxSeqNo = 4, + Done = 5, + }; + + void MoveTo(EState state); + void UpdateMaxSeqNo(uint64_t maxSeqNo); + void LaunchGetMaxSeqNoFutures(std::unique_lock& lock); + void HandleDescribeResult(); + + public: + TSplittedPartitionWorker(TProducer* producer, std::uint32_t partitionId); + void DoWork(); + bool IsDone(); + bool IsInit(); + std::string GetStateName() const; + + private: + TProducer* Producer; + NThreading::TFuture DescribeTopicFuture; + EState State = EState::Init; + std::uint32_t PartitionId; + std::uint64_t MaxSeqNo = 0; + std::vector WriteSessions; + std::vector> GetMaxSeqNoFutures; + std::mutex Lock; + std::uint64_t NotReadyFutures = 0; + size_t Retries = 0; + TInstant DoneAt = TInstant::Max(); + }; + + struct TEventsWorker : public std::enable_shared_from_this { + enum class EEventType { + SessionClosed = 0, + ReadyToAccept = 1, + Ack = 2, + }; + + TEventsWorker(TProducer* producer); + + std::optional> DoWork(); + NThreading::TFuture WaitEvent(); + void UnsubscribeFromPartition(std::uint32_t partition); + void SubscribeToPartition(std::uint32_t partition); + std::optional> HandleNewMessage(); + void HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event); + std::optional GetEvent(bool block, const std::vector& eventTypes = {}); + std::vector GetEvents(bool block, std::optional maxEventsCount = std::nullopt, const std::vector& eventTypes = {}); + std::list::iterator AckQueueBegin(std::uint32_t partition); + std::list::iterator AckQueueEnd(std::uint32_t partition); + std::optional GetContinuationToken(); + std::optional GetSessionClosedEvent(); + + private: + void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition); + void HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); + bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition); + bool TransferEventsToOutputQueue(); + void AddContinuationToken(); + bool AddSessionClosedIfNeeded(); + std::optional GetEventImpl(bool block, const std::vector& eventTypes = {}); + EEventType GetEventType(const TWriteSessionEvent::TEvent& event); + + TProducer* Producer; + + std::unordered_set ReadyFutures; + std::unordered_map> PartitionsEventQueues; + std::list EventsOutputQueue; + std::list TokensQueue; + + using EventsIter = std::list::iterator; + + std::mutex Lock; + + NThreading::TPromise EventsPromise; + NThreading::TFuture EventsFuture; + + std::optional CloseEvent; + + friend class TProducer; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + // Partition chooser + + struct IPartitionChooser { + virtual std::uint32_t ChoosePartition(const std::string_view key) = 0; + virtual ~IPartitionChooser() = default; + }; + + struct TBoundPartitionChooser : IPartitionChooser { + TBoundPartitionChooser(TProducer* producer); + std::uint32_t ChoosePartition(const std::string_view key) override; + private: + TProducer* Producer; + }; + + struct THashPartitionChooser : IPartitionChooser { + THashPartitionChooser(std::vector&& partitions); + std::uint32_t ChoosePartition(const std::string_view key) override; + private: + std::vector Partitions; + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + + struct TMetricGauge { + std::uint64_t MetricCount = 0; + std::uint64_t Sum = 0; + std::uint64_t Max = 0; + + long double Average(); + void Add(std::uint64_t value); + std::uint64_t GetMax() const; + std::uint64_t GetSum() const; + void Clear(); + }; + + struct TMetrics { + TMetrics(TProducer* producer); + + TMetricGauge MainWorkerTimeMs; + TMetricGauge CycleTimeMs; + TMetricGauge WriteLagMs; + TMetricGauge ContinuationTokensSent; + TMetricGauge BufferFull; + TMetricGauge IncomingMessages; + TMetricGauge OutgoingMessages; + std::mutex Lock; + TProducer* Producer; + + void AddMainWorkerTime(std::uint64_t ms); + void AddCycleTime(std::uint64_t ms); + void AddWriteLag(std::uint64_t lagMs); + void IncContinuationTokensSent(); + void IncBufferFull(); + void IncIncomingMessages(); + void IncOutgoingMessages(); + void PrintMetrics(); + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + + void RunMainWorker(std::int64_t owner); + + void NonBlockingClose(); + + void SetCloseDeadline(const TDuration& closeTimeout); + + TDuration GetCloseTimeout(); + + std::string GetProducerId(std::uint32_t partition); + + void HandleAutoPartitioning(std::uint32_t partition); + + bool RunSplittedPartitionWorkers(); + + void RunUserEventLoop(); + + TInstant GetCloseDeadline(); + + void GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent = std::nullopt); + + TStringBuilder LogPrefix(); + + void NextEpoch(); + + TWriteResult WriteInternal(TContinuationToken&&, TWriteMessage&& message); + +public: + TProducer(const TProducerSettings& settings, + std::shared_ptr client, + std::shared_ptr connections, + TDbDriverStatePtr dbDriverState); + + void Write(TContinuationToken&& continuationToken, TWriteMessage&& message); + + [[nodiscard]] TWriteResult Write(TWriteMessage&& message) override; + + [[nodiscard]] NThreading::TFuture Flush() override; + + TWriteStats GetWriteStats() override; + + NThreading::TFuture WaitEvent(); + + std::optional GetEvent(bool block = false); + + std::vector GetEvents(bool block = false, std::optional maxEventsCount = std::nullopt); + + TCloseResult Close(TDuration closeTimeout = TDuration::Max()) override; + + TWriterCounters::TPtr GetCounters(); + + std::vector GetPartitions() const; + + std::unordered_map GetPartitionsMap() const; + + std::map GetPartitionsIndex() const; + + size_t GetSessionsCount(); + + size_t GetIdleSessionsCount(); + + ~TProducer(); + +private: + std::shared_ptr Connections; + std::shared_ptr Client; + TDbDriverStatePtr DbDriverState; + + TMetrics Metrics; + + std::unordered_map Partitions; + std::map PartitionsIndex; + + TProducerSettings Settings; + ESeqNoStrategy SeqNoStrategy = ESeqNoStrategy::NotInitialized; + TProducerSettings::EPartitionChooserStrategy PartitionChooserStrategy = TProducerSettings::EPartitionChooserStrategy::Hash; + + NThreading::TPromise ClosePromise; + NThreading::TFuture CloseFuture; + NThreading::TPromise ShutdownPromise; + NThreading::TFuture ShutdownFuture; + + std::mutex GlobalLock; + std::atomic_bool Closed = false; + std::atomic_bool Done = false; + TInstant CloseDeadline = TInstant::Now(); + + std::unique_ptr PartitionChooser; + + std::function PartitioningKeyHasher; + + std::shared_ptr EventsWorker; + std::shared_ptr SessionsWorker; + std::unordered_map> SplittedPartitionWorkers; + std::unordered_map> ReadySplittedPartitionWorkers; + std::shared_ptr MessagesWorker; + std::shared_ptr RetryPolicy; + + // TFuture::Subscribe may invoke callback synchronously when the future is already ready. + // Also, callbacks may arrive concurrently with the attempt to go idle. + // Use a small state machine to avoid re-entrancy and lost wakeups. + std::atomic MainWorkerState = 0; + // MainWorker has an owner, which can be: + // - user's thread, in this case the value of MainWorkerOwner is -1 + // - subsession's thread, in this case the value of MainWorkerOwner is the subsession's partition ID + std::int64_t MainWorkerOwner = -1; + + std::uint64_t LastWrittenSeqNo = 0; + std::uint64_t MessagesWritten = 0; + + std::atomic Epoch = 0; + + std::list, TFlushResult>> FlushPromises; + + std::vector EventTypesWithHandlers; +}; + +} // namespace NYdb::inline V3::NTopic diff --git a/src/client/topic/impl/topic.cpp b/src/client/topic/impl/topic.cpp index cf8c77e0aac..429f61a5d49 100644 --- a/src/client/topic/impl/topic.cpp +++ b/src/client/topic/impl/topic.cpp @@ -591,14 +591,8 @@ std::shared_ptr TTopicClient::CreateSimpleBlockingW return Impl_->CreateSimpleWriteSession(settings); } -std::shared_ptr TTopicClient::CreateSimpleBlockingKeyedWriteSession( - const TKeyedWriteSessionSettings& settings) { - return Impl_->CreateSimpleKeyedWriteSession(settings); -} - -std::shared_ptr TTopicClient::CreateKeyedWriteSession( - const TKeyedWriteSessionSettings& settings) { - return Impl_->CreateKeyedWriteSession(settings); +std::shared_ptr TTopicClient::CreateProducer(const TProducerSettings& settings) { + return Impl_->CreateProducer(settings); } std::shared_ptr TTopicClient::CreateWriteSession(const TWriteSessionSettings& settings) { diff --git a/src/client/topic/impl/topic_impl.cpp b/src/client/topic/impl/topic_impl.cpp index acfe6bde587..236629cfea5 100644 --- a/src/client/topic/impl/topic_impl.cpp +++ b/src/client/topic/impl/topic_impl.cpp @@ -2,6 +2,7 @@ #include "read_session.h" #include "write_session.h" +#include "producer.h" namespace NYdb::inline V3::NTopic { @@ -61,7 +62,7 @@ std::shared_ptr TTopicClient::TImpl::CreateSimpleWr return std::move(session); } -std::shared_ptr TTopicClient::TImpl::CreateSimpleKeyedWriteSession(const TKeyedWriteSessionSettings& settings) { +std::shared_ptr TTopicClient::TImpl::CreateProducer(const TProducerSettings& settings) { auto alteredSettings = settings; { std::lock_guard guard(Lock); @@ -69,31 +70,27 @@ std::shared_ptr TTopicClient::TImpl::CreateSim alteredSettings.CompressionExecutor(Settings.DefaultCompressionExecutor_); } + bool handlersSet = settings.EventHandlers_.AcksHandler_ || + settings.EventHandlers_.SessionClosedHandler_ || + settings.EventHandlers_.CommonHandler_; + if (!settings.EventHandlers_.HandlersExecutor_) { - alteredSettings.EventHandlers_.HandlersExecutor(Settings.DefaultHandlersExecutor_); + if (handlersSet) { + alteredSettings.EventHandlers_.HandlersExecutor(Settings.DefaultHandlersExecutor_); + } else { + alteredSettings.EventHandlers_.HandlersExecutor(NTopic::CreateSyncExecutor()); + } } - } - auto session = std::make_shared( - alteredSettings, shared_from_this(), Connections_, DbDriverState_ - ); - return session; -} + // As we don't support continuation tokens in IProducer interface + alteredSettings.EventHandlers_.ReadyToAcceptHandler({}); -std::shared_ptr TTopicClient::TImpl::CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings) { - auto alteredSettings = settings; - { - std::lock_guard guard(Lock); - if (!settings.CompressionExecutor_) { - alteredSettings.CompressionExecutor(Settings.DefaultCompressionExecutor_); - } - - if (!settings.EventHandlers_.HandlersExecutor_) { - alteredSettings.EventHandlers_.HandlersExecutor(Settings.DefaultHandlersExecutor_); + if (!settings.EventHandlers_.AcksHandler_) { + alteredSettings.EventHandlers_.AcksHandler([&](TWriteSessionEvent::TAcksEvent&) {}); } } - return std::make_shared( + return std::make_shared( alteredSettings, shared_from_this(), Connections_, DbDriverState_ ); } diff --git a/src/client/topic/impl/topic_impl.h b/src/client/topic/impl/topic_impl.h index 669ae12c39f..7a3a926892d 100644 --- a/src/client/topic/impl/topic_impl.h +++ b/src/client/topic/impl/topic_impl.h @@ -2,6 +2,7 @@ #include "transaction.h" +#include #include #define INCLUDE_YDB_INTERNAL_H @@ -316,8 +317,14 @@ class TTopicClient::TImpl : public TClientImplCommon { // Runtime API. std::shared_ptr CreateReadSession(const TReadSessionSettings& settings); std::shared_ptr CreateSimpleWriteSession(const TWriteSessionSettings& settings); - std::shared_ptr CreateSimpleKeyedWriteSession(const TKeyedWriteSessionSettings& settings); - std::shared_ptr CreateKeyedWriteSession(const TKeyedWriteSessionSettings& settings); + std::shared_ptr CreateProducer(const TProducerSettings& settings); + + template + std::shared_ptr> CreateTypedProducer(const TProducerSettings& settings) { + auto producer = CreateProducer(settings); + return std::make_shared>(producer); + } + std::shared_ptr CreateWriteSession(const TWriteSessionSettings& settings); using IReadSessionConnectionProcessorFactory = diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 51ecaa233b9..bb540a94bf3 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -96,1699 +96,6 @@ TWriteSession::~TWriteSession() { TryGetImpl()->Close(TDuration::Zero()); } -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession - -// TKeyedWriteSessionSettings - -std::string TKeyedWriteSessionSettings::DefaultPartitioningKeyHasher(const std::string_view key) { - const std::uint64_t lo = MurmurHash(key.data(), key.size(), std::uint64_t{0}); - const std::uint64_t hi = MurmurHash(key.data(), key.size(), std::uint64_t{0x9E3779B97F4A7C15ull}); // fixed seed - - const std::uint64_t hiBe = InetToHost(hi); - const std::uint64_t loBe = InetToHost(lo); - - std::string out; - out.resize(16); - memcpy(out.data() + 0, &hiBe, 8); - memcpy(out.data() + 8, &loBe, 8); - return out; // 16 bytes -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TPartitionInfo - -bool TKeyedWriteSession::TPartitionInfo::InRange(const std::string_view key) const { - if (FromBound_ > key) { - return false; - } - if (ToBound_.has_value() && *ToBound_ <= key) { - return false; - } - return true; -} - -bool TKeyedWriteSession::TPartitionInfo::IsSplitted() const { - return !Children_.empty(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TMessageInfo - -TKeyedWriteSession::TMessageInfo::TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx) - : Key(key) - , Data(message.Data) - , Codec(message.Codec) - , OriginalSize(message.OriginalSize) - , SeqNo(message.SeqNo_) - , CreateTimestamp(message.CreateTimestamp_) - , TxInMessage(message.Tx_) - , Tx(tx) - , Partition(partition) -{ - for (const auto& [key, value] : message.MessageMeta_) { - MessageMeta.Fields.emplace_back(key, value); - } -} - -TWriteMessage TKeyedWriteSession::TMessageInfo::BuildMessage() const { - TWriteMessage message(Data); - message.Codec = Codec; - message.OriginalSize = OriginalSize; - message.SeqNo(SeqNo); - message.CreateTimestamp(CreateTimestamp); - for (const auto& [key, value] : MessageMeta.Fields) { - message.MessageMeta_.emplace_back(key, value); - } - message.Tx(TxInMessage); - return message; -} -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TWriteSessionWrapper - -TKeyedWriteSession::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition) - : Session(std::move(session)) - , Partition(partition) - , QueueSize(0) -{} - -bool TKeyedWriteSession::TWriteSessionWrapper::IsQueueEmpty() const { - return QueueSize == 0; -} - -bool TKeyedWriteSession::TWriteSessionWrapper::AddToQueue(std::uint64_t delta) { - bool idle = QueueSize == 0; - QueueSize += delta; - return idle; -} - -bool TKeyedWriteSession::TWriteSessionWrapper::RemoveFromQueue(std::uint64_t delta) { - Y_ABORT_UNLESS(QueueSize >= delta, "RemoveFromQueue: underflow"); - QueueSize -= delta; - return QueueSize == 0; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TIdleSession - -bool TKeyedWriteSession::TIdleSession::Less(const std::shared_ptr& other) const { - if (EmptySince == other->EmptySince) { - return Session->Partition < other->Session->Partition; - } - - return EmptySince < other->EmptySince; -} - -bool TKeyedWriteSession::TIdleSession::Comparator::operator()( - const std::shared_ptr& first, - const std::shared_ptr& second) const { - return first->Less(second); -} - -bool TKeyedWriteSession::TIdleSession::IsExpired() const { - return TInstant::Now() - EmptySince > IdleTimeout; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TSplittedPartitionWorker - -TKeyedWriteSession::TSplittedPartitionWorker::TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId) - : Session(session) - , PartitionId(partitionId) -{ - LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "Creating splitted partition worker for partition " << PartitionId); -} - -std::string TKeyedWriteSession::TSplittedPartitionWorker::GetStateName() const { - switch (State) { - case EState::Init: - return "Init"; - case EState::PendingDescribe: - return "PendingDescribe"; - case EState::GotDescribe: - return "GotDescribe"; - case EState::PendingMaxSeqNo: - return "PendingMaxSeqNo"; - case EState::Done: - return "Done"; - case EState::GotMaxSeqNo: - return "GotMaxSeqNo"; - } -} - -void TKeyedWriteSession::TSplittedPartitionWorker::DoWork() { - std::unique_lock lock(Lock); - std::weak_ptr session = Session->shared_from_this(); - switch (State) { - case EState::Init: - DescribeTopicFuture = Session->Client->DescribeTopic(Session->Settings.Path_, TDescribeTopicSettings()); - lock.unlock(); - DescribeTopicFuture.Subscribe([this, session](const NThreading::TFuture&) { - auto sessionPtr = session.lock(); - if (!sessionPtr) { - return; - } - - { - std::lock_guard lock(Lock); - MoveTo(EState::GotDescribe); - } - - sessionPtr->RunMainWorker(); - }); - lock.lock(); - if (State == EState::Init) { - MoveTo(EState::PendingDescribe); - } - break; - case EState::GotDescribe: - HandleDescribeResult(); - if (State != EState::GotDescribe) { - break; - } - - LaunchGetMaxSeqNoFutures(lock); - if (State == EState::GotDescribe) { - MoveTo(EState::PendingMaxSeqNo); - } - break; - case EState::PendingDescribe: - case EState::PendingMaxSeqNo: - case EState::Done: - break; - case EState::GotMaxSeqNo: - Session->MessagesWorker->RebuildPendingMessagesIndex(PartitionId); - Session->MessagesWorker->ScheduleResendMessages(PartitionId, MaxSeqNo); - for (const auto& child : Session->Partitions[PartitionId].Children_) { - Session->Partitions[child].Locked(false); - } - Session->Partitions[PartitionId].Locked_ = false; - MoveTo(EState::Done); - break; - } -} - -void TKeyedWriteSession::TSplittedPartitionWorker::MoveTo(EState state) { - State = state; - LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "Moving splitted partition worker for partition " << PartitionId << " to state " << GetStateName()); -} - -void TKeyedWriteSession::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t maxSeqNo) { - MaxSeqNo = std::max(MaxSeqNo, maxSeqNo); -} - -bool TKeyedWriteSession::TSplittedPartitionWorker::IsDone() { - std::lock_guard lock(Lock); - DoneAt = TInstant::Now(); - return State == EState::Done; -} - -bool TKeyedWriteSession::TSplittedPartitionWorker::CanBeRemoved() { - std::lock_guard lock(Lock); - if (State != EState::Done) { - return false; - } - - return TInstant::Now() - DoneAt > TDuration::Seconds(10); -} - -bool TKeyedWriteSession::TSplittedPartitionWorker::IsInit() { - std::lock_guard lock(Lock); - return State == EState::Init; -} - -void TKeyedWriteSession::TSplittedPartitionWorker::HandleDescribeResult() { - std::vector newPartitionsIds; - const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); - for (const auto& partition : partitions) { - if (partition.GetPartitionId() != PartitionId) { - continue; - } - - LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Found partition " << partition.GetPartitionId() << " for partition " << PartitionId << " children: " << partition.GetChildPartitionIds().size()); - for (const auto& childPartitionId : partition.GetChildPartitionIds()) { - newPartitionsIds.push_back(childPartitionId); - } - break; - } - - if (newPartitionsIds.empty()) { - // describe response is incomplete, we need to resend describe request - MoveTo(EState::Init); - Y_ABORT_UNLESS(++Retries < 40, "Too many retries for partition %u", PartitionId); - LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Describe response is incomplete, we need to resend describe request for partition " << PartitionId); - return; - } - - std::vector children; - const auto& splittedPartition = Session->Partitions[PartitionId]; - Session->PartitionsIndex.erase(splittedPartition.FromBound_); - - for (const auto& newPartitionId : newPartitionsIds) { - auto partitionDescribeInfo = std::find_if(partitions.begin(), partitions.end(), [newPartitionId](const auto& partition) { - return partition.GetPartitionId() == newPartitionId; - }); - Y_ABORT_UNLESS(partitionDescribeInfo != partitions.end(), "Partition describe info not found"); - Session->PartitionsIndex[partitionDescribeInfo->GetFromBound().value_or("")] = newPartitionId; - Session->Partitions[newPartitionId] = TPartitionInfo() - .PartitionId(newPartitionId) - .FromBound(partitionDescribeInfo->GetFromBound().value_or("")) - .ToBound(partitionDescribeInfo->GetToBound()) - .Locked(true); - children.push_back(newPartitionId); - } - - Session->Partitions[PartitionId].Children(children); -} - -void TKeyedWriteSession::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_lock& lock) { - Y_ABORT_UNLESS(DescribeTopicFuture.IsReady(), "DescribeTopicFuture is not ready yet"); - - std::unordered_map partitionIdToParentId; - const auto& partitions = DescribeTopicFuture.GetValue().GetTopicDescription().GetPartitions(); - for (const auto& partition : partitions) { - auto parentPartitions = partition.GetParentPartitionIds(); - if (parentPartitions.empty()) { - continue; - } - - // we consider here that each partition has only one parent partition - partitionIdToParentId[partition.GetPartitionId()] = parentPartitions.front(); - } - - std::vector ancestors; - std::uint32_t currentPartitionId = PartitionId; - while (true) { - ancestors.push_back(currentPartitionId); - - auto parentPartitionId = partitionIdToParentId.find(currentPartitionId); - if (parentPartitionId == partitionIdToParentId.end()) { - break; - } - currentPartitionId = parentPartitionId->second; - } - - NotReadyFutures = ancestors.size(); - for (const auto& ancestor : ancestors) { - auto wrappedSession = Session->SessionsWorker->GetWriteSession(ancestor, false); - Y_ABORT_UNLESS(wrappedSession, "Write session not found"); - WriteSessions.push_back(wrappedSession); - - auto future = wrappedSession->Session->GetInitSeqNo(); - std::weak_ptr session = Session->shared_from_this(); - lock.unlock(); - future.Subscribe([this, session, wrappedSession, ancestor](const NThreading::TFuture& result) { - auto sessionPtr = session.lock(); - if (!sessionPtr) { - return; - } - - if (IsDone()) { - return; - } - - bool gotMaxSeqNo = false; - { - std::lock_guard lock(Lock); - if (result.HasException()) { - LOG_LAZY(sessionPtr->DbDriverState->Log, TLOG_ERR, sessionPtr->LogPrefix() << "Failed to get max seq no for partition " << ancestor << " for splitted partition " << PartitionId); - TSessionClosedEvent sessionClosedEvent(EStatus::INTERNAL_ERROR, {}); - sessionPtr->GetSessionClosedEventAndDie(wrappedSession, std::move(sessionClosedEvent)); - MoveTo(EState::Done); - return; - } - - UpdateMaxSeqNo(result.GetValue()); - if (--NotReadyFutures == 0) { - MoveTo(EState::GotMaxSeqNo); - gotMaxSeqNo = true; - } - } - - if (gotMaxSeqNo) { - sessionPtr->RunMainWorker(); - } - }); - lock.lock(); - GetMaxSeqNoFutures.push_back(future); - } - - if (ancestors.empty()) { - LOG_LAZY(Session->DbDriverState->Log, TLOG_INFO, Session->LogPrefix() << "No ancestors found for partition " << PartitionId); - MoveTo(EState::Init); - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TEventsWorkerWrapper - -TKeyedWriteSession::TEventsWorker::TEventsWorker(TKeyedWriteSession* session) - : Session(session) -{ - EventsPromise = NThreading::NewPromise(); - EventsFuture = EventsPromise.GetFuture(); - - AddReadyToAcceptEvent(); -} - -void TKeyedWriteSession::TEventsWorker::HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); - queueIt->second.push_back(TWriteSessionEvent::TEvent(std::move(event))); -} - -void TKeyedWriteSession::TEventsWorker::HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event) { - Session->MessagesWorker->HandleContinuationToken(partition, std::move(event.ContinuationToken)); -} - -void TKeyedWriteSession::TEventsWorker::HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition) { - if (event.IsSuccess()) { - return; - } - - Session->Partitions[partition].Locked_ = true; - - if (event.GetStatus() == EStatus::OVERLOADED) { - Session->HandleAutoPartitioning(partition); - return; - } - - if (!CloseEvent.has_value()) { - CloseEvent = std::move(event); - } - Session->NonBlockingClose(); -} - -bool TKeyedWriteSession::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition) { - while (true) { - auto event = wrappedSession->Session->GetEvent(false); - if (!event) { - break; - } - - if (auto sessionClosedEvent = std::get_if(&*event); sessionClosedEvent) { - HandleSessionClosedEvent(std::move(*sessionClosedEvent), partition); - return true; - } - - if (auto readyToAcceptEvent = std::get_if(&*event)) { - HandleReadyToAcceptEvent(partition, std::move(*readyToAcceptEvent)); - continue; - } - - if (auto acksEvent = std::get_if(&*event)) { - Session->SessionsWorker->OnReadFromSession(wrappedSession); - HandleAcksEvent(partition, std::move(*acksEvent)); - continue; - } - } - - return false; -} - -std::optional> TKeyedWriteSession::TEventsWorker::DoWork() { - std::unique_lock lock(Lock); - - while (!ReadyFutures.empty()) { - auto idx = *ReadyFutures.begin(); - ReadyFutures.erase(idx); - lock.unlock(); - // RunEventLoop without Lock: sub-session's WaitEvent() completion may run the Subscribe - // callback (ReadyFutures.insert) synchronously; that callback takes Lock -> same-thread deadlock. - auto isSessionClosed = RunEventLoop(Session->SessionsWorker->GetWriteSession(idx), idx); - if (!isSessionClosed) { - SubscribeToPartition(idx); - } else { - UnsubscribeFromPartition(idx); - } - lock.lock(); - } - - if (!Session->Done.load() && TransferEventsToOutputQueue()) { - return EventsPromise; - } - - return std::nullopt; -} - -void TKeyedWriteSession::TEventsWorker::SubscribeToPartition(std::uint32_t partition) { - if (Session->Partitions[partition].IsSplitted() || Session->SplittedPartitionWorkers.contains(partition)) { - Session->Partitions[partition].Future(NThreading::MakeFuture()); - return; - } - - auto wrappedSession = Session->SessionsWorker->GetWriteSession(partition); - auto newFuture = wrappedSession->Session->WaitEvent(); - std::weak_ptr session = Session->shared_from_this(); - std::weak_ptr self = shared_from_this(); - - newFuture.Subscribe([self, session, partition](const NThreading::TFuture&) { - auto sessionPtr = session.lock(); - if (!sessionPtr) { - return; - } - - auto selfPtr = self.lock(); - if (!selfPtr) { - return; - } - - { - std::lock_guard lock(selfPtr->Lock); - selfPtr->ReadyFutures.insert(partition); - } - sessionPtr->RunMainWorker(); - }); - Session->Partitions[partition].Future(newFuture); -} - -std::optional> TKeyedWriteSession::TEventsWorker::HandleNewMessage() { - std::lock_guard lock(Lock); - if (Session->MessagesWorker->IsMemoryUsageOK()) { - AddReadyToAcceptEvent(); - return EventsPromise; - } - - Session->Metrics.IncBufferFull(); - return std::nullopt; -} - -void TKeyedWriteSession::TEventsWorker::AddReadyToAcceptEvent() { - EventsOutputQueue.push_back(TWriteSessionEvent::TReadyToAcceptEvent(IssueContinuationToken())); - Session->Metrics.IncContinuationTokensSent(); -} - -bool TKeyedWriteSession::TEventsWorker::AddSessionClosedIfNeeded() { - if (!Session->Closed.load()) { - return false; - } - - if (!CloseEvent.has_value()) { - CloseEvent = TSessionClosedEvent(EStatus::SUCCESS, {}); - } - - if (EventsOutputQueue.empty() && (Session->MessagesWorker->IsQueueEmpty() || Session->Done.load())) { - EventsOutputQueue.push_back(*CloseEvent); - return true; - } - - return false; -} - -bool TKeyedWriteSession::TEventsWorker::TransferEventsToOutputQueue() { - bool eventsTransferred = false; - bool shouldAddReadyToAcceptEvent = false; - std::unordered_map> acks; - - auto messagesWorker = Session->MessagesWorker; - auto buildOutputAckEvent = [&](std::deque& acksQueue, std::uint64_t partition, std::optional expectedSeqNo) -> TWriteSessionEvent::TAcksEvent { - TWriteSessionEvent::TAcksEvent ackEvent; - - if (expectedSeqNo.has_value()) { - if (acksQueue.front().SeqNo != expectedSeqNo.value()) { - LOG_LAZY(Session->DbDriverState->Log, TLOG_ERR, Session->LogPrefix() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << partition); - } - Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << Session->Partitions[partition].PartitionId_); - } - - auto ack = std::move(acksQueue.front()); - ackEvent.Acks.push_back(std::move(ack)); - acksQueue.pop_front(); - return ackEvent; - }; - auto finishWithAck = [messagesWorker, &shouldAddReadyToAcceptEvent]() { - bool wasMemoryUsageOk = messagesWorker->IsMemoryUsageOK(); - messagesWorker->HandleAck(); - if (messagesWorker->IsMemoryUsageOK() && !wasMemoryUsageOk) { - shouldAddReadyToAcceptEvent = true; - } - }; - - while (messagesWorker->HasInFlightMessages()) { - const auto& head = messagesWorker->GetFrontInFlightMessage(); - - auto remainingAcks = acks.find(head.Partition); - if (remainingAcks != acks.end() && remainingAcks->second.size() > 0) { - EventsOutputQueue.push_back(buildOutputAckEvent(remainingAcks->second, head.Partition, head.SeqNo)); - finishWithAck(); - continue; - } - - auto eventsQueueIt = PartitionsEventQueues.find(head.Partition); - if (eventsQueueIt == PartitionsEventQueues.end() || eventsQueueIt->second.empty()) { - // No events for this message yet, stop processing (preserve order) - break; - } - - auto event = std::move(eventsQueueIt->second.front()); - auto acksEvent = std::get_if(&event); - Y_ABORT_UNLESS(acksEvent, "Expected AcksEvent only in PartitionsEventQueues"); - - std::deque acksQueue; - std::copy(acksEvent->Acks.begin(), acksEvent->Acks.end(), std::back_inserter(acksQueue)); - EventsOutputQueue.push_back(buildOutputAckEvent(acksQueue, head.Partition, head.SeqNo)); - acks[head.Partition] = std::move(acksQueue); - eventsQueueIt->second.pop_front(); - eventsTransferred = true; - - finishWithAck(); - } - - // this case handles situation: - // 1st message is written to partition 0 - // 2nd message is written to partition 1 - // 3rd message is written to partition 0 - // 4th message is written to partition 1 - // but AcksEvent for partition 0 looks like: - // [ack1, ack3] - // In this case we can not just forget about ack3, because 3rd message is in-flight - // so we will push 'AcksEvent' back to the queue for partition 0 - for (auto& [partition, acksQueue] : acks) { - if (acksQueue.size() > 0) { - TWriteSessionEvent::TAcksEvent ackEvent; - std::copy(acksQueue.begin(), acksQueue.end(), std::back_inserter(ackEvent.Acks)); - PartitionsEventQueues[partition].push_front(std::move(ackEvent)); - } - } - - if (shouldAddReadyToAcceptEvent) { - AddReadyToAcceptEvent(); - Session->Metrics.IncContinuationTokensSent(); - } - - return eventsTransferred; -} - -std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueBegin(std::uint32_t partition) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); - return queueIt->second.begin(); -} - -std::list::iterator TKeyedWriteSession::TEventsWorker::AckQueueEnd(std::uint32_t partition) { - auto [queueIt, _] = PartitionsEventQueues.try_emplace(partition); - return queueIt->second.end(); -} - -TKeyedWriteSession::TEventsWorker::EEventType TKeyedWriteSession::TEventsWorker::GetEventType(const TWriteSessionEvent::TEvent& event) { - if (std::holds_alternative(event)) { - return EEventType::SessionClosed; - } else if (std::holds_alternative(event)) { - return EEventType::ReadyToAccept; - } else if (std::holds_alternative(event)) { - return EEventType::Ack; - } - - Y_ABORT_UNLESS(false, "Unexpected event type"); -} - -std::optional TKeyedWriteSession::TEventsWorker::GetEventImpl(bool block, const std::vector& eventTypes) { - std::unique_lock lock(Lock); - if (EventsOutputQueue.empty() && block) { - lock.unlock(); - WaitEvent().Wait(); - lock.lock(); - } - - if (!EventsOutputQueue.empty()) { - if (!eventTypes.empty() && std::find(eventTypes.begin(), eventTypes.end(), GetEventType(EventsOutputQueue.front())) == eventTypes.end()) { - return std::nullopt; - } - - auto event = std::move(EventsOutputQueue.front()); - EventsOutputQueue.pop_front(); - return event; - } - - return std::nullopt; -} - -std::optional TKeyedWriteSession::TEventsWorker::GetEvent(bool block, const std::vector& eventTypes) { - { - std::unique_lock lock(Lock); - AddSessionClosedIfNeeded(); - } - auto event = GetEventImpl(block, eventTypes); - - return event; -} - -std::vector TKeyedWriteSession::TEventsWorker::GetEvents(bool block, std::optional maxEventsCount, const std::vector& eventTypes) { - if (maxEventsCount.has_value() && maxEventsCount.value() == 0) { - return {}; - } - - { - std::unique_lock lock(Lock); - AddSessionClosedIfNeeded(); - } - - std::vector events; - while (true) { - auto event = GetEventImpl(block, eventTypes); - if (!event) { - break; - } - - events.push_back(std::move(*event)); - if (maxEventsCount.has_value() && events.size() >= maxEventsCount.value()) { - break; - } - } - - return events; -} - -NThreading::TFuture TKeyedWriteSession::TEventsWorker::WaitEvent() { - std::unique_lock lock(Lock); - - AddSessionClosedIfNeeded(); - if (!EventsOutputQueue.empty()) { - return NThreading::MakeFuture(); - } - - if (EventsFuture.IsReady() && !Session->Closed.load()) { - EventsPromise = NThreading::NewPromise(); - EventsFuture = EventsPromise.GetFuture(); - } - - return EventsFuture; -} - -void TKeyedWriteSession::TEventsWorker::UnsubscribeFromPartition(std::uint32_t partition) { - ReadyFutures.erase(partition); - Session->Partitions[partition].Future(NThreading::MakeFuture()); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TSessionsWorker - -TKeyedWriteSession::TSessionsWorker::TSessionsWorker(TKeyedWriteSession* session) - : Session(session) -{} - -TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::GetWriteSession(std::uint32_t partition, bool directToPartition) { - auto sessionIter = SessionsIndex.find(partition); - if (sessionIter == SessionsIndex.end() || !directToPartition) { - return CreateWriteSession(partition, directToPartition); - } - - return sessionIter->second; -} - -std::string TKeyedWriteSession::TSessionsWorker::GetProducerId(std::uint32_t partitionId) { - return std::format("{}_{}", Session->Settings.ProducerIdPrefix_, partitionId); -} - -TKeyedWriteSession::WrappedWriteSessionPtr TKeyedWriteSession::TSessionsWorker::CreateWriteSession(std::uint32_t partition, bool directToPartition) { - auto partitionId = Session->Partitions[partition].PartitionId_; - auto producerId = GetProducerId(partitionId); - TWriteSessionSettings alteredSettings = Session->Settings; - - alteredSettings - .ProducerId(producerId) - .MessageGroupId(producerId) - .MaxMemoryUsage(std::numeric_limits::max()) - .RetryPolicy(Session->RetryPolicy) - .EventHandlers(TWriteSessionSettings::TEventHandlers() - .ReadyToAcceptHandler({}) - .AcksHandler({}) - .SessionClosedHandler({})); - - if (directToPartition) { - alteredSettings.DirectWriteToPartition(true); - alteredSettings.PartitionId(partitionId); - } - auto writeSession = std::make_shared( - Session->Client->CreateWriteSession(alteredSettings), - partition); - - if (directToPartition) { - SessionsIndex.emplace(partition, writeSession); - Session->EventsWorker->SubscribeToPartition(partition); - } - return writeSession; -} - -void TKeyedWriteSession::TSessionsWorker::DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout, bool mustBeEmpty) { - if (it == SessionsIndex.end() || !it->second) { - return; - } - - auto closeResult = it->second->Session->Close(closeTimeout); - Y_ABORT_UNLESS(!mustBeEmpty || closeResult, "There are still messages in flight"); - const auto partition = it->second->Partition; - it = SessionsIndex.erase(it); - Session->EventsWorker->UnsubscribeFromPartition(partition); -} - -void TKeyedWriteSession::TSessionsWorker::OnReadFromSession(WrappedWriteSessionPtr wrappedSession) { - if (wrappedSession->RemoveFromQueue(1)) { - Y_ABORT_UNLESS(!wrappedSession->IdleSession, "IdleSession is already set"); - auto idleSessionPtr = std::make_shared(wrappedSession.get(), TInstant::Now(), Session->Settings.SubSessionIdleTimeout_); - auto [itIdle, inserted] = IdlerSessions.insert(idleSessionPtr); - Y_ABORT_UNLESS(inserted, "Duplicate idle session for partition"); - IdlerSessionsIndex[wrappedSession->Partition] = itIdle; - wrappedSession->IdleSession = idleSessionPtr; - } -} - -void TKeyedWriteSession::TSessionsWorker::OnWriteToSession(WrappedWriteSessionPtr wrappedSession) { - if (wrappedSession->AddToQueue(1) && wrappedSession->IdleSession) { - auto itIdle = IdlerSessionsIndex.find(wrappedSession->Partition); - if (itIdle != IdlerSessionsIndex.end()) { - IdlerSessions.erase(itIdle->second); - IdlerSessionsIndex.erase(itIdle); - } - wrappedSession->IdleSession.reset(); - } -} - -void TKeyedWriteSession::TSessionsWorker::DoWork() { - while (!IdlerSessions.empty()) { - auto it = IdlerSessions.begin(); - if (!(*it)->IsExpired()) { - break; - } - - const auto partition = (*it)->Session->Partition; - if (Session->Partitions[partition].Locked_) { - continue; - } - - // Remove idle tracking first to keep containers consistent even if the session - // is already absent from SessionsIndex. - IdlerSessions.erase(it); - IdlerSessionsIndex.erase(partition); - - auto sessionIter = SessionsIndex.find(partition); - if (sessionIter != SessionsIndex.end()) { - sessionIter->second->IdleSession.reset(); - DestroyWriteSession(sessionIter, TDuration::Zero(), !Session->SplittedPartitionWorkers.contains(partition)); - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TMessagesWorker - -TKeyedWriteSession::TMessagesWorker::TMessagesWorker(TKeyedWriteSession* session) - : Session(session) -{ -} - -void TKeyedWriteSession::TMessagesWorker::RechoosePartitionIfNeeded(MessageIter message) { - const auto& partitionInfo = Session->Partitions[message->Partition]; - if (partitionInfo.Children_.empty()) { - return; - } - - // this case means that partition was split, so we need to rechoose the partition for the message - auto newPartition = Session->PartitionChooser->ChoosePartition(message->Key); - message->Partition = newPartition; -} - -void TKeyedWriteSession::TMessagesWorker::DoWork() { - auto sessionsWorker = Session->SessionsWorker; - - auto iterateMessagesIndex = [&](std::unordered_map>& messagesIndex, auto stopCondition) { - std::vector partitionsProcessed; - for (auto& [partition, messages] : messagesIndex) { - while (!messages.empty()) { - auto head = messages.front(); - if (stopCondition(head)) { - break; - } - - auto wrappedSession = sessionsWorker->GetWriteSession(head->Partition); - if (!SendMessage(wrappedSession, *head)) { - break; - } - - Session->Metrics.AddWriteLag((TInstant::Now() - head->CreateTimestamp.value_or(TInstant::Now())).MilliSeconds()); - head->Sent = true; - sessionsWorker->OnWriteToSession(wrappedSession); - messages.pop_front(); - } - - if (messages.empty()) { - partitionsProcessed.push_back(partition); - } - } - - for (const auto& partition : partitionsProcessed) { - messagesIndex.erase(partition); - } - }; - - iterateMessagesIndex( - MessagesToResendIndex, - [](MessageIter) { - return false; - } - ); - - iterateMessagesIndex( - PendingMessagesIndex, - [this](MessageIter head) { - return Session->Partitions[head->Partition].Locked_ || - MessagesToResendIndex.contains(head->Partition); - } - ); -} - -bool TKeyedWriteSession::TMessagesWorker::SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message) { - if (!wrappedSession) { - return false; - } - - auto continuationToken = GetContinuationToken(message.Partition); - if (!continuationToken) { - return false; - } - - Session->Metrics.IncOutgoingMessages(); - wrappedSession->Session->Write(std::move(*continuationToken), message.BuildMessage(), message.Tx); - return true; -} - -void TKeyedWriteSession::TMessagesWorker::PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message) { - auto iter = InFlightMessages.insert(InFlightMessages.end(), std::move(message)); - auto [inFlightMessagesIndexIt, _] = InFlightMessagesIndex.try_emplace(partition); - inFlightMessagesIndexIt->second.push_back(iter); - - auto [pendingMessagesIndexIt, __] = PendingMessagesIndex.try_emplace(partition); - pendingMessagesIndexIt->second.push_back(iter); -} - -void TKeyedWriteSession::TMessagesWorker::HandleAck() { - PopInFlightMessage(); -} - -void TKeyedWriteSession::TMessagesWorker::PopInFlightMessage() { - Y_ABORT_UNLESS(!InFlightMessages.empty()); - const std::uint64_t partition = InFlightMessages.front().Partition; - const auto it = InFlightMessages.begin(); - - auto mapIt = InFlightMessagesIndex.find(partition); - if (mapIt != InFlightMessagesIndex.end()) { - auto& list = mapIt->second; - for (auto listIt = list.begin(); listIt != list.end(); ++listIt) { - if (*listIt == it) { - list.erase(listIt); - break; - } - } - if (list.empty()) { - InFlightMessagesIndex.erase(mapIt); - } - } - - Y_ABORT_UNLESS(it->Data.size() <= MemoryUsage, "MemoryUsage is less than the size of the message"); - MemoryUsage -= it->Data.size(); - InFlightMessages.pop_front(); -} - -bool TKeyedWriteSession::TMessagesWorker::IsMemoryUsageOK() const { - return MemoryUsage <= Session->Settings.MaxMemoryUsage_ / 2; -} - -void TKeyedWriteSession::TMessagesWorker::AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx) { - MemoryUsage += message.Data.size(); - PushInFlightMessage(partition, TMessageInfo(key, std::move(message), partition, tx)); -} - -std::optional TKeyedWriteSession::TMessagesWorker::GetContinuationToken(std::uint32_t partition) { - auto it = ContinuationTokens.find(partition); - if (it != ContinuationTokens.end() && !it->second.empty()) { - auto token = std::move(it->second.front()); - it->second.pop_front(); - if (it->second.empty()) { - ContinuationTokens.erase(it); - } - return token; - } - - return std::nullopt; -} - -void TKeyedWriteSession::TMessagesWorker::HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken) { - auto [it, _] = ContinuationTokens.try_emplace(partition); - it->second.push_back(std::move(continuationToken)); -} - -bool TKeyedWriteSession::TMessagesWorker::IsQueueEmpty() const { - return InFlightMessages.empty(); -} - -const TKeyedWriteSession::TMessageInfo& TKeyedWriteSession::TMessagesWorker::GetFrontInFlightMessage() const { - Y_ABORT_UNLESS(!InFlightMessages.empty()); - return InFlightMessages.front(); -} - -bool TKeyedWriteSession::TMessagesWorker::HasInFlightMessages() const { - return !InFlightMessages.empty(); -} - -void TKeyedWriteSession::TMessagesWorker::ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo) { - auto it = InFlightMessagesIndex.find(partition); - if (it == InFlightMessagesIndex.end()) { - return; - } - - auto& list = it->second; - auto resendIt = list.begin(); - auto ackQueueIt = Session->EventsWorker->AckQueueBegin(partition); - size_t ackIdx = 0; - auto ackQueueEnd = Session->EventsWorker->AckQueueEnd(partition); - std::vector acksToSend; - - while (resendIt != list.end()) { - if (!(*resendIt)->SeqNo.has_value() || (*resendIt)->SeqNo.value() > afterSeqNo) { - break; - } - - auto seqNo = (*resendIt)->SeqNo.value(); - if (ackQueueIt == ackQueueEnd) { - // this case can happen if the message was sent, but session was closed before the ack was received - TWriteSessionEvent::TWriteAck ack; - ack.SeqNo = seqNo; - acksToSend.push_back(std::move(ack)); - } else { - auto acksEvent = std::get_if(&*ackQueueIt); - if (ackIdx == acksEvent->Acks.size()) { - ++ackQueueIt; - ackIdx = 0; - continue; - } - - if (acksEvent->Acks[ackIdx].SeqNo > seqNo) { - // this case can happen if the message was sent, but session was closed before the ack was received - TWriteSessionEvent::TWriteAck ack; - ack.SeqNo = seqNo; - acksEvent->Acks.insert(acksEvent->Acks.begin() + ackIdx, std::move(ack)); - } - ++ackIdx; - } - ++resendIt; - } - - if (!acksToSend.empty()) { - TWriteSessionEvent::TAcksEvent event; - event.Acks = std::move(acksToSend); - Session->EventsWorker->HandleAcksEvent(partition, std::move(event)); - } - - // IMPORTANT: do not mutate InFlightMessagesIndex while holding references/iterators to its elements. - // try_emplace()/rehash may invalidate 'it' and 'list' -> use-after-free and segfaults. - std::vector> messagesFromOldPartition; - messagesFromOldPartition.reserve(std::distance(resendIt, list.end())); - auto currentSeqNo = resendIt != list.end() ? (*resendIt)->SeqNo.value_or(0) : 0; - for (auto iter = resendIt; iter != list.end(); ++iter) { - if (iter != resendIt && currentSeqNo != 0) { - Y_ABORT_UNLESS((*iter)->SeqNo.value_or(0) > currentSeqNo, "SeqNo is not increasing for partition %d", partition); - } - - auto newPartition = Session->PartitionChooser->ChoosePartition((*iter)->Key); - (*iter)->Partition = newPartition; - messagesFromOldPartition.emplace_back(newPartition, *iter); - - currentSeqNo = (*iter)->SeqNo.value_or(0); - } - - list.erase(resendIt, list.end()); - for (const auto& [newPartition, msgIt] : messagesFromOldPartition) { - auto [inFlightMessagesIndexChainIt, _] = InFlightMessagesIndex.try_emplace(newPartition); - inFlightMessagesIndexChainIt->second.push_back(msgIt); - - if (msgIt->Sent) { - auto [messagesToResendChainIt, __] = MessagesToResendIndex.try_emplace(newPartition); - messagesToResendChainIt->second.push_back(msgIt); - } - } - - InFlightMessagesIndex.erase(partition); -} - -void TKeyedWriteSession::TMessagesWorker::RebuildPendingMessagesIndex(std::uint32_t partition) { - auto [oldPendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(partition); - std::unordered_map> pendingMessagesForNewPartitions; - for (auto it = oldPendingMessagesIndexChainIt->second.begin(); it != oldPendingMessagesIndexChainIt->second.end(); ++it) { - auto newPartition = Session->PartitionChooser->ChoosePartition((*it)->Key); - auto [pendingMessagesForNewPartitionsIt, __] = pendingMessagesForNewPartitions.try_emplace(newPartition); - pendingMessagesForNewPartitionsIt->second.push_back(*it); - } - - for (const auto& [newPartition, pendingMessagesForNewPartition] : pendingMessagesForNewPartitions) { - auto [pendingMessagesIndexChainIt, __] = PendingMessagesIndex.try_emplace(newPartition); - for (auto reverseIt = pendingMessagesForNewPartition.rbegin(); reverseIt != pendingMessagesForNewPartition.rend(); ++reverseIt) { - pendingMessagesIndexChainIt->second.push_front(*reverseIt); - } - } - - PendingMessagesIndex.erase(partition); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::TKeyedWriteSessionRetryPolicy - -TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::TKeyedWriteSessionRetryPolicy(TKeyedWriteSession* session) - : Session(session) -{} - -typename TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::IRetryState::TPtr TKeyedWriteSession::TKeyedWriteSessionRetryPolicy::CreateRetryState() const { - struct TRetryState : public IRetryState { - TRetryState(TKeyedWriteSession* session) - : Session(session) - {} - ~TRetryState() = default; - TMaybe GetNextRetryDelay(EStatus status) override { - if (status == EStatus::OVERLOADED) { - return Nothing(); - } - - if (!UserRetryState) { - auto policy = Session->Settings.RetryPolicy_ ? Session->Settings.RetryPolicy_ : NYdb::NTopic::IRetryPolicy::GetDefaultPolicy(); - UserRetryState = policy->CreateRetryState(); - } - - return UserRetryState->GetNextRetryDelay(status); - } - - private: - TKeyedWriteSession* Session; - IRetryState::TPtr UserRetryState; - }; - - return std::make_unique(Session); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession::Metrics - -void TKeyedWriteSession::TMetricGauge::Add(std::uint64_t value) { - Sum += value; - Max = std::max(Max, value); - MetricCount++; -} - -void TKeyedWriteSession::TMetricGauge::Clear() { - Sum = 0; - MetricCount = 0; - Max = 0; -} - -long double TKeyedWriteSession::TMetricGauge::Average() { - if (MetricCount == 0) { - return 0; - } - - return (long double)Sum / (long double)MetricCount; -} - -std::uint64_t TKeyedWriteSession::TMetricGauge::GetMax() const { - return Max; -} - -std::uint64_t TKeyedWriteSession::TMetricGauge::GetSum() const { - return Sum; -} - -TKeyedWriteSession::TMetrics::TMetrics(TKeyedWriteSession* session): Session(session) {} - -void TKeyedWriteSession::TMetrics::AddMainWorkerTime(std::uint64_t ms) { - std::lock_guard lock(Lock); - MainWorkerTimeMs.Add(ms); -} - -void TKeyedWriteSession::TMetrics::AddCycleTime(std::uint64_t ms) { - std::lock_guard lock(Lock); - CycleTimeMs.Add(ms); -} - -void TKeyedWriteSession::TMetrics::AddWriteLag(std::uint64_t lagMs) { - std::lock_guard lock(Lock); - WriteLagMs.Add(lagMs); -} - -void TKeyedWriteSession::TMetrics::IncContinuationTokensSent() { - std::lock_guard lock(Lock); - ContinuationTokensSent.Add(1); -} - -void TKeyedWriteSession::TMetrics::IncBufferFull() { - std::lock_guard lock(Lock); - BufferFull.Add(1); -} - -void TKeyedWriteSession::TMetrics::IncIncomingMessages() { - std::lock_guard lock(Lock); - IncomingMessages.Add(1); -} - -void TKeyedWriteSession::TMetrics::IncOutgoingMessages() { - std::lock_guard lock(Lock); - OutgoingMessages.Add(1); -} - -void TKeyedWriteSession::TMetrics::PrintMetrics() { - std::lock_guard lock(Lock); - LOG_LAZY( - Session->DbDriverState->Log, - TLOG_ERR, - Session->LogPrefix() - << "METRICS: average MainWorkerTimeMs: " << MainWorkerTimeMs.Average() - << " ms, average CycleTimeMs: " << CycleTimeMs.Average() - << " ms, average WriteLagMs: " << WriteLagMs.Average() << " ms, " - << "max MainWorkerTimeMs: " << MainWorkerTimeMs.GetMax() << " ms, " - << "max CycleTimeMs: " << CycleTimeMs.GetMax() << " ms, " - << "max WriteLagMs: " << WriteLagMs.GetMax() << " ms, " - << "ContinuationTokensSent: " << ContinuationTokensSent.GetSum() << " tokens, " - << "BufferFull: " << BufferFull.GetSum() << " times, " - << "IncomingMessages: " << IncomingMessages.GetSum() << " messages, " - << "OutgoingMessages: " << OutgoingMessages.GetSum() << " messages"); - MainWorkerTimeMs.Clear(); - CycleTimeMs.Clear(); - WriteLagMs.Clear(); - ContinuationTokensSent.Clear(); - BufferFull.Clear(); - IncomingMessages.Clear(); - OutgoingMessages.Clear(); -} - - -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession - -TKeyedWriteSession::TKeyedWriteSession( - const TKeyedWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState) - : Connections(connections), - Client(client), - DbDriverState(dbDriverState), - Metrics(this), - Settings(settings) -{ - if (settings.ProducerIdPrefix_.empty()) { - ythrow TContractViolation("ProducerIdPrefix is required for KeyedWriteSession"); - } - - if (!settings.ProducerId_.empty()) { - ythrow TContractViolation("ProducerId should be empty for KeyedWriteSession, use ProducerIdPrefix instead"); - } - - if (!settings.MessageGroupId_.empty()) { - ythrow TContractViolation("MessageGroupId should be empty for KeyedWriteSession"); - } - - TDescribeTopicSettings describeTopicSettings; - auto topicConfig = client->DescribeTopic(settings.Path_, describeTopicSettings).GetValueSync(); - auto partitions = topicConfig.GetTopicDescription().GetPartitions(); - std::sort(partitions.begin(), partitions.end(), [](const auto& a, const auto& b) -> bool { - return a.GetPartitionId() < b.GetPartitionId(); - }); - - auto partitionChooserStrategy = settings.PartitionChooserStrategy_; - auto strategy = topicConfig.GetTopicDescription().GetPartitioningSettings().GetAutoPartitioningSettings().GetStrategy(); - auto autoPartitioningEnabled = (strategy != EAutoPartitioningStrategy::Disabled && - strategy != EAutoPartitioningStrategy::Unspecified); - - for (const auto& partition : partitions) { - auto partitionId = partition.GetPartitionId(); - auto fromBound = partition.GetFromBound().value_or(""); - auto toBound = partition.GetToBound(); - LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Adding partition " << partitionId << " from bound " << fromBound << " to bound " << (toBound.has_value() ? toBound.value() : "null")); - Partitions[partitionId] = TPartitionInfo() - .PartitionId(partitionId) - .FromBound(fromBound) - .ToBound(toBound); - } - - for (const auto& partition : partitions) { - auto children = partition.GetChildPartitionIds(); - - std::vector childrenIndices; - childrenIndices.reserve(children.size()); - for (auto child : children) { - childrenIndices.push_back(child); - } - Partitions[partition.GetPartitionId()].Children(childrenIndices); - } - - if (Settings.EventHandlers_.CommonHandler_) { - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); - } else { - if (Settings.EventHandlers_.SessionClosedHandler_) { - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::SessionClosed); - } - if (Settings.EventHandlers_.ReadyToAcceptHandler_) { - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::ReadyToAccept); - } - if (Settings.EventHandlers_.AcksHandler_) { - EventTypesWithHandlers.push_back(TEventsWorker::EEventType::Ack); - } - } - - switch (partitionChooserStrategy) { - case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound: - PartitioningKeyHasher = settings.PartitioningKeyHasher_; - PartitionChooser = std::make_unique(this); - for (size_t i = 0; i < Partitions.size(); ++i) { - if (i > 0 && Partitions[i].FromBound_.empty() && !Partitions[i].ToBound_.has_value()) { - ythrow TContractViolation("Unbounded partition is not supported for Bound partition chooser strategy"); - } - - if (!Partitions[i].Children_.empty()) { - continue; - } - - PartitionsIndex[Partitions[i].FromBound_] = Partitions[i].PartitionId_; - } - break; - case TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash: - if (autoPartitioningEnabled) { - ythrow TContractViolation("Hash partition chooser strategy is not supported for topic with auto partitioning"); - } - - std::vector partitionsIds; - partitionsIds.reserve(partitions.size()); - for (const auto& partition : partitions) { - partitionsIds.push_back(partition.GetPartitionId()); - } - - PartitionChooser = std::make_unique(std::move(partitionsIds)); - break; - } - - ClosePromise = NThreading::NewPromise(); - CloseFuture = ClosePromise.GetFuture(); - ShutdownPromise = NThreading::NewPromise(); - ShutdownFuture = ShutdownPromise.GetFuture(); - - SessionsWorker = std::make_shared(this); - MessagesWorker = std::make_shared(this); - EventsWorker = std::make_shared(this); - RetryPolicy = std::make_shared(this); - - // Start handlers executor for user callbacks (Acks/ReadyToAccept/SessionClosed/Common). - Settings.EventHandlers_.HandlersExecutor_->Start(); - - CloseFuture.Subscribe([this](const NThreading::TFuture&) { - RunMainWorker(); - }); - - RunMainWorker(); - - LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Keyed write session created"); -} - -std::vector TKeyedWriteSession::GetPartitions() const { - std::vector partitions; - partitions.reserve(Partitions.size()); - for (const auto& [partitionId, partitionInfo] : Partitions) { - partitions.push_back(partitionInfo); - } - return partitions; -} - -std::unordered_map TKeyedWriteSession::GetPartitionsMap() const { - return Partitions; -} - -std::map TKeyedWriteSession::GetPartitionsIndex() const { - return PartitionsIndex; -} - -void TKeyedWriteSession::Write(TContinuationToken&&, const std::string& key, TWriteMessage&& message, TTransactionBase* tx) { - std::optional> eventsPromise; - { - std::lock_guard lock(GlobalLock); - Metrics.IncIncomingMessages(); - if (Closed.load()) { - return; - } - - if ((message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithoutSeqNo) - || (!message.SeqNo_.has_value() && SeqNoStrategy == ESeqNoStrategy::WithSeqNo)) { - ythrow TContractViolation("Can not mix messages with and without seqNo"); - } - - if (SeqNoStrategy == ESeqNoStrategy::NotInitialized) { - SeqNoStrategy = message.SeqNo_.has_value() ? ESeqNoStrategy::WithSeqNo : ESeqNoStrategy::WithoutSeqNo; - } - - auto partition = PartitionChooser->ChoosePartition(key); - MessagesWorker->AddMessage(key, std::move(message), partition, tx); - eventsPromise = EventsWorker->HandleNewMessage(); - RunUserEventLoop(); - } - - RunMainWorker(); - if (eventsPromise) { - eventsPromise->TrySetValue(); - } -} - -bool TKeyedWriteSession::Close(TDuration closeTimeout) { - if (Closed.exchange(true)) { - std::lock_guard lock(GlobalLock); - return MessagesWorker->IsQueueEmpty(); - } - - SetCloseDeadline(closeTimeout); - - ClosePromise.TrySetValue(); - ShutdownFuture.Wait(CloseDeadline); - RunUserEventLoop(); - Done.store(true); - - // No need to lock here, because we are waiting for the shutdown future and it will block until the main worker is done - return MessagesWorker->IsQueueEmpty(); -} - -void TKeyedWriteSession::NonBlockingClose() { - Closed.store(true); - Done.store(true); -} - -void TKeyedWriteSession::SetCloseDeadline(const TDuration& closeTimeout) { - std::lock_guard lock(GlobalLock); - CloseDeadline = TInstant::Now() + closeTimeout; -} - -TKeyedWriteSession::~TKeyedWriteSession() { - Close(TDuration::Zero()); - Settings.EventHandlers_.HandlersExecutor_->Stop(); - ShutdownFuture.Wait(); -} - -NThreading::TFuture TKeyedWriteSession::WaitEvent() { - return EventsWorker->WaitEvent(); -} - -std::optional TKeyedWriteSession::GetEvent(bool block) { - if (Settings.EventHandlers_.CommonHandler_) { - return std::nullopt; - } - - return EventsWorker->GetEvent(block); -} - -std::vector TKeyedWriteSession::GetEvents(bool block, std::optional maxEventsCount) { - if (Settings.EventHandlers_.CommonHandler_) { - return {}; - } - - return EventsWorker->GetEvents(block, maxEventsCount); -} - -TDuration TKeyedWriteSession::GetCloseTimeout() { - std::lock_guard lock(GlobalLock); - auto now = TInstant::Now(); - if (CloseDeadline <= now) { - return TDuration::Zero(); - } - return CloseDeadline - now; -} - -bool TKeyedWriteSession::RunSplittedPartitionWorkers() { - if (SplittedPartitionWorkers.empty() && ReadySplittedPartitionWorkers.empty()) { - return false; - } - - bool needRerun = false; - std::unordered_map> readySplittedPartitionWorkers; - for (const auto& [partition, splittedPartitionWorker] : SplittedPartitionWorkers) { - if (splittedPartitionWorker->IsDone()) { - readySplittedPartitionWorkers[partition] = splittedPartitionWorker; - continue; - } - - splittedPartitionWorker->DoWork(); - needRerun = needRerun || splittedPartitionWorker->IsInit(); - needRerun = needRerun || splittedPartitionWorker->IsDone(); - } - - for (const auto& [partition, splittedPartitionWorker] : readySplittedPartitionWorkers) { - ReadySplittedPartitionWorkers[partition] = splittedPartitionWorker; - SplittedPartitionWorkers.erase(partition); - } - - std::vector partitionsToRemove; - for (const auto& [partition, splittedPartitionWorker] : ReadySplittedPartitionWorkers) { - if (splittedPartitionWorker->CanBeRemoved()) { - partitionsToRemove.push_back(partition); - } - } - - for (const auto& partition : partitionsToRemove) { - ReadySplittedPartitionWorkers.erase(partition); - } - - return needRerun; -} - -void TKeyedWriteSession::RunUserEventLoop() { - if (!Settings.EventHandlers_.AcksHandler_ && - !Settings.EventHandlers_.ReadyToAcceptHandler_ && - !Settings.EventHandlers_.SessionClosedHandler_ && - !Settings.EventHandlers_.CommonHandler_) { - return; - } - - auto handlersExecutor = Settings.EventHandlers_.HandlersExecutor_; - if (!handlersExecutor) { - return; - } - - while (true) { - auto event = EventsWorker->GetEvent(false, EventTypesWithHandlers); - if (!event) { - break; - } - - if (auto* readyToAcceptEvent = std::get_if(&*event)) { - if (Settings.EventHandlers_.ReadyToAcceptHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*readyToAcceptEvent)]() mutable { - Settings.EventHandlers_.ReadyToAcceptHandler_(ev); - }); - } else if (Settings.EventHandlers_.CommonHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*event)]() mutable { - Settings.EventHandlers_.CommonHandler_(ev); - }); - } - continue; - } - - if (auto* acksEvent = std::get_if(&*event)) { - if (Settings.EventHandlers_.AcksHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*acksEvent)]() mutable { - Settings.EventHandlers_.AcksHandler_(ev); - }); - } else if (Settings.EventHandlers_.CommonHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*event)]() mutable { - Settings.EventHandlers_.CommonHandler_(ev); - }); - } - continue; - } - - if (auto* sessionClosedEvent = std::get_if(&*event)) { - if (Settings.EventHandlers_.SessionClosedHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*sessionClosedEvent)]() mutable { - Settings.EventHandlers_.SessionClosedHandler_(ev); - }); - } else if (Settings.EventHandlers_.CommonHandler_) { - handlersExecutor->Post( - [this, ev = std::move(*event)]() mutable { - Settings.EventHandlers_.CommonHandler_(ev); - }); - } - break; - } - } -} - -void TKeyedWriteSession::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent) { - std::optional receivedSessionClosedEvent; - while (true) { - auto event = wrappedSession->Session->GetEvent(false); - if (!event) { - break; - } - - if (auto* closedEvent = std::get_if(&*event)) { - receivedSessionClosedEvent = std::move(*closedEvent); - break; - } - } - - if (!receivedSessionClosedEvent || receivedSessionClosedEvent->GetStatus() == EStatus::SUCCESS || receivedSessionClosedEvent->GetStatus() == EStatus::OVERLOADED) { - LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Failed to get session closed event"); - EventsWorker->HandleSessionClosedEvent(std::move(*sessionClosedEvent), wrappedSession->Partition); - } else { - EventsWorker->HandleSessionClosedEvent(std::move(*receivedSessionClosedEvent), wrappedSession->Partition); - } -} - -TStringBuilder TKeyedWriteSession::LogPrefix() { - return TStringBuilder() << " SessionId: " << Settings.SessionId_ << " Epoch: " << Epoch.load() << " "; -} - -void TKeyedWriteSession::NextEpoch() { - auto maxEpoch = MAX_EPOCH - 1; - if (Epoch.compare_exchange_weak(maxEpoch, 0)) { - LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Epoch overflow, resetting to 0"); - return; - } - - Epoch.fetch_add(1); -} - -void TKeyedWriteSession::RunMainWorker() { - // This function is both "request to run" and the runner itself. - // We must handle two properties: - // - TFuture::Subscribe may call back synchronously when future is already ready. - // - A callback may race with the runner trying to go idle (avoid lost wakeups). - enum : std::uint8_t { - Idle = 0, - Running = 1, - Rerun = 2, - }; - - // Try to become the runner. If already running, just request a rerun. - std::uint8_t state = MainWorkerState.load(std::memory_order_acquire); - for (;;) { - if (state & Running) { - if (MainWorkerState.compare_exchange_weak(state, std::uint8_t(state | Rerun), - std::memory_order_acq_rel, - std::memory_order_acquire)) { - return; - } - continue; - } else { - if (MainWorkerState.compare_exchange_weak(state, Running, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - break; // we are the runner now - } - continue; - } - } - - NextEpoch(); - - auto startWorkerTime = TInstant::Now(); - // Runner loop: process, arm subscription, then either go idle or loop again. - for (;;) { - auto startIter = TInstant::Now(); - // Clear rerun request for this iteration. - MainWorkerState.fetch_and(std::uint8_t(~Rerun), std::memory_order_acq_rel); - bool needRerun = false; - std::optional> eventsPromise; - - { - std::unique_lock lock(GlobalLock); - eventsPromise = EventsWorker->DoWork(); - RunUserEventLoop(); - needRerun = RunSplittedPartitionWorkers(); - if (!Done.load()) { - SessionsWorker->DoWork(); - MessagesWorker->DoWork(); - } - } - - if (eventsPromise) { - eventsPromise->TrySetValue(); - } - - const auto isClosed = Closed.load(); - const auto closeTimeout = GetCloseTimeout(); - if (isClosed && (Done.load() || MessagesWorker->IsQueueEmpty() || closeTimeout == TDuration::Zero())) { - ShutdownPromise.TrySetValue(); - EventsWorker->EventsPromise.TrySetValue(); - ClosePromise.TrySetValue(); - MainWorkerState.store(Idle, std::memory_order_release); - return; - } - - if (needRerun) { - // we need this case to start resending messages if there are any - Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); - continue; - } - - // Try to go idle. If someone requested rerun concurrently, keep running. - std::uint8_t cur = MainWorkerState.load(std::memory_order_acquire); - for (;;) { - if (cur & Rerun) { - Metrics.AddCycleTime((TInstant::Now() - startIter).MilliSeconds()); - break; // continue outer loop - } - if (MainWorkerState.compare_exchange_weak(cur, Idle, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - auto workerFinished = TInstant::Now(); - Metrics.AddCycleTime((workerFinished - startIter).MilliSeconds()); - Metrics.AddMainWorkerTime((workerFinished - startWorkerTime).MilliSeconds()); - return; // successfully went idle - } - } - // Rerun was requested; continue the loop without recursion. - } -} - -TInstant TKeyedWriteSession::GetCloseDeadline() { - std::lock_guard lock(GlobalLock); - return CloseDeadline; -} - -void TKeyedWriteSession::HandleAutoPartitioning(std::uint32_t partition) { - LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "HandleAutoPartitioning: " << partition); - auto splittedPartitionWorker = std::make_shared(this, partition); - SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); -} - -std::string TKeyedWriteSession::GetProducerId(std::uint32_t partition) { - return std::format("{}_{}", Settings.ProducerIdPrefix_, partition); -} - -TWriterCounters::TPtr TKeyedWriteSession::GetCounters() { - return nullptr; -} - -TKeyedWriteSession::TBoundPartitionChooser::TBoundPartitionChooser(TKeyedWriteSession* session) - : Session(session) -{} - -std::uint32_t TKeyedWriteSession::TBoundPartitionChooser::ChoosePartition(const std::string_view key) { - auto hashedKey = Session->PartitioningKeyHasher(key); - - auto lowerBound = Session->PartitionsIndex.lower_bound(hashedKey); - if (lowerBound != Session->PartitionsIndex.end() && lowerBound->first == hashedKey) { - return lowerBound->second; - } - - Y_ABORT_IF(lowerBound == Session->PartitionsIndex.begin(), "Lower bound is the first element"); - return std::prev(lowerBound)->second; -} - -TKeyedWriteSession::THashPartitionChooser::THashPartitionChooser(std::vector&& partitions) - : Partitions(std::move(partitions)) -{} - -std::uint32_t TKeyedWriteSession::THashPartitionChooser::ChoosePartition(const std::string_view key) { - auto hash = MurmurHash(key.data(), key.size()); - return Partitions[hash % Partitions.size()]; -} - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TSimpleBlockingWriteSession @@ -1857,129 +164,4 @@ bool TSimpleBlockingWriteSession::Close(TDuration closeTimeout) { return Writer->Close(std::move(closeTimeout)); } -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TSimpleBlockingKeyedWriteSession - -TSimpleBlockingKeyedWriteSession::TSimpleBlockingKeyedWriteSession( - const TKeyedWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState) - : Writer(std::make_shared(settings, client, connections, dbDriverState)) -{ - ClosePromise = NThreading::NewPromise(); - CloseFuture = ClosePromise.GetFuture(); -} - -void TSimpleBlockingKeyedWriteSession::RunEventLoop() { - while (true) { - auto event = Writer->GetEvent(false); - if (!event) { - break; - } - - if (auto readyToAcceptEvent = std::get_if(&*event)) { - ContinuationTokensQueue.push(std::move(readyToAcceptEvent->ContinuationToken)); - continue; - } - if (std::get_if(&*event)) { - Closed.store(true); - return; - } - if (auto acksEvent = std::get_if(&*event)) { - HandleAcksEvent(std::move(*acksEvent)); - } - } -} - -void TSimpleBlockingKeyedWriteSession::HandleAcksEvent(const TWriteSessionEvent::TAcksEvent& acksEvent) { - for (auto ack : acksEvent.Acks) { - AckedSeqNos.insert(ack.SeqNo); - } -} - -template -bool TSimpleBlockingKeyedWriteSession::Wait(const TDuration& timeout, F&& stopFunc) { - std::unique_lock lock(Lock); - - auto deadline = TInstant::Now() + timeout; - while (true) { - if (TInstant::Now() > deadline) { - return false; - } - - RunEventLoop(); - - if (stopFunc()) { - return true; - } - - if (Closed.load()) { - return false; - } - - std::vector> futures; - futures.push_back(CloseFuture); - futures.push_back(Writer->WaitEvent()); - lock.unlock(); - NThreading::NWait::WaitAny(futures).Wait(deadline); - lock.lock(); - } -} - -std::optional TSimpleBlockingKeyedWriteSession::GetContinuationToken(TDuration timeout) { - std::optional token; - - Wait(timeout, [&]() { - if (!ContinuationTokensQueue.empty()) { - token = std::move(ContinuationTokensQueue.front()); - ContinuationTokensQueue.pop(); - return true; - } - return false; - }); - - return token; -} - -bool TSimpleBlockingKeyedWriteSession::WaitForAck(std::optional seqNo, TDuration timeout) { - return Wait(timeout, [&]() { - if (!seqNo.has_value()) { - if (AckedSeqNos.empty()) { - return false; - } - - AckedSeqNos.erase(AckedSeqNos.begin()); - return true; - } - - if (AckedSeqNos.contains(*seqNo)) { - AckedSeqNos.erase(*seqNo); - return true; - } - return false; - }); -} - -bool TSimpleBlockingKeyedWriteSession::Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx, TDuration blockTimeout) { - auto continuationToken = GetContinuationToken(blockTimeout); - if (!continuationToken) { - return false; - } - - auto seqNo = message.SeqNo_; - Writer->Write(std::move(*continuationToken), std::move(key), std::move(message), tx); - return WaitForAck(seqNo, blockTimeout); -} - -bool TSimpleBlockingKeyedWriteSession::Close(TDuration closeTimeout) { - Closed.store(true); - ClosePromise.TrySetValue(); - return Writer->Close(closeTimeout); -} - -TWriterCounters::TPtr TSimpleBlockingKeyedWriteSession::GetCounters() { - return nullptr; -} - } // namespace NYdb::inline V3::NTopic \ No newline at end of file diff --git a/src/client/topic/impl/write_session.h b/src/client/topic/impl/write_session.h index 5d47017f21b..f975b75ab5b 100644 --- a/src/client/topic/impl/write_session.h +++ b/src/client/topic/impl/write_session.h @@ -1,7 +1,5 @@ #pragma once -#include -#include #include #include #include @@ -11,8 +9,6 @@ #include #include -#include -#include #include namespace NYdb::inline V3::NTopic { @@ -62,411 +58,6 @@ class TWriteSession : public IWriteSession, void Start(const TDuration& delay); }; -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TKeyedWriteSession - -class TKeyedWriteSession : public IKeyedWriteSession, - public TContinuationTokenIssuer, - public std::enable_shared_from_this { -private: - using WriteSessionPtr = std::shared_ptr; - - struct TPartitionInfo { - using TSelf = TPartitionInfo; - - bool InRange(const std::string_view key) const; - bool IsSplitted() const; - - FLUENT_SETTING(std::string, FromBound); - FLUENT_SETTING(std::optional, ToBound); - FLUENT_SETTING(std::uint32_t, PartitionId); - FLUENT_SETTING(std::vector, Children); - FLUENT_SETTING_DEFAULT(bool, Locked, false); - FLUENT_SETTING_DEFAULT(NThreading::TFuture, Future, NThreading::MakeFuture()); - }; - - struct TMessageInfo { - TMessageInfo(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx); - - std::string Key; - std::string Data; - std::optional Codec; - uint32_t OriginalSize = 0; - std::optional SeqNo; - std::optional CreateTimestamp; - TMessageMeta MessageMeta; - std::optional> TxInMessage; - TTransactionBase* Tx; - std::uint32_t Partition; - bool Sent = false; - - TWriteMessage BuildMessage() const; - }; - - struct TIdleSession; - - struct TWriteSessionWrapper { - WriteSessionPtr Session; - const std::uint32_t Partition; - std::uint64_t QueueSize = 0; - std::shared_ptr IdleSession = nullptr; - - TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition); - - bool IsQueueEmpty() const; - bool AddToQueue(std::uint64_t delta); - bool RemoveFromQueue(std::uint64_t delta); - }; - - using WrappedWriteSessionPtr = std::shared_ptr; - - struct TIdleSession { - TIdleSession(TWriteSessionWrapper* session, TInstant emptySince, TDuration idleTimeout) - : Session(session) - , EmptySince(emptySince) - , IdleTimeout(idleTimeout) - {} - - const TWriteSessionWrapper* Session; - const TInstant EmptySince; - const TDuration IdleTimeout; - - bool Less(const std::shared_ptr& other) const; - bool IsExpired() const; - - struct Comparator { - bool operator()(const std::shared_ptr& first, const std::shared_ptr& second) const; - }; - }; - - using IdleSessionPtr = std::shared_ptr; - - enum class ESeqNoStrategy { - NotInitialized, - WithoutSeqNo, - WithSeqNo, - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // Custom retry policy - - struct TKeyedWriteSessionRetryPolicy : public ::IRetryPolicy { - using TSelf = TKeyedWriteSessionRetryPolicy; - using TPtr = std::shared_ptr; - - TKeyedWriteSessionRetryPolicy(TKeyedWriteSession* session); - ~TKeyedWriteSessionRetryPolicy() = default; - typename IRetryState::TPtr CreateRetryState() const override; - - private: - TKeyedWriteSession* Session; - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // Workers - - struct TEventsWorker; - - struct TSessionsWorker { - TSessionsWorker(TKeyedWriteSession* session); - WrappedWriteSessionPtr GetWriteSession(std::uint32_t partition, bool directToPartition = true); - void OnReadFromSession(WrappedWriteSessionPtr wrappedSession); - void OnWriteToSession(WrappedWriteSessionPtr wrappedSession); - void DoWork(); - - private: - void AddIdleSession(WrappedWriteSessionPtr wrappedSession, TInstant emptySince, TDuration idleTimeout); - void RemoveIdleSession(std::uint32_t partition); - WrappedWriteSessionPtr CreateWriteSession(std::uint32_t partition, bool directToPartition = true); - - using TSessionsIndexIterator = std::unordered_map::iterator; - void DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout, bool mustBeEmpty = true); - - std::string GetProducerId(std::uint32_t partitionId); - - TKeyedWriteSession* Session; - std::set IdlerSessions; - using IdlerSessionsIterator = std::set::iterator; - std::unordered_map IdlerSessionsIndex; - std::unordered_map SessionsIndex; - }; - - struct TMessagesWorker { - TMessagesWorker(TKeyedWriteSession* session); - - void DoWork(); - - void AddMessage(const std::string& key, TWriteMessage&& message, std::uint32_t partition, TTransactionBase* tx); - void ScheduleResendMessages(std::uint32_t partition, std::uint64_t afterSeqNo); - void RebuildPendingMessagesIndex(std::uint32_t partition); - void HandleAck(); - void HandleContinuationToken(std::uint32_t partition, TContinuationToken&& continuationToken); - bool IsMemoryUsageOK() const; - bool IsQueueEmpty() const; - bool HasInFlightMessages() const; - const TMessageInfo& GetFrontInFlightMessage() const; - - private: - using MessageIter = std::list::iterator; - - void PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message); - void PopInFlightMessage(); - bool SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message); - std::optional GetContinuationToken(std::uint32_t partition); - void RechoosePartitionIfNeeded(MessageIter message); - - TKeyedWriteSession* Session; - - std::list InFlightMessages; - std::unordered_map> InFlightMessagesIndex; - std::unordered_map> PendingMessagesIndex; - std::unordered_map> MessagesToResendIndex; - std::unordered_map> ContinuationTokens; - - std::uint64_t MemoryUsage = 0; - - friend class TKeyedWriteSession; - }; - - struct TSplittedPartitionWorker : public std::enable_shared_from_this { - private: - enum class EState { - Init = 0, - PendingDescribe = 1, - GotDescribe = 2, - PendingMaxSeqNo = 3, - GotMaxSeqNo = 4, - Done = 5, - }; - - void MoveTo(EState state); - void UpdateMaxSeqNo(uint64_t maxSeqNo); - void LaunchGetMaxSeqNoFutures(std::unique_lock& lock); - void HandleDescribeResult(); - - public: - TSplittedPartitionWorker(TKeyedWriteSession* session, std::uint32_t partitionId); - void DoWork(); - bool IsDone(); - bool IsInit(); - bool CanBeRemoved(); - std::string GetStateName() const; - - private: - TKeyedWriteSession* Session; - NThreading::TFuture DescribeTopicFuture; - EState State = EState::Init; - std::uint32_t PartitionId; - std::uint64_t MaxSeqNo = 0; - std::vector WriteSessions; - std::vector> GetMaxSeqNoFutures; - std::mutex Lock; - std::uint64_t NotReadyFutures = 0; - size_t Retries = 0; - TInstant DoneAt = TInstant::Max(); - }; - - struct TEventsWorker : public std::enable_shared_from_this { - enum class EEventType { - SessionClosed = 0, - ReadyToAccept = 1, - Ack = 2, - }; - - TEventsWorker(TKeyedWriteSession* session); - - std::optional> DoWork(); - NThreading::TFuture WaitEvent(); - void UnsubscribeFromPartition(std::uint32_t partition); - void SubscribeToPartition(std::uint32_t partition); - std::optional> HandleNewMessage(); - void HandleAcksEvent(std::uint64_t partition, TWriteSessionEvent::TAcksEvent&& event); - std::optional GetEvent(bool block, const std::vector& eventTypes = {}); - std::vector GetEvents(bool block, std::optional maxEventsCount = std::nullopt, const std::vector& eventTypes = {}); - std::list::iterator AckQueueBegin(std::uint32_t partition); - std::list::iterator AckQueueEnd(std::uint32_t partition); - - private: - void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition); - void HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); - bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition); - bool TransferEventsToOutputQueue(); - void AddReadyToAcceptEvent(); - bool AddSessionClosedIfNeeded(); - std::optional GetEventImpl(bool block, const std::vector& eventTypes = {}); - EEventType GetEventType(const TWriteSessionEvent::TEvent& event); - - TKeyedWriteSession* Session; - - std::unordered_set ReadyFutures; - std::unordered_map> PartitionsEventQueues; - std::list EventsOutputQueue; - std::mutex Lock; - - NThreading::TPromise EventsPromise; - NThreading::TFuture EventsFuture; - - std::optional CloseEvent; - - friend class TKeyedWriteSession; - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - // Partition chooser - - struct IPartitionChooser { - virtual std::uint32_t ChoosePartition(const std::string_view key) = 0; - virtual ~IPartitionChooser() = default; - }; - - struct TBoundPartitionChooser : IPartitionChooser { - TBoundPartitionChooser(TKeyedWriteSession* session); - std::uint32_t ChoosePartition(const std::string_view key) override; - private: - TKeyedWriteSession* Session; - }; - - struct THashPartitionChooser : IPartitionChooser { - THashPartitionChooser(std::vector&& partitions); - std::uint32_t ChoosePartition(const std::string_view key) override; - private: - std::vector Partitions; - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - - struct TMetricGauge { - std::uint64_t MetricCount = 0; - std::uint64_t Sum = 0; - std::uint64_t Max = 0; - - long double Average(); - void Add(std::uint64_t value); - std::uint64_t GetMax() const; - std::uint64_t GetSum() const; - void Clear(); - }; - - struct TMetrics { - TMetrics(TKeyedWriteSession* session); - - TMetricGauge MainWorkerTimeMs; - TMetricGauge CycleTimeMs; - TMetricGauge WriteLagMs; - TMetricGauge ContinuationTokensSent; - TMetricGauge BufferFull; - TMetricGauge IncomingMessages; - TMetricGauge OutgoingMessages; - std::mutex Lock; - TKeyedWriteSession* Session; - - void AddMainWorkerTime(std::uint64_t ms); - void AddCycleTime(std::uint64_t ms); - void AddWriteLag(std::uint64_t lagMs); - void IncContinuationTokensSent(); - void IncBufferFull(); - void IncIncomingMessages(); - void IncOutgoingMessages(); - void PrintMetrics(); - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// - - void RunMainWorker(); - - void NonBlockingClose(); - - void SetCloseDeadline(const TDuration& closeTimeout); - - TDuration GetCloseTimeout(); - - std::string GetProducerId(std::uint32_t partition); - - void HandleAutoPartitioning(std::uint32_t partition); - - bool RunSplittedPartitionWorkers(); - - void RunUserEventLoop(); - - TInstant GetCloseDeadline(); - - void GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent = std::nullopt); - - TStringBuilder LogPrefix(); - - void NextEpoch(); - -public: - TKeyedWriteSession(const TKeyedWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState); - - void Write(TContinuationToken&& continuationToken, const std::string& key, TWriteMessage&& message, - TTransactionBase* tx = nullptr) override; - - NThreading::TFuture WaitEvent() override; - - std::optional GetEvent(bool block = false) override; - - std::vector GetEvents(bool block = false, std::optional maxEventsCount = std::nullopt) override; - - bool Close(TDuration closeTimeout = TDuration::Max()) override; - - TWriterCounters::TPtr GetCounters() override; - - std::vector GetPartitions() const; - - std::unordered_map GetPartitionsMap() const; - - std::map GetPartitionsIndex() const; - - ~TKeyedWriteSession(); - -private: - std::shared_ptr Connections; - std::shared_ptr Client; - TDbDriverStatePtr DbDriverState; - - TMetrics Metrics; - - std::unordered_map Partitions; - std::map PartitionsIndex; - - TKeyedWriteSessionSettings Settings; - ESeqNoStrategy SeqNoStrategy = ESeqNoStrategy::NotInitialized; - - NThreading::TPromise ClosePromise; - NThreading::TFuture CloseFuture; - NThreading::TPromise ShutdownPromise; - NThreading::TFuture ShutdownFuture; - - std::mutex GlobalLock; - std::atomic_bool Closed = false; - std::atomic_bool Done = false; - TInstant CloseDeadline = TInstant::Now(); - - std::unique_ptr PartitionChooser; - - std::function PartitioningKeyHasher; - - std::shared_ptr EventsWorker; - std::shared_ptr SessionsWorker; - std::unordered_map> SplittedPartitionWorkers; - std::unordered_map> ReadySplittedPartitionWorkers; - std::shared_ptr MessagesWorker; - std::shared_ptr RetryPolicy; - - // TFuture::Subscribe may invoke callback synchronously when the future is already ready. - // Also, callbacks may arrive concurrently with the attempt to go idle. - // Use a small state machine to avoid re-entrancy and lost wakeups. - std::atomic MainWorkerState = 0; - std::atomic Epoch = 0; - static constexpr size_t MAX_EPOCH = 1'000'000'000; - - std::vector EventTypesWithHandlers; -}; - //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TSimpleBlockingWriteSession @@ -501,47 +92,4 @@ class TSimpleBlockingWriteSession : public ISimpleBlockingWriteSession { std::atomic_bool Closed = false; }; -//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -// TSimpleBlockingKeyedWriteSession - -class TSimpleBlockingKeyedWriteSession : public ISimpleBlockingKeyedWriteSession { -private: - std::optional GetContinuationToken(TDuration timeout); - - void HandleAcksEvent(const TWriteSessionEvent::TAcksEvent& acksEvent); - - bool WaitForAck(std::optional seqNo, TDuration timeout); - - template - bool Wait(const TDuration& timeout, F&& stopFunc); - - void RunEventLoop(); - -public: - TSimpleBlockingKeyedWriteSession( - const TKeyedWriteSessionSettings& settings, - std::shared_ptr client, - std::shared_ptr connections, - TDbDriverStatePtr dbDriverState); - - - bool Write(const std::string& key, TWriteMessage&& message, TTransactionBase* tx = nullptr, - TDuration blockTimeout = TDuration::Max()) override; - - bool Close(TDuration closeTimeout = TDuration::Max()) override; - - TWriterCounters::TPtr GetCounters() override; - -protected: - std::shared_ptr Writer; - std::unordered_set AckedSeqNos; - std::queue ContinuationTokensQueue; - - NThreading::TPromise ClosePromise; - NThreading::TFuture CloseFuture; - - std::mutex Lock; - std::atomic_bool Closed = false; -}; - } // namespace NYdb::NTopic diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index d9d4b920702..f4c9367780d 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -1,6 +1,5 @@ #include "ut_utils/topic_sdk_test_setup.h" -#include #include #include @@ -14,6 +13,7 @@ #include #include #include +#include #include #include @@ -31,9 +31,6 @@ #include #include -#include -#include -#include #include using namespace std::chrono_literals; @@ -55,6 +52,22 @@ TString SerializeDataChunk(ui64 seqNo, const TString& payload) { return result; } +TWriteMessage CreateMessage(std::string_view payload, const std::string& key, ui64 seqNo) { + TWriteMessage msg(payload); + msg.SeqNo(seqNo); + msg.Key(key); + return msg; +} + +struct TExample { + std::string Payload; + std::string Serialize() const { return Payload; } +}; + +std::string Serialize(const TExample& value) { + return value.Payload; +} + // Write a message with binary (non-UTF8) producer ID using direct tablet communication // This bypasses gRPC string validation by sending directly to the PQ tablet // The SourceId field in TCmdWrite is defined as 'bytes' in protobuf, so it supports binary data @@ -971,41 +984,26 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } - Y_UNIT_TEST(KeyedWriteSession_UserEventHandlers) { + Y_UNIT_TEST(Producer_UserEventHandlers) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.MaxBlock(TDuration::Seconds(30)); - std::atomic readyCount{0}; std::atomic acksCount{0}; std::atomic closedCount{0}; - std::atomic commonCount{0}; - - std::mutex tokensMutex; - std::condition_variable tokensCv; - std::deque readyTokens; writeSettings.EventHandlers_.HandlersExecutor(std::make_shared()); - writeSettings.EventHandlers_.ReadyToAcceptHandler( - [&](TWriteSessionEvent::TReadyToAcceptEvent& ev) { - readyCount.fetch_add(1); - { - std::lock_guard lock(tokensMutex); - readyTokens.emplace_back(std::move(ev.ContinuationToken)); - } - tokensCv.notify_one(); - }); - writeSettings.EventHandlers_.AcksHandler( [&](TWriteSessionEvent::TAcksEvent& ev) { Y_UNUSED(ev); @@ -1018,92 +1016,65 @@ Y_UNIT_TEST_SUITE(BasicUsage) { closedCount.fetch_add(1); }); - writeSettings.EventHandlers_.CommonHandler( - [&](TWriteSessionEvent::TEvent& ev) { - Y_UNUSED(ev); - commonCount.fetch_add(1); - }); - - auto getReadyToken = [&]() -> std::optional { - std::unique_lock lock(tokensMutex); - tokensCv.wait_for(lock, std::chrono::seconds(30), [&]() { return !readyTokens.empty(); }); - if (readyTokens.empty()) { - return std::nullopt; - } - auto token = std::move(readyTokens.front()); - readyTokens.pop_front(); - return token; - }; - - auto session = client.CreateKeyedWriteSession(writeSettings); + auto session = client.CreateProducer(writeSettings); const ui64 messages = 5; for (ui64 i = 0; i < messages; ++i) { - auto token = getReadyToken(); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "payload"; TWriteMessage msg(payload); msg.SeqNo(i + 1); - session->Write(std::move(*token), "key-" + ToString(i), std::move(msg)); + msg.Key("key-" + ToString(i)); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT_C(session->Close(TDuration::Seconds(30)), "Failed to close keyed write session"); - - UNIT_ASSERT_C(readyCount.load() > 0, "ReadyToAcceptHandler was not called"); - UNIT_ASSERT_C(acksCount.load() == messages, "AcksHandler does not work properly"); + UNIT_ASSERT_C(session->Close(TDuration::Seconds(30)).IsSuccess(), "Failed to close keyed write session"); + auto acks = acksCount.load(); + UNIT_ASSERT_C(acks == messages, "AcksHandler does not work properly " + std::to_string(acks)); UNIT_ASSERT_C(closedCount.load() > 0, "SessionClosedHandler was not called"); - UNIT_ASSERT_C(commonCount.load() == 0, "CommonHandler should not be called when type-specific handlers are set"); } - Y_UNIT_TEST(KeyedWriteSession_ProducerIdPrefixRequired) { + Y_UNIT_TEST(Producer_ProducerIdPrefixRequired) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 1); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - UNIT_ASSERT_EXCEPTION(setup.MakeClient().CreateKeyedWriteSession(writeSettings), TContractViolation); + UNIT_ASSERT_EXCEPTION(setup.MakeClient().CreateProducer(writeSettings), TContractViolation); } - Y_UNIT_TEST(KeyedWriteSession_SessionClosedDueToUserError) { + Y_UNIT_TEST(Producer_SessionClosedDueToUserError) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); auto publicClient = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = publicClient.CreateKeyedWriteSession(writeSettings); - TKeyedWriteSessionEventLoop eventLoop(session); - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); + auto session = publicClient.CreateProducer(writeSettings); std::string payload = "msg0"; TWriteMessage msg(payload); msg.SeqNo(0); - session->Write(std::move(*token), "key", std::move(msg)); - - auto readyToAcceptEvent = session->GetEvent(false); - UNIT_ASSERT_C(std::holds_alternative(*readyToAcceptEvent), "ReadyToAcceptEvent is not received"); - - UNIT_ASSERT_C(session->WaitEvent().Wait(TDuration::Seconds(1000)), "Timed out waiting for event"); - auto event = session->GetEvent(false); - UNIT_ASSERT_C(event, "Event is not received"); - auto sessionClosedEvent = std::get_if(&*event); - UNIT_ASSERT_C(sessionClosedEvent, "SessionClosedEvent is not received"); - UNIT_ASSERT_C(sessionClosedEvent->GetStatus() == EStatus::BAD_REQUEST, "Status is not BAD_REQUEST"); - UNIT_ASSERT(!session->Close(TDuration::Seconds(10))); + msg.Key("key"); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + auto flushResult = session->Flush().GetValueSync(); + UNIT_ASSERT_C(flushResult.IsClosed(), "Failed to flush producer"); + UNIT_ASSERT_C(flushResult.ClosedDescription->GetStatus() == EStatus::BAD_REQUEST, "Status is not BAD_REQUEST"); + UNIT_ASSERT_C(session->Close(TDuration::Seconds(10)).IsAlreadyClosed(), "Failed to close producer"); } - Y_UNIT_TEST(KeyedWriteSession_NoAutoPartitioning_HashPartitionChooser) { + Y_UNIT_TEST(Producer_NoAutoPartitioning_HashPartitionChooser) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); @@ -1118,15 +1089,16 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const ui64 partitionId0 = beforePartitions[0].GetPartitionId(); const ui64 partitionId1 = beforePartitions[1].GetPartitionId(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = publicClient.CreateKeyedWriteSession(writeSettings); + auto session = publicClient.CreateProducer(writeSettings); const std::string key0 = FindKeyForBucket(0, 2); const std::string key1 = FindKeyForBucket(1, 2); @@ -1134,28 +1106,23 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const ui64 count0 = 7; const ui64 count1 = 11; - TKeyedWriteSessionEventLoop eventLoop(session); - auto seqNo = 1; for (ui64 i = 0; i < count0; ++i) { - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); - std::string payload = "msg0"; TWriteMessage msg(payload); msg.SeqNo(seqNo++); - session->Write(std::move(*token), key0, std::move(msg)); + msg.Key(key0); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } for (ui64 i = 0; i < count1; ++i) { - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "msg1"; TWriteMessage msg(payload); msg.SeqNo(seqNo++); - session->Write(std::move(*token), key1, std::move(msg)); + msg.Key(key1); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + UNIT_ASSERT_C(session->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close keyed write session"); auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); @@ -1188,7 +1155,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { ); } - Y_UNIT_TEST(KeyedWriteSession_NoAutoPartitioning_BoundPartitionChooser) { + Y_UNIT_TEST(Producer_NoAutoPartitioning_BoundPartitionChooser) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopicWithAutoscale(TEST_TOPIC, TEST_CONSUMER, 5, 10); @@ -1200,22 +1167,21 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const auto& beforePartitions = before.GetTopicDescription().GetPartitions(); UNIT_ASSERT_VALUES_EQUAL(beforePartitions.size(), 5); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = publicClient.CreateKeyedWriteSession(writeSettings); - auto keyedSession = std::dynamic_pointer_cast(session); - const auto& partitions = keyedSession->GetPartitions(); - - TKeyedWriteSessionEventLoop eventLoop(session); + auto producer = publicClient.CreateProducer(writeSettings); + auto rawProducer = std::dynamic_pointer_cast(producer); + const auto& partitions = rawProducer->GetPartitions(); std::unordered_map keysCount; for (const auto& p : partitions) { @@ -1231,15 +1197,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } } - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "msg"; TWriteMessage msg(payload); msg.SeqNo(i + 1); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); @@ -1258,132 +1223,170 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } } - Y_UNIT_TEST(KeyedWriteSession_EventLoop_Acks) { + Y_UNIT_TEST(Producer_EventLoop_Acks) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); - TKeyedWriteSessionEventLoop eventLoop(session); + auto producer = client.CreateProducer(writeSettings); const ui64 count = 3000; for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(i); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + } + + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + } + + Y_UNIT_TEST(Producer_WriteToClosedProducer) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); + + auto client = setup.MakeClient(); + + TProducerSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(30)); + + auto producer = client.CreateProducer(writeSettings); + + const ui64 count = 10; + for (ui64 i = 1; i <= count; ++i) { + auto key = CreateGuidAsString(); + std::string payload = "data"; + TWriteMessage msg(payload); + msg.SeqNo(i); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT(eventLoop.WaitForAcks(count, TDuration::Seconds(60))); - eventLoop.CheckAcksOrder(); - UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + + std::string payload = "data"; + TWriteMessage msg(payload); + msg.SeqNo(count + 1); + msg.Key(CreateGuidAsString()); + auto writeResult = producer->Write(std::move(msg)); + UNIT_ASSERT_C(writeResult.IsError(), "Failed to write message"); + UNIT_ASSERT_C(writeResult.ErrorMessage == "producer is closed", "Error message is not correct"); + UNIT_ASSERT_C(writeResult.ClosedDescription->GetStatus() == EStatus::SUCCESS, "Status is not SUCCESS"); } - Y_UNIT_TEST(KeyedWriteSession_MultiThreadedWrite_Acks) { + Y_UNIT_TEST(Producer_MultiThreadedWrite_Acks) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); + auto producer = client.CreateProducer(writeSettings); constexpr ui64 threadsCount = 4; constexpr ui64 perThread = 25; - constexpr ui64 total = threadsCount * perThread; std::atomic nextSeqNo{1}; std::vector threads; threads.reserve(threadsCount); - TKeyedWriteSessionEventLoop eventLoop(session); - for (ui64 t = 0; t < threadsCount; ++t) { threads.emplace_back([&, t]() { auto key = TStringBuilder() << "key-" << t; for (ui64 i = 0; i < perThread; ++i) { - std::cout << "thread " << t << " writing message " << i << std::endl; - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); const ui64 seqNo = nextSeqNo.fetch_add(1); std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(seqNo); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + auto writeResult = producer->Write(std::move(msg)); + UNIT_ASSERT_C( + writeResult.IsSuccess(), + "Failed to write message" + ); } }); } - UNIT_ASSERT(eventLoop.WaitForAcks(total, TDuration::Seconds(60))); - UNIT_ASSERT(session->Close(TDuration::Seconds(10))); + for (auto& thread : threads) { + thread.join(); + } + + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); } - Y_UNIT_TEST(KeyedWriteSession_IdleSessionsTimeout) { + Y_UNIT_TEST(Producer_IdleSessionsTimeout) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(5)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); + auto producer = client.CreateProducer(writeSettings); - TKeyedWriteSessionEventLoop eventLoop(session); constexpr ui64 messages = 100; ui64 seqNo = 1; for (ui64 i = 0; i < messages; ++i) { auto key = CreateGuidAsString(); - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(seqNo++); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT(eventLoop.WaitForAcks(messages, TDuration::Seconds(60))); - eventLoop.CheckAcksOrder(); - + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); Sleep(TDuration::Seconds(6)); for (ui64 i = 0; i < messages; ++i) { auto key = CreateGuidAsString(); - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(seqNo++); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } - UNIT_ASSERT(eventLoop.WaitForAcks(messages * 2, TDuration::Seconds(60))); - eventLoop.CheckAcksOrder(); + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); } Y_UNIT_TEST(KeyedWriteSession_BoundPartitionChooser_SplitPartition_MultiThreadedAcksOrder) { @@ -1392,31 +1395,30 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); + auto producer = client.CreateProducer(writeSettings); constexpr ui64 messages = 1000; - TKeyedWriteSessionEventLoop eventLoop(session); std::jthread writer([&]() { for (ui64 i = 1; i <= messages; ++i) { - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(30)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); auto key = CreateGuidAsString(); std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(i); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } }); @@ -1429,108 +1431,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writer.join(); splitter.join(); - UNIT_ASSERT(eventLoop.WaitForAcks(messages, TDuration::Seconds(60))); - eventLoop.CheckAcksOrder(); - UNIT_ASSERT(session->Close(TDuration::Seconds(30))); - } - - Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_BasicWrite) { - TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; - setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 5); - - auto client = setup.MakeClient(); - - TKeyedWriteSessionSettings writeSettings; - writeSettings - .Path(setup.GetTopicPath(TEST_TOPIC)) - .Codec(ECodec::RAW); - writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); - - auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); - - const std::string key1 = "key1"; - const std::string key2 = "key2"; - - // Write several messages with different keys - size_t seqNo = 1; - for (int i = 0; i < 5; ++i) { - std::string payload = "message1-" + ToString(i); - TWriteMessage msg(payload); - msg.SeqNo(seqNo++); - bool res = session->Write(key1, std::move(msg)); - UNIT_ASSERT(res); - } - - for (int i = 0; i < 5; ++i) { - std::string payload = "message2-" + ToString(i); - TWriteMessage msg(payload); - msg.SeqNo(seqNo++); - bool res = session->Write(key2, std::move(msg)); - UNIT_ASSERT(res); - } - - UNIT_ASSERT(session->Close(TDuration::Seconds(10))); - } - - Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_NoSeqNo) { - TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; - setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); - - auto client = setup.MakeClient(); - - TKeyedWriteSessionSettings writeSettings; - writeSettings - .Path(setup.GetTopicPath(TEST_TOPIC)) - .Codec(ECodec::RAW); - writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); - - auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); - - const ui64 messages = 10; - for (ui64 i = 0; i < messages; ++i) { - std::string payload = "payload-" + ToString(i); - TWriteMessage msg(payload); - bool res = session->Write("key-" + ToString(i % 3), std::move(msg)); - UNIT_ASSERT(res); - } - - bool closeRes = session->Close(TDuration::Seconds(30)); - UNIT_ASSERT(closeRes); - } - - Y_UNIT_TEST(SimpleBlockingKeyedWriteSession_ManyMessages) { - TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; - setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); - - auto client = setup.MakeClient(); - - TKeyedWriteSessionSettings writeSettings; - writeSettings - .Path(setup.GetTopicPath(TEST_TOPIC)) - .Codec(ECodec::RAW); - writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); - - auto session = client.CreateSimpleBlockingKeyedWriteSession(writeSettings); - - ui64 seqNo = 1; - - for (ui64 i = 0; i < 1000; ++i) { - auto key = CreateGuidAsString(); - std::string payload = "payload-" + ToString(seqNo); - TWriteMessage msg(payload); - msg.SeqNo(seqNo++); - bool res = session->Write(key, std::move(msg)); - UNIT_ASSERT(res); - } - - bool closeRes = session->Close(TDuration::Seconds(60)); - UNIT_ASSERT(closeRes); + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(30)).IsSuccess(), "Failed to close producer"); } Y_UNIT_TEST(KeyedWriteSession_CloseTimeout) { @@ -1539,31 +1441,36 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto client = setup.MakeClient(); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); - - TKeyedWriteSessionEventLoop eventLoop(session); + auto producer = client.CreateProducer(writeSettings); for (int i = 0; i < 1000; ++i) { - auto token = eventLoop.GetContinuationToken(TDuration::Seconds(10)); - UNIT_ASSERT_C(token, "Timed out waiting for ReadyToAcceptEvent"); std::string payload = "message-" + ToString(i); TWriteMessage msg(payload); msg.SeqNo(i + 1); - session->Write(std::move(*token), "key1", std::move(msg)); + msg.Key("key1"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); } // Test Close timeout - const TDuration closeTimeout = TDuration::Seconds(2); + const TDuration closeTimeout = TDuration::Seconds(5); const TInstant startTime = TInstant::Now(); - session->Close(closeTimeout); + auto result = producer->Close(closeTimeout); + // Close may legitimately return Timeout if there are still in-flight messages + // at the moment of deadline. For this test we only require that Close doesn't + // block longer than the timeout and that the session eventually closes successfully. + UNIT_ASSERT_C( + result.IsSuccess() || result.IsTimeout(), + TStringBuilder() << "Failed to close keyed write session, status: " << static_cast(result.Status) + ); const TDuration actualDuration = TInstant::Now() - startTime; // Verify that Close didn't block longer than timeout (with some tolerance) @@ -1572,75 +1479,33 @@ Y_UNIT_TEST_SUITE(BasicUsage) { actualDuration <= maxExpectedDuration + maxExpectedDuration / 10, TStringBuilder() << "Close() took " << actualDuration << " but timeout was " << closeTimeout ); - - int attempts = 0; - constexpr int maxAttempts = 1100; - for (attempts = 0; attempts < maxAttempts; ++attempts) { - auto event = session->GetEvent(false); - if (!event) { - break; - } - - auto sessionClosedEvent = std::get_if(&*event); - if (!sessionClosedEvent) { - continue; - } - - UNIT_ASSERT(sessionClosedEvent->IsSuccess()); - break; - } - - UNIT_ASSERT(attempts < maxAttempts); } - Y_UNIT_TEST(AutoPartitioning_KeyedWriteSession) { + Y_UNIT_TEST(AutoPartitioning_Producer) { auto settings = TTopicSdkTestSetup::MakeServerSettings(); settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; TTopicClient client = setup.MakeClient(); - std::queue readyTokens1; - std::queue readyTokens2; - std::optional sessionClosedEvent; - std::unordered_set ackedSeqNos; - bool closed = false; - - auto createMessage = [](std::string_view payload, ui64 seqNo) -> TWriteMessage { - TWriteMessage msg(payload); - msg.SeqNo(seqNo); - return msg; - }; - - TCreateTopicSettings createSettings; - createSettings - .BeginConfigurePartitioningSettings() - .MinActivePartitions(2) - .MaxActivePartitions(100) - .BeginConfigureAutoPartitioningSettings() - .UpUtilizationPercent(2) - .DownUtilizationPercent(1) - .StabilizationWindow(TDuration::Seconds(2)) - .Strategy(EAutoPartitioningStrategy::ScaleUp) - .EndConfigureAutoPartitioningSettings() - .EndConfigurePartitioningSettings(); - client.CreateTopic(TEST_TOPIC, createSettings).Wait(); + CreateTopicWithAutoPartitioning(client); auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync(); UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 2); - TKeyedWriteSessionSettings writeSettings1; + TProducerSettings writeSettings1; writeSettings1 .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings1.ProducerIdPrefix("autopartitioning_keyed_1"); - writeSettings1.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings1.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings1.MaxBlock(TDuration::Seconds(30)); - TKeyedWriteSessionSettings writeSettings2 = writeSettings1; + TProducerSettings writeSettings2 = writeSettings1; writeSettings2.ProducerIdPrefix("autopartitioning_keyed_2"); - auto session1 = client.CreateKeyedWriteSession(writeSettings1); - auto session2 = client.CreateKeyedWriteSession(writeSettings2); + auto producer1 = client.CreateProducer(writeSettings1); + auto producer2 = client.CreateProducer(writeSettings2); auto msgData = TString(1_MB, 'a'); std::vector keys; @@ -1648,94 +1513,36 @@ Y_UNIT_TEST_SUITE(BasicUsage) { keys.push_back(partition.GetFromBound().value_or("")); } - auto getQueue = [&](const std::shared_ptr& s) -> std::queue& { - if (s == session1) { - return readyTokens1; - } - if (s == session2) { - return readyTokens2; - } - Y_ABORT("Unknown session pointer in AutoPartitioning_KeyedWriteSession"); - }; - - auto eventLoop = [&](std::shared_ptr s) { - while (true) { - auto event = s->GetEvent(false); - if (!event) { - break; - } - if (auto* ready = std::get_if(&*event)) { - getQueue(s).push(std::move(ready->ContinuationToken)); - continue; - } - if (auto* closedEv = std::get_if(&*event)) { - sessionClosedEvent = std::move(*closedEv); - closed = true; - break; - } - if (auto* acks = std::get_if(&*event)) { - for (const auto& ack : acks->Acks) { - UNIT_ASSERT_C( - ackedSeqNos.insert(ack.SeqNo).second, - "Duplicate ack for seqNo " << ack.SeqNo); - } - } - } - }; - - auto getReadyToken = [&](std::shared_ptr s) -> std::optional { - auto& q = getQueue(s); - while (q.empty() && !closed) { - s->WaitEvent().Wait(TDuration::Seconds(5)); - eventLoop(s); - } - if (q.empty()) { - return std::nullopt; - } - auto t = std::move(q.front()); - q.pop(); - return t; - }; - - auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { - auto token = getReadyToken(s); - UNIT_ASSERT(token); + auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { auto key = keys[seqNo % keys.size()]; if (key.empty()) { key = "lalala"; } - s->Write(std::move(*token), key, createMessage(payload, seqNo)); + UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsSuccess(), "Failed to write message"); }; { - writeMessage(session1, msgData, 1); - writeMessage(session1, msgData, 2); + writeMessage(producer1, msgData, 1); + writeMessage(producer1, msgData, 2); Sleep(TDuration::Seconds(5)); auto d = client.DescribeTopic(TEST_TOPIC).GetValueSync(); UNIT_ASSERT_EQUAL(d.GetTopicDescription().GetPartitions().size(), 2); } { - writeMessage(session1, msgData, 3); - writeMessage(session1, msgData, 4); - writeMessage(session1, msgData, 5); - writeMessage(session1, msgData, 6); - writeMessage(session1, msgData, 7); - writeMessage(session2, msgData, 8); - writeMessage(session1, msgData, 9); - writeMessage(session1, msgData, 10); - writeMessage(session2, msgData, 11); - writeMessage(session1, msgData, 12); - Sleep(TDuration::Seconds(30)); - for (int i = 0; i < 50 && ackedSeqNos.size() < 12 && !closed; ++i) { - eventLoop(session1); - eventLoop(session2); - if (ackedSeqNos.size() < 12) { - Sleep(TDuration::MilliSeconds(200)); - } - } - UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), 12, - "Expected exactly 12 distinct acks, each seqNo exactly once; got " << ackedSeqNos.size()); + writeMessage(producer1, msgData, 3); + writeMessage(producer1, msgData, 4); + writeMessage(producer1, msgData, 5); + writeMessage(producer1, msgData, 6); + writeMessage(producer1, msgData, 7); + writeMessage(producer2, msgData, 8); + writeMessage(producer1, msgData, 9); + writeMessage(producer1, msgData, 10); + writeMessage(producer2, msgData, 11); + writeMessage(producer1, msgData, 12); + UNIT_ASSERT_C(producer1->Flush().GetValueSync().IsSuccess(), "Failed to flush producer1"); + UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer2"); + Sleep(TDuration::Seconds(5)); } auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); @@ -1743,23 +1550,17 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(partitionsCount >= 4, TStringBuilder() << "Partitions count: " << partitionsCount << ", expected at least 4"); - writeMessage(session1, msgData, 13); - writeMessage(session1, msgData, 14); - Sleep(TDuration::Seconds(20)); - for (int i = 0; i < 50 && ackedSeqNos.size() < 14 && !closed; ++i) { - eventLoop(session1); - eventLoop(session2); - if (ackedSeqNos.size() < 14) { - Sleep(TDuration::MilliSeconds(200)); - } - } + writeMessage(producer1, msgData, 13); + writeMessage(producer1, msgData, 14); + UNIT_ASSERT_C(producer1->Flush().GetValueSync().IsSuccess(), "Failed to flush producer1"); + UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer2"); - TKeyedWriteSessionSettings writeSettings3 = writeSettings1; + TProducerSettings writeSettings3 = writeSettings1; writeSettings3.ProducerIdPrefix("autopartitioning_keyed_3"); - auto session3 = client.CreateKeyedWriteSession(writeSettings3); + auto producer3 = client.CreateProducer(writeSettings3); - auto partitionsMap1 = dynamic_cast(session1.get())->GetPartitionsMap(); - auto partitionsMap3 = dynamic_cast(session3.get())->GetPartitionsMap(); + auto partitionsMap1 = dynamic_cast(producer1.get())->GetPartitionsMap(); + auto partitionsMap3 = dynamic_cast(producer3.get())->GetPartitionsMap(); for (const auto& [partitionId, partitionInfo] : partitionsMap1) { auto partitionInfo3 = partitionsMap3.find(partitionId); @@ -1771,8 +1572,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { "To bound is not equal for partition " << partitionId); } - auto partitionsIndex1 = dynamic_cast(session1.get())->GetPartitionsIndex(); - auto partitionsIndex3 = dynamic_cast(session3.get())->GetPartitionsIndex(); + auto partitionsIndex1 = dynamic_cast(producer1.get())->GetPartitionsIndex(); + auto partitionsIndex3 = dynamic_cast(producer3.get())->GetPartitionsIndex(); for (const auto& [key, partitionId] : partitionsIndex1) { auto partitionId3 = partitionsIndex3.find(key); UNIT_ASSERT_C(partitionId3 != partitionsIndex3.end(), @@ -1781,64 +1582,40 @@ Y_UNIT_TEST_SUITE(BasicUsage) { "Partition id is not equal for key " << key); } - UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), 14, - "Expected exactly 14 distinct acks, each seqNo exactly once; got " << ackedSeqNos.size()); - auto sessionPartitions = dynamic_cast(session1.get())->GetPartitions(); + auto sessionPartitions = dynamic_cast(producer1.get())->GetPartitions(); UNIT_ASSERT_EQUAL_C(sessionPartitions.size(), partitionsCount, "Expected exactly" << partitionsCount << " partitions, actual: " << sessionPartitions.size()); - UNIT_ASSERT(session1->Close(TDuration::Seconds(30))); - UNIT_ASSERT(session2->Close(TDuration::Seconds(30))); + UNIT_ASSERT(producer1->Close(TDuration::Seconds(30)).IsSuccess()); + UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); + UNIT_ASSERT(producer3->Close(TDuration::Seconds(30)).IsSuccess()); } - Y_UNIT_TEST(AutoPartitioning_KeyedWriteSession_SmallMessages) { + Y_UNIT_TEST(AutoPartitioning_Producer_SmallMessages) { auto settings = TTopicSdkTestSetup::MakeServerSettings(); settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; TTopicClient client = setup.MakeClient(); - std::queue readyTokens1; - std::queue readyTokens2; - std::optional sessionClosedEvent; - std::unordered_set ackedSeqNos; - bool closed = false; - - auto createMessage = [](std::string_view payload, ui64 seqNo) -> TWriteMessage { - TWriteMessage msg(payload); - msg.SeqNo(seqNo); - return msg; - }; - - TCreateTopicSettings createSettings; - createSettings - .BeginConfigurePartitioningSettings() - .MinActivePartitions(2) - .MaxActivePartitions(100) - .BeginConfigureAutoPartitioningSettings() - .UpUtilizationPercent(2) - .DownUtilizationPercent(1) - .StabilizationWindow(TDuration::Seconds(2)) - .Strategy(EAutoPartitioningStrategy::ScaleUp) - .EndConfigureAutoPartitioningSettings() - .EndConfigurePartitioningSettings(); - client.CreateTopic(TEST_TOPIC, createSettings).Wait(); + CreateTopicWithAutoPartitioning(client); auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync(); UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 2); - TKeyedWriteSessionSettings writeSettings1; + TProducerSettings writeSettings1; writeSettings1 .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings1.ProducerIdPrefix("autopartitioning_keyed_small_1"); - writeSettings1.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Bound); + writeSettings1.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings1.MaxBlock(TDuration::Seconds(30)); - TKeyedWriteSessionSettings writeSettings2 = writeSettings1; + TProducerSettings writeSettings2 = writeSettings1; writeSettings2.ProducerIdPrefix("autopartitioning_keyed_small_2"); - auto session1 = client.CreateKeyedWriteSession(writeSettings1); - auto session2 = client.CreateKeyedWriteSession(writeSettings2); + auto producer1 = client.CreateProducer(writeSettings1); + auto producer2 = client.CreateProducer(writeSettings2); const size_t msgSize = 256_KB; auto msgData = TString(msgSize, 'a'); const ui64 totalMessages = 44; @@ -1848,57 +1625,15 @@ Y_UNIT_TEST_SUITE(BasicUsage) { keys.push_back(partition.GetFromBound().value_or("")); } - auto getQueue = [&](const std::shared_ptr& s) -> std::queue& { - if (s == session1) return readyTokens1; - if (s == session2) return readyTokens2; - Y_ABORT("Unknown session pointer in AutoPartitioning_KeyedWriteSession_SmallMessages"); - }; - - auto eventLoop = [&](std::shared_ptr s) { - while (true) { - auto event = s->GetEvent(false); - if (!event) break; - if (auto* ready = std::get_if(&*event)) { - getQueue(s).push(std::move(ready->ContinuationToken)); - continue; - } - if (auto* closedEv = std::get_if(&*event)) { - sessionClosedEvent = std::move(*closedEv); - closed = true; - break; - } - if (auto* acks = std::get_if(&*event)) { - for (const auto& ack : acks->Acks) { - UNIT_ASSERT_C(ackedSeqNos.insert(ack.SeqNo).second, - "Duplicate ack for seqNo " << ack.SeqNo); - } - } - } - }; - - auto getReadyToken = [&](std::shared_ptr s) -> std::optional { - auto& q = getQueue(s); - while (q.empty() && !closed) { - s->WaitEvent().Wait(TDuration::Seconds(5)); - eventLoop(s); - } - if (q.empty()) return std::nullopt; - auto t = std::move(q.front()); - q.pop(); - return t; - }; - - auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { - auto token = getReadyToken(s); - UNIT_ASSERT(token); + auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { auto key = keys[seqNo % keys.size()]; if (key.empty()) key = "a"; - s->Write(std::move(*token), key, createMessage(payload, seqNo)); + UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsSuccess(), "Failed to write message"); }; { - writeMessage(session1, msgData, 1); - writeMessage(session1, msgData, 2); + writeMessage(producer1, msgData, 1); + writeMessage(producer1, msgData, 2); Sleep(TDuration::Seconds(5)); auto d = client.DescribeTopic(TEST_TOPIC).GetValueSync(); UNIT_ASSERT_EQUAL(d.GetTopicDescription().GetPartitions().size(), 2); @@ -1906,17 +1641,12 @@ Y_UNIT_TEST_SUITE(BasicUsage) { { for (ui64 seq = 3; seq <= totalMessages - 2; ++seq) { - auto s = (seq % 4 == 0) ? session2 : session1; + auto s = (seq % 4 == 0) ? producer2 : producer1; writeMessage(s, msgData, seq); } - Sleep(TDuration::Seconds(30)); - for (int i = 0; i < 80 && ackedSeqNos.size() < totalMessages - 2 && !closed; ++i) { - eventLoop(session1); - eventLoop(session2); - if (ackedSeqNos.size() < totalMessages - 2) Sleep(TDuration::MilliSeconds(200)); - } - UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), totalMessages - 2, - "Expected " << totalMessages - 2 << " acks; got " << ackedSeqNos.size()); + Sleep(TDuration::Seconds(5)); + UNIT_ASSERT_C(producer1->Flush().GetValueSync().IsSuccess(), "Failed to flush producer1"); + UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer2"); } auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); @@ -1924,25 +1654,308 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(partitionsCount >= 3, TStringBuilder() << "Partitions count: " << partitionsCount << ", expected at least 3 (auto-partitioning)"); - writeMessage(session1, msgData, totalMessages - 1); - writeMessage(session1, msgData, totalMessages); - Sleep(TDuration::Seconds(20)); - for (int i = 0; i < 80 && ackedSeqNos.size() < totalMessages && !closed; ++i) { - eventLoop(session1); - eventLoop(session2); - if (ackedSeqNos.size() < totalMessages) Sleep(TDuration::MilliSeconds(200)); - } + writeMessage(producer1, msgData, totalMessages - 1); + writeMessage(producer1, msgData, totalMessages); + UNIT_ASSERT_C(producer1->Flush().GetValueSync().IsSuccess(), "Failed to flush producer1"); + UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer2"); - UNIT_ASSERT_EQUAL_C(ackedSeqNos.size(), totalMessages, - "Expected " << totalMessages << " acks; got " << ackedSeqNos.size()); - auto sessionPartitions = dynamic_cast(session1.get())->GetPartitions(); + auto sessionPartitions = dynamic_cast(producer1.get())->GetPartitions(); UNIT_ASSERT_EQUAL_C(sessionPartitions.size(), partitionsCount, "Session partitions " << sessionPartitions.size() << " != topic partitions " << partitionsCount); - UNIT_ASSERT(session1->Close(TDuration::Seconds(30))); - UNIT_ASSERT(session2->Close(TDuration::Seconds(30))); + UNIT_ASSERT(producer1->Close(TDuration::Seconds(30)).IsSuccess()); + UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); } + + Y_UNIT_TEST(Producer_BasicWrite) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 10); + + TProducerSettings writeSettings; + writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); + writeSettings.Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix("producer_basic_write"); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + + auto producer = client.CreateProducer(writeSettings); + auto msgData = TString(10_KB, 'a'); + + for (ui64 i = 0; i < 100; ++i) { + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + } + + UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); + + auto describe = client.DescribeTopic(TEST_TOPIC, TDescribeTopicSettings().IncludeStats(true)).GetValueSync(); + UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 10); + + ui64 messagesWritten = 0; + for (const auto& partition : describe.GetTopicDescription().GetPartitions()) { + auto stats = partition.GetPartitionStats(); + UNIT_ASSERT(stats); + messagesWritten += stats->GetEndOffset() - stats->GetStartOffset(); + } + UNIT_ASSERT_EQUAL(messagesWritten, 100); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(1)).IsSuccess(), "Failed to close producer"); + } + + Y_UNIT_TEST(TypedProducer_BasicWrite) { + constexpr ui64 messageCount = 100; + + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 10); + + TProducerSettings writeSettings; + writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); + writeSettings.Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix("producer_basic_write"); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.KeyProducer([](const TWriteMessage& message) -> std::string { + return ToString(MurmurHash(message.Data.data(), message.Data.size())); + }); + writeSettings.MaxBlock(TDuration::Seconds(1)); + + auto producer = client.CreateTypedProducer(writeSettings); + + std::vector sentPayloads; + sentPayloads.reserve(messageCount); + for (ui64 i = 0; i < messageCount; ++i) { + auto payload = CreateGuidAsString(); + sentPayloads.push_back(payload); + UNIT_ASSERT(producer->Write(TExample{.Payload = payload}).IsSuccess()); + } + + UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().MessagesWritten, messageCount); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(1)).IsSuccess(), "Failed to close producer"); + + std::vector receivedPayloads; + receivedPayloads.reserve(messageCount); + NThreading::TPromise allReadPromise = NThreading::NewPromise(); + + auto readSettings = TReadSessionSettings() + .ConsumerName(setup.GetConsumerName()) + .AppendTopics(setup.GetTopicPath(TEST_TOPIC)); + + readSettings.EventHandlers_.SimpleDataHandlers([&](TReadSessionEvent::TDataReceivedEvent& ev) { + for (auto& msg : ev.GetMessages()) { + receivedPayloads.push_back(std::string(msg.GetData())); + msg.Commit(); + } + if (receivedPayloads.size() >= messageCount) { + allReadPromise.SetValue(); + } + }); + + auto readSession = client.CreateReadSession(readSettings); + UNIT_ASSERT(allReadPromise.GetFuture().Wait(TDuration::Seconds(30))); + readSession->Close(TDuration::Seconds(5)); + + UNIT_ASSERT_VALUES_EQUAL(receivedPayloads.size(), messageCount); + + std::sort(sentPayloads.begin(), sentPayloads.end()); + std::sort(receivedPayloads.begin(), receivedPayloads.end()); + UNIT_ASSERT_VALUES_EQUAL(sentPayloads, receivedPayloads); + } + + Y_UNIT_TEST(Producer_SmallSessionIdleTimeout) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 5); + + TProducerSettings writeSettings; + writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); + writeSettings.Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix("producer_basic_write"); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::MilliSeconds(500)); + writeSettings.MaxBlock(TDuration::Seconds(1)); + + auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); + const auto& partitions = describeResult.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_EQUAL(partitions.size(), 5); + + auto producer = client.CreateProducer(writeSettings); + auto producerRaw = dynamic_cast(producer.get()); + auto msgData = TString(10_KB, 'a'); + + for (ui64 i = 0; i < 3; ++i) { + for (const auto& partition : partitions) { + for (ui64 i = 0; i < 10; ++i) { + TWriteMessage msg(msgData); + msg.Partition(partition.GetPartitionId()); + UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); + } + UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); + UNIT_ASSERT((producerRaw->GetIdleSessionsCount() == 1 && producerRaw->GetSessionsCount() == 1) || + (producerRaw->GetIdleSessionsCount() == 0 && producerRaw->GetSessionsCount() == 0)); + Sleep(TDuration::Seconds(1)); + } + } + + { + auto describeResult = client.DescribeTopic(TEST_TOPIC, TDescribeTopicSettings().IncludeStats(true)).GetValueSync(); + ui64 messagesWritten = 0; + for (const auto& partition : describeResult.GetTopicDescription().GetPartitions()) { + auto stats = partition.GetPartitionStats(); + UNIT_ASSERT(stats); + messagesWritten += stats->GetEndOffset() - stats->GetStartOffset(); + } + UNIT_ASSERT_EQUAL(messagesWritten, 150); + } + UNIT_ASSERT(producer->Close(TDuration::Seconds(1)).IsSuccess()); + } + + Y_UNIT_TEST(Producer_CustomKeyProducerFunction) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); + + // Capture partition ids in the same order as DescribeTopic returns them + // (the keyed session uses the same DescribeTopic ordering to map hash bucket -> partition id). + auto publicClient = setup.MakeClient(); + auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); + auto before = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + UNIT_ASSERT_C(before.IsSuccess(), before.GetIssues().ToOneLineString()); + const auto& beforePartitions = before.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_VALUES_EQUAL(beforePartitions.size(), 2); + const ui64 partitionId0 = beforePartitions[0].GetPartitionId(); + const ui64 partitionId1 = beforePartitions[1].GetPartitionId(); + + constexpr auto keyAttributeName = "__key"; + + TProducerSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + writeSettings.KeyProducer([](const TWriteMessage& message) -> std::string { + for (const auto& [attributeName, attributeValue] : message.MessageMeta_) { + if (attributeName == keyAttributeName) { + return attributeValue; + } + } + return ""; + }); + writeSettings.MaxBlock(TDuration::Seconds(1)); + + auto producer = publicClient.CreateProducer(writeSettings); + + const std::string key0 = FindKeyForBucket(0, 2); + const std::string key1 = FindKeyForBucket(1, 2); + + const ui64 count0 = 7; + const ui64 count1 = 11; + + auto seqNo = 1; + for (ui64 i = 0; i < count0; ++i) { + std::string payload = "msg0"; + TWriteMessage msg(payload); + msg.SeqNo(seqNo++); + msg.MessageMeta_.emplace_back(keyAttributeName, key0); + UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); + } + for (ui64 i = 0; i < count1; ++i) { + std::string payload = "msg1"; + TWriteMessage msg(payload); + msg.SeqNo(seqNo++); + msg.MessageMeta_.emplace_back(keyAttributeName, key1); + UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); + } + + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().MessagesWritten, count0 + count1); + UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().LastWrittenSeqNo, seqNo - 1); + + auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); + UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); + const auto& afterPartitions = after.GetTopicDescription().GetPartitions(); + UNIT_ASSERT_VALUES_EQUAL(afterPartitions.size(), 2); + + std::unordered_map endOffsets; + for (const auto& p : afterPartitions) { + auto stats = p.GetPartitionStats(); + UNIT_ASSERT(stats.has_value()); + endOffsets[p.GetPartitionId()] = stats->GetEndOffset(); + } + + auto it0 = endOffsets.find(partitionId0); + auto it1 = endOffsets.find(partitionId1); + UNIT_ASSERT(it0 != endOffsets.end()); + UNIT_ASSERT(it1 != endOffsets.end()); + + const ui64 endOffset0 = it0->second; + const ui64 endOffset1 = it1->second; + + // Partition ordering in DescribeTopic is not a part of public API contract, so allow swapping. + UNIT_ASSERT_VALUES_EQUAL(endOffset0 + endOffset1, count0 + count1); + UNIT_ASSERT_C( + (endOffset0 == count0 && endOffset1 == count1) || (endOffset0 == count1 && endOffset1 == count0), + TStringBuilder() << "Unexpected end offsets distribution: " + << "partitionId0=" << partitionId0 << " endOffset0=" << endOffset0 << ", " + << "partitionId1=" << partitionId1 << " endOffset1=" << endOffset1 << ", " + << "expected (" << count0 << "," << count1 << ") in any order" + ); + } + + Y_UNIT_TEST(Producer_BlockingWrite) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 10); + + TProducerSettings writeSettings; + writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); + writeSettings.Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix("simple_blocking_producer_basic_write"); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(1)); + + auto producer = client.CreateProducer(writeSettings); + auto msgData = TString(10_KB, 'a'); + + for (ui64 i = 0; i < 100; ++i) { + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + } + + UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().MessagesWritten, 100); + UNIT_ASSERT(producer->Close(TDuration::Seconds(1)).IsSuccess()); + } + + Y_UNIT_TEST(Producer_TimeoutError) { + auto settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); + TTopicSdkTestSetup setup{TEST_CASE_NAME, settings, false}; + TTopicClient client = setup.MakeClient(); + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 1); + + TProducerSettings writeSettings; + writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); + writeSettings.Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix("simple_blocking_producer_basic_write"); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxMemoryUsage(100_KB); + writeSettings.MaxBlock(TDuration::MilliSeconds(1)); + + auto producer = client.CreateProducer(writeSettings); + auto msgData = TString(1_MB, 'a'); + + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsTimeout()); + UNIT_ASSERT(producer->Close(TDuration::Seconds(10)).IsSuccess()); + } } // Y_UNIT_TEST_SUITE(BasicUsage) } // namespace diff --git a/src/client/topic/ut/ut_utils/event_loop.cpp b/src/client/topic/ut/ut_utils/event_loop.cpp deleted file mode 100644 index 233b2863e71..00000000000 --- a/src/client/topic/ut/ut_utils/event_loop.cpp +++ /dev/null @@ -1,93 +0,0 @@ -#include "event_loop.h" -#include - -namespace NYdb::inline V3::NTopic::NTests { - -TKeyedWriteSessionEventLoop::TKeyedWriteSessionEventLoop(std::shared_ptr session) - : Session_(std::move(session)) -{} - -void TKeyedWriteSessionEventLoop::Run() { - while (true) { - auto event = Session_->GetEvent(false); - if (!event) { - break; - } - if (auto* ready = std::get_if(&*event)) { - std::lock_guard lk(Lock_); - ReadyTokens_.push(std::move(ready->ContinuationToken)); - continue; - } - if (std::get_if(&*event)) { - break; - } - if (auto* acks = std::get_if(&*event)) { - std::lock_guard lk(Lock_); - for (const auto& ack : acks->Acks) { - auto [it, inserted] = AckedSeqNos_.insert(ack.SeqNo); - UNIT_ASSERT_C(inserted, TStringBuilder() << "Ack already received: " << ack.SeqNo); - AckOrder_.push_back(ack.SeqNo); - } - } - } -} - -std::optional TKeyedWriteSessionEventLoop::GetContinuationToken(TDuration timeout) { - const TInstant deadline = TInstant::Now() + timeout; - while (TInstant::Now() < deadline) { - { - std::lock_guard lock(Lock_); - if (!ReadyTokens_.empty()) { - auto token = std::move(ReadyTokens_.front()); - ReadyTokens_.pop(); - return token; - } - } - Session_->WaitEvent().Wait(deadline); - Run(); - } - - std::lock_guard lock(Lock_); - if (!ReadyTokens_.empty()) { - auto token = std::move(ReadyTokens_.front()); - ReadyTokens_.pop(); - return token; - } - return std::nullopt; -} - -bool TKeyedWriteSessionEventLoop::WaitForAcks(size_t count, TDuration timeout) { - const TInstant deadline = TInstant::Now() + timeout; - while (TInstant::Now() < deadline) { - { - std::lock_guard lock(Lock_); - if (AckedSeqNos_.size() >= count) { - return true; - } - } - Session_->WaitEvent().Wait(deadline); - Run(); - } - return false; -} - -void TKeyedWriteSessionEventLoop::CheckAcksOrder() { - std::lock_guard lock(Lock_); - size_t expectedAck = 1; - UNIT_ASSERT_C(AckedSeqNos_.size() == AckOrder_.size(), TStringBuilder() << "Unexpected number of acks: got " << AckOrder_.size() << ", expected " << AckedSeqNos_.size()); - size_t index = 0; - for (const auto& ack : AckOrder_) { - TStringBuilder sb; - if (ack != expectedAck) { - for (size_t i = std::min(size_t(0), index - 10); i < std::min(index + 10, AckOrder_.size()); i++) { - sb << "Ack " << i << ": " << AckOrder_[i] << " "; - } - - sb << "Unexpected ack order: got " << ack << ", expected " << expectedAck; - } - UNIT_ASSERT_VALUES_EQUAL_C(ack, expectedAck, sb); - expectedAck++; - } -} - -} // namespace NYdb::inline V3::NTopic::NTests diff --git a/src/client/topic/ut/ut_utils/event_loop.h b/src/client/topic/ut/ut_utils/event_loop.h deleted file mode 100644 index a1fa6ca60f6..00000000000 --- a/src/client/topic/ut/ut_utils/event_loop.h +++ /dev/null @@ -1,37 +0,0 @@ -#pragma once - -#include - -#include - -#include -#include -#include -#include - -namespace NYdb::inline V3::NTopic::NTests { - -//! Helper for keyed write session tests: runs event loop and provides continuation tokens. -class TKeyedWriteSessionEventLoop { -public: - explicit TKeyedWriteSessionEventLoop(std::shared_ptr session); - - //! Block until a continuation token is available or timeout. Calls Run() while waiting. - //! Returns nullopt on timeout or if session was closed before a token appeared. - std::optional GetContinuationToken(TDuration timeout); - - bool WaitForAcks(size_t count, TDuration timeout); - void CheckAcksOrder(); - -private: - //! Process all currently available events. Returns true if SessionClosed was seen. - void Run(); - - std::shared_ptr Session_; - std::queue ReadyTokens_; - std::unordered_set AckedSeqNos_; - std::vector AckOrder_; - std::mutex Lock_; -}; - -} // namespace NYdb::inline V3::NTopic::NTests diff --git a/tests/integration/topic/basic_usage_it.cpp b/tests/integration/topic/basic_usage_it.cpp index d4e05d0ce47..d8f63572e1b 100644 --- a/tests/integration/topic/basic_usage_it.cpp +++ b/tests/integration/topic/basic_usage_it.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -40,83 +41,6 @@ std::uint64_t TSimpleWriteSessionTestAdapter::GetAcquiredMessagesCount() const { return 0; } -class TKeyedWriteSessionTestAdapter { -public: - TKeyedWriteSessionTestAdapter(NTopic::IKeyedWriteSession* session); - - void WaitForAcks(size_t count, TDuration timeout); - std::optional GetContinuationToken(TDuration timeout); - size_t GetAckedSeqNosCount() const; - bool ValidateAcksOrder() const; - -private: - void RunEventLoop(TDuration timeout, size_t stopOnAcksCount, bool stopOnContinuationToken = false); - - NTopic::IKeyedWriteSession* Session; - std::queue tokens; - std::vector ackedSeqNos; -}; - -TKeyedWriteSessionTestAdapter::TKeyedWriteSessionTestAdapter(NTopic::IKeyedWriteSession* session) - : Session(session) -{} - -void TKeyedWriteSessionTestAdapter::WaitForAcks(size_t count, TDuration timeout) { - RunEventLoop(timeout, count, false); -} - -bool TKeyedWriteSessionTestAdapter::ValidateAcksOrder() const { - size_t expectedSeqNo = 1; - for (const auto& seqNo : ackedSeqNos) { - if (seqNo != expectedSeqNo) { - return false; - } - expectedSeqNo++; - } - return true; -} - -size_t TKeyedWriteSessionTestAdapter::GetAckedSeqNosCount() const { - return ackedSeqNos.size(); -} - -std::optional TKeyedWriteSessionTestAdapter::GetContinuationToken(TDuration timeout) { - RunEventLoop(timeout, 0, true); - if (tokens.empty()) { - return std::nullopt; - } - auto token = std::move(tokens.front()); - tokens.pop(); - return token; -} - -void TKeyedWriteSessionTestAdapter::RunEventLoop(TDuration timeout, size_t stopOnAcksCount, bool stopOnContinuationToken) { - auto deadline = TInstant::Now() + timeout; - while (TInstant::Now() < deadline) { - Session->WaitEvent().Wait(deadline); - auto event = Session->GetEvent(false); - if (!event) { - continue; - } - if (auto ev = std::get_if(&*event)) { - tokens.push(std::move(ev->ContinuationToken)); - if (stopOnContinuationToken) { - return; - } - continue; - } - if (auto ev = std::get_if(&*event)) { - for (const auto& ack : ev->Acks) { - ackedSeqNos.push_back(ack.SeqNo); - } - - if (ackedSeqNos.size() >= stopOnAcksCount) { - return; - } - } - } -} - } namespace NYdb::inline V3::NTopic::NTests { @@ -932,7 +856,7 @@ TEST_F(BasicUsage, TEST_NAME(TWriteSession_WriteEncoded_Broken)) { } } -TEST_F(BasicUsage, TEST_NAME(TKeyedWriteSessionBasicWrite_NoAutoPartitioning)) { +TEST_F(BasicUsage, TEST_NAME(TProducerBasicWrite_NoAutoPartitioning)) { // Basic write test for keyed write session. // Write 10 messages with different keys and check that they are written to different partitions. // Check that the order of messages is preserved. @@ -946,34 +870,30 @@ TEST_F(BasicUsage, TEST_NAME(TKeyedWriteSessionBasicWrite_NoAutoPartitioning)) { auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); - TKeyedWriteSessionSettings writeSettings; + TProducerSettings writeSettings; writeSettings .Path(GetTopicPath(TOPIC_NAME)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TKeyedWriteSessionSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); + writeSettings.MaxBlock(TDuration::Seconds(30)); - auto session = client.CreateKeyedWriteSession(writeSettings); - auto keyedSession = std::dynamic_pointer_cast(session); + auto producer = client.CreateProducer(writeSettings); + auto keyedSession = std::dynamic_pointer_cast(producer); - NPersQueue::NTests::TKeyedWriteSessionTestAdapter testAdapter(session.get()); for (size_t i = 0; i < 100; ++i) { auto key = CreateGuidAsString(); - auto token = testAdapter.GetContinuationToken(TDuration::Seconds(30)); - ASSERT_TRUE(token.has_value()) << "Timed out waiting for ReadyToAcceptEvent"; TWriteMessage msg("msg"); msg.SeqNo(i + 1); - session->Write(std::move(*token), key, std::move(msg)); + msg.Key(key); + ASSERT_TRUE(producer->Write(std::move(msg)).IsSuccess()); } - testAdapter.WaitForAcks(100, TDuration::Seconds(30)); - ASSERT_TRUE(session->Close(TDuration::Seconds(10))); - ASSERT_EQ(testAdapter.GetAckedSeqNosCount(), 100ull); - ASSERT_TRUE(testAdapter.ValidateAcksOrder()); + ASSERT_TRUE(producer->Close(TDuration::Seconds(10)).IsSuccess()); auto after = client.DescribeTopic(GetTopicPath(TOPIC_NAME), describeTopicSettings).GetValueSync(); ASSERT_TRUE(after.IsSuccess()) << after.GetIssues().ToOneLineString(); From ce374f2aecc784eceb1bace0a5b98a03f9213707 Mon Sep 17 00:00:00 2001 From: Ivan Nikolaev Date: Tue, 10 Mar 2026 07:52:24 +0000 Subject: [PATCH 62/93] SchemeShard: add metrics and status handling for forced compaction (#35430) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_table.proto | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index f0e7add639d..e6dfcdc7172 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -395280d887bc0d1c34863836ee9136d8ffdfa81d +12b938fb40f92d00ec41c9f089ae00d9654675cb diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 9477c6b5856..86f806de757 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -975,6 +975,8 @@ message CompactMetadata { optional uint32 max_shards_in_flight = 3; optional CompactState.State state = 4; optional float progress = 5; + optional uint32 shards_total = 6; + optional uint32 shards_done = 7; } // Alter table with given path From 12556cce6594a0aa9fad25da58843c92bf1262a8 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Tue, 10 Mar 2026 07:52:31 +0000 Subject: [PATCH 63/93] LOGBROKER-10314 Fix race in read session (#35470) --- .github/last_commit.txt | 2 +- src/client/topic/impl/read_session_impl.ipp | 19 ++++++++++++++----- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e6dfcdc7172..d93e67a88c2 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -12b938fb40f92d00ec41c9f089ae00d9654675cb +007c64774876f18b35178692b866f6c52e2791f5 diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index 563a44edf6b..78c3b6bbee5 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -2253,14 +2253,23 @@ inline void TSingleClusterReadSessionImpl::ConfirmPartitionStreamEnd(TPar std::lock_guard guard(HierarchyDataLock); ReadingFinishedData.insert(partitionStream->GetPartitionSessionId()); } - for (auto& [_, s] : PartitionStreams) { - for (auto partitionId : childIds) { - if (s->GetPartitionId() == partitionId) { - EventsQueue->SignalReadyEvents(s); - break; + + std::vector>> partitionStreams; + { + std::lock_guard guard(Lock); + for (auto& [_, s] : PartitionStreams) { + for (auto partitionId : childIds) { + if (s->GetPartitionId() == partitionId) { + partitionStreams.push_back(s); + break; + } } } } + + for (auto& s : partitionStreams) { + EventsQueue->SignalReadyEvents(s); + } } template <> From 33d015b68c28702a16259f274fbe78c39e5fb926 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Tue, 10 Mar 2026 07:52:38 +0000 Subject: [PATCH 64/93] Fix: TTableClient destructor can hang indefinitely on Drain().Wait() (#35532) --- .github/last_commit.txt | 2 +- src/client/impl/session/session_pool.cpp | 33 +++++++++++++++++------- src/client/table/impl/table_client.cpp | 2 +- src/client/table/impl/table_client.h | 3 +++ 4 files changed, 29 insertions(+), 11 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index d93e67a88c2..4e27be1413a 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -007c64774876f18b35178692b866f6c52e2791f5 +abde8e11540c2ccb2c8b2d9d15b2bed3d3fbaf7e diff --git a/src/client/impl/session/session_pool.cpp b/src/client/impl/session/session_pool.cpp index 877566a34fe..5907e2ce12a 100644 --- a/src/client/impl/session/session_pool.cpp +++ b/src/client/impl/session/session_pool.cpp @@ -257,16 +257,31 @@ void TSessionPool::IncrementActiveCounterUnsafe() { } void TSessionPool::Drain(std::function&&)> cb, bool close) { - std::lock_guard guard(Mtx_); - Closed_ = close; - for (auto it = Sessions_.begin(); it != Sessions_.end();) { - it->second->UpdateServerCloseHandler(nullptr); - const bool cont = cb(std::move(it->second)); - it = Sessions_.erase(it); - if (!cont) - break; + std::vector> waitersToReplyError; + { + std::lock_guard guard(Mtx_); + Closed_ = close; + for (auto it = Sessions_.begin(); it != Sessions_.end();) { + it->second->UpdateServerCloseHandler(nullptr); + const bool cont = cb(std::move(it->second)); + it = Sessions_.erase(it); + if (!cont) + break; + } + if (close) { + // Collect all pending waiters to reply with error outside the lock. + // When the pool is permanently closed, all pending session requests + // must be rejected to avoid them waiting indefinitely. + while (auto waiter = WaitersQueue_.TryGet()) { + waitersToReplyError.push_back(std::move(waiter)); + } + } + UpdateStats(); + } + for (auto& waiter : waitersToReplyError) { + FakeSessionsCounter_.Inc(); + waiter->ReplyError(CLIENT_RESOURCE_EXHAUSTED_ACTIVE_SESSION_LIMIT); } - UpdateStats(); } TPeriodicCb TSessionPool::CreatePeriodicTask(std::weak_ptr weakClient, diff --git a/src/client/table/impl/table_client.cpp b/src/client/table/impl/table_client.cpp index 92876f3d4fe..5b5d6d2c838 100644 --- a/src/client/table/impl/table_client.cpp +++ b/src/client/table/impl/table_client.cpp @@ -38,7 +38,7 @@ TTableClient::TImpl::TImpl(std::shared_ptr&& connections, TTableClient::TImpl::~TImpl() { if (Connections_->GetDrainOnDtors()) { - Drain().Wait(); + Drain().Wait(DRAIN_TIMEOUT); } } diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 0e3e4734462..9c7566498ab 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -30,6 +30,9 @@ namespace NTable { //How ofter run host scan to perform session balancing constexpr TDeadline::Duration HOSTSCAN_PERIODIC_ACTION_INTERVAL = std::chrono::seconds(2); constexpr TDuration KEEP_ALIVE_CLIENT_TIMEOUT = TDuration::Seconds(5); +// Max wait time for Drain() to complete. Sessions are being closed with 2 seconds +// timeout each (in parallel), plus up to 10 seconds for discovery if needed. +constexpr TDuration DRAIN_TIMEOUT = TDuration::Seconds(30); TDuration GetMinTimeToTouch(const TSessionPoolSettings& settings); TDuration GetMaxTimeToTouch(const TSessionPoolSettings& settings); From e92250e513bb56053b58f1519fbc5ab0c6c2b0ae Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 10 Mar 2026 07:52:38 +0000 Subject: [PATCH 65/93] Update import generation: 36 --- .github/import_generation.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/import_generation.txt b/.github/import_generation.txt index 7facc89938b..81b5c5d06cc 100644 --- a/.github/import_generation.txt +++ b/.github/import_generation.txt @@ -1 +1 @@ -36 +37 From bce3a46c653795245767e46d09967a15c0dbe724 Mon Sep 17 00:00:00 2001 From: Nikolay Perfilov Date: Tue, 10 Mar 2026 12:45:16 +0300 Subject: [PATCH 66/93] Update CMake files --- examples/CMakeLists.txt | 1 + .../producer/basic_write/CMakeLists.txt | 36 +++++++++++++++++++ .../ydb-cpp-sdk/client/topic/CMakeLists.txt | 6 ++++ src/client/topic/impl/CMakeLists.txt | 1 + 4 files changed, 44 insertions(+) create mode 100644 examples/topic_writer/producer/basic_write/CMakeLists.txt diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 380b49e8e99..1e2c1a0f8a1 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -6,6 +6,7 @@ add_subdirectory(secondary_index_builtin) add_subdirectory(time) add_subdirectory(topic_reader) add_subdirectory(topic_writer/transaction) +add_subdirectory(topic_writer/producer/basic_write) add_subdirectory(ttl) add_subdirectory(vector_index) add_subdirectory(vector_index_builtin) diff --git a/examples/topic_writer/producer/basic_write/CMakeLists.txt b/examples/topic_writer/producer/basic_write/CMakeLists.txt new file mode 100644 index 00000000000..f95bbedad67 --- /dev/null +++ b/examples/topic_writer/producer/basic_write/CMakeLists.txt @@ -0,0 +1,36 @@ +add_executable(topic_writer_producer_buffer_overloaded) + +target_link_libraries(topic_writer_producer_buffer_overloaded + PUBLIC + yutil + YDB-CPP-SDK::Topic + YDB-CPP-SDK::Query +) + +target_sources(topic_writer_producer_buffer_overloaded + PRIVATE + main.cpp +) + +vcs_info(topic_writer_producer_buffer_overloaded) + +if (CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64" OR CMAKE_SYSTEM_PROCESSOR STREQUAL "AMD64") + target_link_libraries(topic_writer_producer_buffer_overloaded PUBLIC + cpuid_check + ) +endif() + +if (CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_link_options(topic_writer_producer_buffer_overloaded PRIVATE + -ldl + -lrt + -Wl,--no-as-needed + -lpthread + ) +elseif (CMAKE_SYSTEM_NAME STREQUAL "Darwin") + target_link_options(topic_writer_producer_buffer_overloaded PRIVATE + -Wl,-platform_version,macos,11.0,11.0 + -framework + CoreFoundation + ) +endif() diff --git a/include/ydb-cpp-sdk/client/topic/CMakeLists.txt b/include/ydb-cpp-sdk/client/topic/CMakeLists.txt index 9a6bede21ed..1bb401bc9dc 100644 --- a/include/ydb-cpp-sdk/client/topic/CMakeLists.txt +++ b/include/ydb-cpp-sdk/client/topic/CMakeLists.txt @@ -26,4 +26,10 @@ generate_enum_serilization(client-ydb_topic-include include/ydb-cpp-sdk/client/topic/write_events.h ) +generate_enum_serilization(client-ydb_topic-include + ${YDB_SDK_SOURCE_DIR}/include/ydb-cpp-sdk/client/topic/producer.h + INCLUDE_HEADERS + include/ydb-cpp-sdk/client/topic/producer.h +) + _ydb_sdk_install_targets(TARGETS client-ydb_topic-include) \ No newline at end of file diff --git a/src/client/topic/impl/CMakeLists.txt b/src/client/topic/impl/CMakeLists.txt index 96fc3b6ba8b..7bb76da30b7 100644 --- a/src/client/topic/impl/CMakeLists.txt +++ b/src/client/topic/impl/CMakeLists.txt @@ -36,6 +36,7 @@ target_sources(client-ydb_topic-impl read_session_event.cpp read_session.cpp topic_impl.cpp + producer.cpp topic.cpp transaction.cpp write_session_impl.cpp From 8ba9df81657119fb1033266a19bb468fd79e47e8 Mon Sep 17 00:00:00 2001 From: Nikita Vasilev Date: Mon, 23 Mar 2026 10:00:10 +0000 Subject: [PATCH 67/93] Remove EnableOltpSink flag from tests (Part 1) (#35499) --- .github/last_commit.txt | 2 +- .../topic/ut/ut_utils/txusage_fixture.cpp | 22 ------------------- .../topic/ut/ut_utils/txusage_fixture.h | 4 ---- 3 files changed, 1 insertion(+), 27 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 4e27be1413a..4a9d09bd3bc 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -abde8e11540c2ccb2c8b2d9d15b2bed3d3fbaf7e +56c1b180cdd3554b3d14744c7c8ecc93aa81f7fa diff --git a/src/client/topic/ut/ut_utils/txusage_fixture.cpp b/src/client/topic/ut/ut_utils/txusage_fixture.cpp index 3b33e5f5064..f92a261d439 100644 --- a/src/client/topic/ut/ut_utils/txusage_fixture.cpp +++ b/src/client/topic/ut/ut_utils/txusage_fixture.cpp @@ -38,8 +38,6 @@ void TFixture::SetUp(NUnitTest::TTestContext&) NKikimr::Tests::TServerSettings settings = TTopicSdkTestSetup::MakeServerSettings(); settings.SetEnableTopicServiceTx(true); settings.SetEnableTopicSplitMerge(true); - settings.SetEnableOltpSink(GetEnableOltpSink()); - settings.SetEnableOlapSink(GetEnableOlapSink()); settings.SetEnableHtapTx(GetEnableHtapTx()); settings.SetAllowOlapDataQuery(GetAllowOlapDataQuery()); @@ -1088,16 +1086,6 @@ auto TFixture::GetAvgWriteBytes(const std::string& topicName, return result; } -bool TFixture::GetEnableOltpSink() const -{ - return false; -} - -bool TFixture::GetEnableOlapSink() const -{ - return false; -} - bool TFixture::GetEnableHtapTx() const { return false; @@ -2103,16 +2091,6 @@ void TFixtureSinks::CreateColumnTable(const std::string& tablePath) UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); } -bool TFixtureSinks::GetEnableOltpSink() const -{ - return true; -} - -bool TFixtureSinks::GetEnableOlapSink() const -{ - return true; -} - bool TFixtureSinks::GetEnableHtapTx() const { return true; diff --git a/src/client/topic/ut/ut_utils/txusage_fixture.h b/src/client/topic/ut/ut_utils/txusage_fixture.h index 5ee506491b4..634492186b7 100644 --- a/src/client/topic/ut/ut_utils/txusage_fixture.h +++ b/src/client/topic/ut/ut_utils/txusage_fixture.h @@ -268,8 +268,6 @@ class TFixture : public NUnitTest::TBaseFixture { std::uint32_t partitionId, const std::string& boundary); - virtual bool GetEnableOltpSink() const; - virtual bool GetEnableOlapSink() const; virtual bool GetEnableHtapTx() const; virtual bool GetAllowOlapDataQuery() const; @@ -404,8 +402,6 @@ class TFixtureSinks : public TFixture { void CreateRowTable(const std::string& path); void CreateColumnTable(const std::string& tablePath); - bool GetEnableOltpSink() const override; - bool GetEnableOlapSink() const override; bool GetEnableHtapTx() const override; bool GetAllowOlapDataQuery() const override; From e592180d4290a9f4eafabe1da9cfad72afd67c08 Mon Sep 17 00:00:00 2001 From: Andrey Zaspa Date: Mon, 23 Mar 2026 10:00:17 +0000 Subject: [PATCH 68/93] Make TServer driver private (#35592) --- .github/last_commit.txt | 2 +- src/client/persqueue_public/ut/ut_utils/test_server.h | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 4a9d09bd3bc..11034a44fe8 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -56c1b180cdd3554b3d14744c7c8ecc93aa81f7fa +338fb286a42e44caab2d0cbe89b6d3758bbc3234 diff --git a/src/client/persqueue_public/ut/ut_utils/test_server.h b/src/client/persqueue_public/ut/ut_utils/test_server.h index a2c731384c6..492a9477829 100644 --- a/src/client/persqueue_public/ut/ut_utils/test_server.h +++ b/src/client/persqueue_public/ut/ut_utils/test_server.h @@ -23,6 +23,7 @@ class TTestServer { : PortManager(portManager.GetOrElse(MakeSimpleShared())) , Port(PortManager->GetPort(2134)) , GrpcPort(PortManager->GetPort(2135)) + , Endpoint("localhost:" + ToString(GrpcPort)) , ServerSettings(settings) , GrpcServerOptions(NYdbGrpc::TServerOptions().SetHost("[::1]").SetPort(GrpcPort)) { @@ -62,6 +63,11 @@ class TTestServer { CleverServer = MakeHolder(ServerSettings); CleverServer->EnableGRpc(GrpcServerOptions); + auto driverConfig = NYdb::TDriverConfig() + .SetEndpoint(Endpoint) + .SetDatabase("/" + ServerSettings.DomainName); + Driver = MakeHolder(driverConfig); + Log << TLOG_INFO << "TTestServer started on Port " << Port << " GrpcPort " << GrpcPort; AnnoyingClient = MakeHolder(ServerSettings, GrpcPort, databaseName); @@ -121,7 +127,7 @@ class TTestServer { } const NYdb::TDriver& GetDriver() const { - return CleverServer->GetDriver(); + return *Driver; } void KillTopicPqrbTablet(const TString& topicPath) { @@ -164,6 +170,7 @@ class TTestServer { TSimpleSharedPtr PortManager; ui16 Port; ui16 GrpcPort; + TString Endpoint; THolder CleverServer; NKikimr::Tests::TServerSettings ServerSettings; @@ -176,6 +183,8 @@ class TTestServer { static const TVector LOGGED_SERVICES; +private: + THolder Driver; }; } // namespace NPersQueue From 19b5eaa7a301e4ab816493bd52d728037ee8a015 Mon Sep 17 00:00:00 2001 From: FloatingCrowbar Date: Mon, 23 Mar 2026 10:00:24 +0000 Subject: [PATCH 69/93] Transfer counters and metrics (#30345) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_replication.proto | 80 ++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 11034a44fe8..2e05f7d5a03 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -338fb286a42e44caab2d0cbe89b6d3758bbc3234 +fe48fe1c86b935af406a1157df5489c18dc0f070 diff --git a/src/api/protos/draft/ydb_replication.proto b/src/api/protos/draft/ydb_replication.proto index 62106c4d027..51d03d9b5ac 100644 --- a/src/api/protos/draft/ydb_replication.proto +++ b/src/api/protos/draft/ydb_replication.proto @@ -7,6 +7,7 @@ import "src/api/protos/ydb_operation.proto"; import "src/api/protos/ydb_scheme.proto"; import "google/protobuf/duration.proto"; +import "google/protobuf/timestamp.proto"; package Ydb.Replication; option java_package = "com.yandex.ydb.replication"; @@ -101,6 +102,8 @@ message DescribeTransferRequest { Ydb.Operations.OperationParams operation_params = 1; // Replication path. string path = 2 [(required) = true]; + // Include detailed statistics for each worker. + optional bool include_stats = 3; } message DescribeTransferResponse { @@ -108,6 +111,14 @@ message DescribeTransferResponse { Ydb.Operations.Operation operation = 1; } +// Message representing sliding window statistics by several windows. +message MultipleWindowsStat { + // Average per minute. + google.protobuf.Duration avg_per_minute = 1; + // Average per hour. + google.protobuf.Duration avg_per_hour = 2; +} + message DescribeTransferResult { message RunningState { } @@ -122,6 +133,74 @@ message DescribeTransferResult { message PausedState { } + message Stats { + enum WorkState { + // No reported state. + STATE_UNSPECIFIED = 0; + // Reading data from topic. + STATE_READ = 1; + // Decompressing read data. + STATE_DECOMPRESS = 2; + // Rrocessing data with lambda. + STATE_PROCESS = 3; + // Writing data to table. + STATE_WRITE = 4; + } + + message WorkerStats { + // unique worker identifier + string worker_id = 1; + // last reported work state + WorkState state = 2; + // timestamp of last state change + google.protobuf.Timestamp last_state_change = 3; + // partition of the topic that worker is reading + int64 partition_id = 4; + // current read offset in partition + int64 read_offset = 5; + // how long this worker is being run for + google.protobuf.Duration uptime = 6; + // total cumulative number of worker restarts + int64 restart_count = 7; + // worker restarts per time, sliding window + MultipleWindowsStat restarts = 8; + // topic read speed in sliding window, bytes + MultipleWindowsStat read_bytes = 9; + // topic read speed in sliding window, messages + MultipleWindowsStat read_messages = 10; + // table write speed in sliding window, bytes + MultipleWindowsStat write_bytes = 11; + // table write speed in sliding window, rows + MultipleWindowsStat write_rows = 12; + + // cpu usage time on decompression, microseconds + MultipleWindowsStat decompression_cpu_time = 13; + // cpu usage time on processing, microseconds + MultipleWindowsStat processing_cpu_time = 14; + } + + // detailed stats of every worker, only included if DescribeTransferRequest.include_stats flag = true in request + repeated WorkerStats workers_stats = 1; + // minimal uptime of all current workers + google.protobuf.Duration min_worker_uptime = 2; + + // topic read speed in sliding window, overall for transfer, bytes + MultipleWindowsStat read_bytes = 3; + // topic read speed in sliding window, overall for transfer, messages + MultipleWindowsStat read_messages = 4; + // table write speed in sliding window, overall for transfer, bytes + MultipleWindowsStat write_bytes = 5; + // table write speed in sliding window, overall for transfer, rows + MultipleWindowsStat write_rows = 6; + + // cpu usage time on decompression, overall for transfer, microseconds + MultipleWindowsStat decompression_cpu_time = 7; + // cpu usage time on processing, overall for transfer, microseconds + MultipleWindowsStat processing_cpu_time = 8; + // moment when stats collection was last started (or reset) + google.protobuf.Timestamp stats_collection_start = 9; + } + // Description of scheme object. Ydb.Scheme.Entry self = 1; @@ -145,4 +224,5 @@ message DescribeTransferResult { } optional BatchSettings batch_settings = 11; + optional Stats stats = 12; } From 25055e0d254ee15e4e10474e7e629ea536afca8f Mon Sep 17 00:00:00 2001 From: Maksim Date: Mon, 23 Mar 2026 10:00:31 +0000 Subject: [PATCH 70/93] [NBS-6805] NBS 2 tests (#35566) --- .github/last_commit.txt | 2 +- src/api/protos/draft/ydb_nbs.proto | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 2e05f7d5a03..af9d1095156 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -fe48fe1c86b935af406a1157df5489c18dc0f070 +aecc3e23bc4a676e23e0d8fbad6558664deb8e08 diff --git a/src/api/protos/draft/ydb_nbs.proto b/src/api/protos/draft/ydb_nbs.proto index 411c6567689..3352ed432da 100644 --- a/src/api/protos/draft/ydb_nbs.proto +++ b/src/api/protos/draft/ydb_nbs.proto @@ -71,8 +71,8 @@ message DeletePartitionResult { message GetLoadActorAdapterActorIdRequest { Ydb.Operations.OperationParams operation_params = 1; - // Partition tablet id (actor id string). - string TabletId = 2; + // Disk id. + string DiskId = 2; } message GetLoadActorAdapterActorIdResponse { From 199cd800f3e46c9c075a03eada40a8639e55fc78 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:00:38 +0000 Subject: [PATCH 71/93] LOGBROKER-10206 Add seqNo initialization (#35653) --- .github/last_commit.txt | 2 +- src/client/topic/impl/producer.cpp | 294 +++++++++++++++++++++---- src/client/topic/impl/producer.h | 46 +++- src/client/topic/ut/basic_usage_ut.cpp | 170 ++++++++++++-- 4 files changed, 442 insertions(+), 70 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index af9d1095156..6399adab2f6 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -aecc3e23bc4a676e23e0d8fbad6558664deb8e08 +748f926df551161c7b935fba9cdbf64ed5590fbb diff --git a/src/client/topic/impl/producer.cpp b/src/client/topic/impl/producer.cpp index 8057902e0fa..f43eb71a3b7 100644 --- a/src/client/topic/impl/producer.cpp +++ b/src/client/topic/impl/producer.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include @@ -71,9 +72,10 @@ TWriteMessage TProducer::TMessageInfo::BuildMessage() const { //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TProducer::TWriteSessionWrapper -TProducer::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition) +TProducer::TWriteSessionWrapper::TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition, bool directToPartition) : Session(std::move(session)) , Partition(partition) + , DirectToPartition(directToPartition) {} //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -171,6 +173,10 @@ void TProducer::TSplittedPartitionWorker::DoWork() { Producer->Partitions[child].Locked(false); } Producer->Partitions[PartitionId].Locked_ = false; + + for (const auto& [partitionId, maxSeqNo] : CachedMaxSeqNos) { + Producer->Partitions[partitionId].CachedMaxSeqNo = maxSeqNo; + } MoveTo(EState::Done); break; } @@ -181,7 +187,8 @@ void TProducer::TSplittedPartitionWorker::MoveTo(EState state) { LOG_LAZY(Producer->DbDriverState->Log, TLOG_INFO, Producer->LogPrefix() << "Moving splitted partition worker for partition " << PartitionId << " to state " << GetStateName()); } -void TProducer::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint64_t maxSeqNo) { +void TProducer::TSplittedPartitionWorker::UpdateMaxSeqNo(std::uint32_t partitionId, std::uint64_t maxSeqNo) { + CachedMaxSeqNos[partitionId] = maxSeqNo; MaxSeqNo = std::max(MaxSeqNo, maxSeqNo); } @@ -269,7 +276,12 @@ void TProducer::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_l NotReadyFutures = ancestors.size(); for (const auto& ancestor : ancestors) { - auto wrappedSession = Producer->SessionsWorker->GetWriteSession(ancestor, false); + if (Producer->Partitions[ancestor].CachedMaxSeqNo.has_value()) { + --NotReadyFutures; + UpdateMaxSeqNo(ancestor, Producer->Partitions[ancestor].CachedMaxSeqNo.value()); + continue; + } + auto wrappedSession = Producer->SessionsWorker->GetOrCreateWriteSession(ancestor, false); Y_ABORT_UNLESS(wrappedSession, "Write session not found"); WriteSessions.push_back(wrappedSession); @@ -297,7 +309,7 @@ void TProducer::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_l return; } - UpdateMaxSeqNo(result.GetValue()); + UpdateMaxSeqNo(ancestor, result.GetValue()); if (--NotReadyFutures == 0) { MoveTo(EState::GotMaxSeqNo); gotMaxSeqNo = true; @@ -307,6 +319,7 @@ void TProducer::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_l if (gotMaxSeqNo) { producerPtr->RunMainWorker(static_cast(PartitionId)); } + producerPtr->SessionsWorker->DestroyWriteSession(ancestor); }); lock.lock(); GetMaxSeqNoFutures.push_back(future); @@ -315,6 +328,12 @@ void TProducer::TSplittedPartitionWorker::LaunchGetMaxSeqNoFutures(std::unique_l if (ancestors.empty()) { LOG_LAZY(Producer->DbDriverState->Log, TLOG_INFO, Producer->LogPrefix() << "No ancestors found for partition " << PartitionId); MoveTo(EState::Init); + return; + } + + if (NotReadyFutures == 0) { + MoveTo(EState::GotMaxSeqNo); + Producer->RunMainWorker(static_cast(PartitionId)); } } @@ -373,7 +392,6 @@ bool TProducer::TEventsWorker::RunEventLoop(WrappedWriteSessionPtr wrappedSessio } if (auto acksEvent = std::get_if(&*event)) { - // Producer->SessionsWorker->OnReadFromSession(wrappedSession, acksEvent->Acks.size()); HandleAcksEvent(partition, std::move(*acksEvent)); continue; } @@ -386,16 +404,22 @@ std::optional> TProducer::TEventsWorker::DoWork() { std::unique_lock lock(Lock); while (!ReadyFutures.empty()) { - auto idx = *ReadyFutures.begin(); - ReadyFutures.erase(idx); + auto partition = *ReadyFutures.begin(); + ReadyFutures.erase(partition); + auto session = Producer->SessionsWorker->GetWriteSession(partition); + if (!session) { + continue; + } + lock.unlock(); // RunEventLoop without Lock: sub-session's WaitEvent() completion may run the Subscribe // callback (ReadyFutures.insert) synchronously; that callback takes Lock -> same-thread deadlock. - auto isSessionClosed = RunEventLoop(Producer->SessionsWorker->GetWriteSession(idx), idx); + auto isSessionClosed = RunEventLoop(session, partition); if (!isSessionClosed) { - SubscribeToPartition(idx); + SubscribeToPartition(partition); } else { - UnsubscribeFromPartition(idx); + session->Closed = true; + Producer->SessionsWorker->DestroyWriteSession(partition); } lock.lock(); } @@ -413,7 +437,7 @@ void TProducer::TEventsWorker::SubscribeToPartition(std::uint32_t partition) { return; } - auto wrappedSession = Producer->SessionsWorker->GetWriteSession(partition); + auto wrappedSession = Producer->SessionsWorker->GetOrCreateWriteSession(partition); auto newFuture = wrappedSession->Session->WaitEvent(); std::weak_ptr producer = Producer->shared_from_this(); std::weak_ptr self = shared_from_this(); @@ -490,7 +514,7 @@ bool TProducer::TEventsWorker::TransferEventsToOutputQueue() { TWriteSessionEvent::TAcksEvent ackEvent; if (expectedSeqNo.has_value()) { - Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << Producer->Partitions[partition].PartitionId_); + Y_ENSURE(acksQueue.front().SeqNo == expectedSeqNo.value(), TStringBuilder() << "Expected seqNo=" << expectedSeqNo.value() << " but got " << acksQueue.front().SeqNo << " for partition " << partition); } auto ack = std::move(acksQueue.front()); @@ -682,12 +706,36 @@ TProducer::TSessionsWorker::TSessionsWorker(TProducer* producer) : Producer(producer) {} -TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::GetWriteSession(std::uint32_t partition, bool directToPartition) { +TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::GetOrCreateWriteSession(std::uint32_t partition, bool directToPartition) { auto sessionIter = SessionsIndex.find(partition); - if (sessionIter == SessionsIndex.end() || !directToPartition) { + if (sessionIter == SessionsIndex.end()) { return CreateWriteSession(partition, directToPartition); } + SessionsToRemove.erase(partition); + if (!directToPartition && sessionIter->second->DirectToPartition) { + Y_ABORT_UNLESS(sessionIter->second->Closed, "Session is not closed for partition: %u", partition); + ClosedSessionsToRemove.push_back(sessionIter->second); + SessionsIndex.erase(sessionIter); + return CreateWriteSession(partition, directToPartition); + } + + if (!sessionIter->second->DirectToPartition) { + sessionIter->second->NonDirectToPartitionOwnership++; + } + + return sessionIter->second; +} + +TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::GetWriteSession(std::uint32_t partition, bool directToPartition) { + auto sessionIter = SessionsIndex.find(partition); + if (sessionIter == SessionsIndex.end()) { + return nullptr; + } + + SessionsToRemove.erase(partition); + Y_ABORT_UNLESS(directToPartition == sessionIter->second->DirectToPartition, "DirectToPartition mismatch: %s != %s", directToPartition ? "true" : "false", sessionIter->second->DirectToPartition ? "true" : "false"); + return sessionIter->second; } @@ -716,27 +764,40 @@ TProducer::WrappedWriteSessionPtr TProducer::TSessionsWorker::CreateWriteSession } auto writeSession = std::make_shared( Producer->Client->CreateWriteSession(alteredSettings), - partition); + partition, + directToPartition + ); + SessionsIndex.emplace(partition, writeSession); if (directToPartition) { - SessionsIndex.emplace(partition, writeSession); Producer->EventsWorker->SubscribeToPartition(partition); } return writeSession; } -void TProducer::TSessionsWorker::DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout) { +void TProducer::TSessionsWorker::DestroyWriteSession(std::uint32_t partition) { + auto it = SessionsIndex.find(partition); if (it == SessionsIndex.end() || !it->second) { return; } - it->second->Session->Close(closeTimeout); - const auto partition = it->second->Partition; + if (!it->second->DirectToPartition && --it->second->NonDirectToPartitionOwnership > 0) { + return; + } + + if (it->second->DirectToPartition) { + Producer->EventsWorker->UnsubscribeFromPartition(partition); + } + + // Remove idle bookkeeping before erasing the session from SessionsIndex so stale + // idle markers cannot later evict a new session created for the same partition. + RemoveIdleSession(partition); + if (static_cast(partition) == Producer->MainWorkerOwner) { - SessionsToRemove.push_back(it->second); + SessionsToRemove.emplace(partition); + } else { + SessionsIndex.erase(it); } - it = SessionsIndex.erase(it); - Producer->EventsWorker->UnsubscribeFromPartition(partition); } size_t TProducer::TSessionsWorker::GetSessionsCount() const { @@ -770,23 +831,40 @@ void TProducer::TSessionsWorker::RemoveIdleSession(std::uint32_t partition) { return; } + const auto idleSession = *itIdle->second; + IdlerSessions.erase(itIdle->second); + IdlerSessionsIndex.erase(itIdle); + auto wrappedSession = SessionsIndex.find(partition); if (wrappedSession == SessionsIndex.end()) { return; } - IdlerSessions.erase(itIdle->second); - IdlerSessionsIndex.erase(itIdle); - wrappedSession->second->IdleSession.reset(); + if (wrappedSession->second.get() == idleSession->Session) { + wrappedSession->second->IdleSession.reset(); + } } void TProducer::TSessionsWorker::DoWork() { - while (!SessionsToRemove.empty()) { - if (static_cast(SessionsToRemove.front()->Partition) == Producer->MainWorkerOwner) { - break; + for (auto it = SessionsToRemove.begin(); it != SessionsToRemove.end();) { + auto partition = *it; + if (static_cast(partition) == Producer->MainWorkerOwner) { + ++it; + continue; + } + + SessionsIndex.erase(partition); + it = SessionsToRemove.erase(it); + } + + for (auto it = ClosedSessionsToRemove.begin(); it != ClosedSessionsToRemove.end();) { + auto session = *it; + if (static_cast(session->Partition) == Producer->MainWorkerOwner) { + ++it; + continue; } - SessionsToRemove.pop_front(); + it = ClosedSessionsToRemove.erase(it); } while (!IdlerSessions.empty()) { @@ -795,11 +873,10 @@ void TProducer::TSessionsWorker::DoWork() { break; } - LOG_LAZY(Producer->DbDriverState->Log, TLOG_DEBUG, TStringBuilder() << Producer->LogPrefix() << "Removing idle session for partition " << (*it)->Session->Partition); - - const auto partition = (*it)->Session->Partition; + auto expiredIdleSession = *it; + const auto partition = expiredIdleSession->Session->Partition; if (Producer->Partitions[partition].Locked_) { - continue; + break; } // Remove idle tracking first to keep containers consistent even if the session @@ -809,8 +886,11 @@ void TProducer::TSessionsWorker::DoWork() { auto sessionIter = SessionsIndex.find(partition); if (sessionIter != SessionsIndex.end()) { + if (sessionIter->second.get() != expiredIdleSession->Session) { + continue; + } sessionIter->second->IdleSession.reset(); - DestroyWriteSession(sessionIter, TDuration::Zero()); + DestroyWriteSession(partition); } } } @@ -834,9 +914,114 @@ void TProducer::TMessagesWorker::RechoosePartitionIfNeeded(MessageIter message) message->Partition = newPartition; } +void TProducer::TMessagesWorker::HandleReadyInitSeqNoFutures() { + std::unique_lock lock(InitLock); + for (const auto& partition : GotInitSeqNoPartitions) { + auto it = InitGetMaxSeqNoFutures.find(partition); + Y_ABORT_UNLESS(it != InitGetMaxSeqNoFutures.end(), "Init get max seq no future not found"); + Y_ABORT_UNLESS(it->second.IsReady(), "Init get max seq no future is not ready"); + + auto gotMaxSeqNo = it->second.GetValue(); + CurrentSeqNo = std::max(CurrentSeqNo, gotMaxSeqNo); + Producer->Partitions[partition].CachedMaxSeqNo = gotMaxSeqNo; + InitGetMaxSeqNoFutures.erase(it); + } + + GotInitSeqNoPartitions.clear(); +} + +void TProducer::TMessagesWorker::FinishInit() { + for (const auto& partition : Producer->Partitions) { + if (!partition.second.IsSplitted() && !InFlightMessagesIndex.contains(partition.first)) { + Producer->SessionsWorker->AddIdleSession(partition.first); + } + + if (partition.second.IsSplitted()) { + Producer->SessionsWorker->DestroyWriteSession(partition.first); + } + } + + InitWriteSessions.clear(); + InitGetMaxSeqNoFutures.clear(); +} + +bool TProducer::TMessagesWorker::LazyInit() { + if (State == EState::Ready) { + return true; + } + + if (Producer->SeqNoStrategy == ESeqNoStrategy::WithSeqNo) { + MoveTo(EState::Ready); + return true; + } + + if (State == EState::PendingSeqNo) { + HandleReadyInitSeqNoFutures(); + if (InitGetMaxSeqNoFutures.empty()) { + FinishInit(); + MoveTo(EState::Ready); + return true; + } + + return false; + } + + std::weak_ptr producer = Producer->shared_from_this(); + std::weak_ptr self = shared_from_this(); + for (const auto& partition : Producer->Partitions) { + auto partitionId = partition.first; + WrappedWriteSessionPtr wrappedSession = nullptr; + if (partition.second.IsSplitted()) { + wrappedSession = Producer->SessionsWorker->GetOrCreateWriteSession(partition.first, false); + } else { + wrappedSession = Producer->SessionsWorker->GetOrCreateWriteSession(partition.first); + } + + InitWriteSessions.push_back(wrappedSession); + auto initGetMaxSeqNoFuture = wrappedSession->Session->GetInitSeqNo(); + + initGetMaxSeqNoFuture.Subscribe([self, producer, partitionId](NThreading::TFuture future) { + auto selfPtr = self.lock(); + if (!selfPtr) { + return; + } + + auto producerPtr = producer.lock(); + if (!producerPtr) { + return; + } + + if (!future.IsReady()) { + return; + } + + { + std::lock_guard lock(selfPtr->InitLock); + selfPtr->GotInitSeqNoPartitions.push_back(partitionId); + } + producerPtr->RunMainWorker(partitionId); + }); + InitGetMaxSeqNoFutures.emplace(partition.first, initGetMaxSeqNoFuture); + } + + MoveTo(EState::PendingSeqNo); + return false; +} + +void TProducer::TMessagesWorker::MoveTo(EState state) { + State = state; +} + void TProducer::TMessagesWorker::DoWork() { - auto sessionsWorker = Producer->SessionsWorker; + if (MessagesToResendIndex.empty() && PendingMessagesIndex.empty()) { + return; + } + if (!LazyInit()) { + return; + } + + auto sessionsWorker = Producer->SessionsWorker; auto iterateMessagesIndex = [&](std::unordered_map>& messagesIndex, auto stopCondition) { std::vector partitionsProcessed; for (auto& [partition, messages] : messagesIndex) { @@ -846,14 +1031,17 @@ void TProducer::TMessagesWorker::DoWork() { break; } - auto wrappedSession = sessionsWorker->GetWriteSession(head->Partition); + if (!head->SeqNo.has_value()) { + head->SeqNo.emplace(++CurrentSeqNo); + } + + auto wrappedSession = sessionsWorker->GetOrCreateWriteSession(head->Partition); if (!SendMessage(wrappedSession, *head)) { break; } Producer->Metrics.AddWriteLag((TInstant::Now() - head->CreateTimestamp.value_or(TInstant::Now())).MilliSeconds()); head->Sent = true; - // sessionsWorker->OnWriteToSession(wrappedSession); messages.pop_front(); } @@ -1245,15 +1433,19 @@ TProducer::TProducer( Settings(settings) { if (settings.ProducerIdPrefix_.empty()) { - ythrow TContractViolation("ProducerIdPrefix is required for KeyedWriteSession"); + ythrow TContractViolation("ProducerIdPrefix is required for Producer"); } if (!settings.ProducerId_.empty()) { - ythrow TContractViolation("ProducerId should be empty for KeyedWriteSession, use ProducerIdPrefix instead"); + ythrow TContractViolation("ProducerId should be empty for Producer, use ProducerIdPrefix instead"); } if (!settings.MessageGroupId_.empty()) { - ythrow TContractViolation("MessageGroupId should be empty for KeyedWriteSession"); + ythrow TContractViolation("MessageGroupId should be empty for Producer"); + } + + if (IsFederation(DbDriverState->DiscoveryEndpoint)) { + ythrow TContractViolation("Producer is not supported for federation"); } TDescribeTopicSettings describeTopicSettings; @@ -1306,16 +1498,17 @@ TProducer::TProducer( case TProducerSettings::EPartitionChooserStrategy::Bound: PartitioningKeyHasher = settings.PartitioningKeyHasher_; PartitionChooser = std::make_unique(this); - for (size_t i = 0; i < Partitions.size(); ++i) { - if (i > 0 && Partitions[i].FromBound_.empty() && !Partitions[i].ToBound_.has_value()) { + for (size_t i = 0; i < partitions.size(); ++i) { + const auto& partition = partitions[i]; + if (i > 0 && !partition.GetFromBound().has_value() && !partition.GetToBound().has_value()) { ythrow TContractViolation("Unbounded partition is not supported for Bound partition chooser strategy"); } - if (!Partitions[i].Children_.empty()) { + if (!partition.GetChildPartitionIds().empty()) { continue; } - PartitionsIndex[Partitions[i].FromBound_] = Partitions[i].PartitionId_; + PartitionsIndex[partition.GetFromBound().value_or("")] = partition.GetPartitionId(); } break; case TProducerSettings::EPartitionChooserStrategy::Hash: @@ -1353,8 +1546,7 @@ TProducer::TProducer( }); RunMainWorker(-1); - - LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Keyed write session created"); + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Producer created"); } std::vector TProducer::GetPartitions() const { @@ -1375,11 +1567,15 @@ std::map TProducer::GetPartitionsIndex() const { } size_t TProducer::GetSessionsCount() { + RunMainWorker(-1); + std::lock_guard lock(GlobalLock); return SessionsWorker->GetSessionsCount(); } size_t TProducer::GetIdleSessionsCount() { + RunMainWorker(-1); + std::lock_guard lock(GlobalLock); return SessionsWorker->GetIdleSessionsCount(); } @@ -1565,6 +1761,11 @@ void TProducer::RunUserEventLoop() { } } +bool TProducer::IsFederation(const std::string& endpoint) { + std::string_view host = GetHost(endpoint); + return host == "logbroker.yandex.net" || host == "logbroker-prestable.yandex.net"; +} + void TProducer::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSession, std::optional sessionClosedEvent) { std::optional receivedSessionClosedEvent; while (true) { @@ -1841,7 +2042,6 @@ TInstant TProducer::GetCloseDeadline() { } void TProducer::HandleAutoPartitioning(std::uint32_t partition) { - LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "HandleAutoPartitioning: " << partition); auto splittedPartitionWorker = std::make_shared(this, partition); SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); } diff --git a/src/client/topic/impl/producer.h b/src/client/topic/impl/producer.h index ff0c997132a..5760567aa9e 100644 --- a/src/client/topic/impl/producer.h +++ b/src/client/topic/impl/producer.h @@ -37,6 +37,8 @@ class TProducer : public IProducer, FLUENT_SETTING(std::vector, Children); FLUENT_SETTING_DEFAULT(bool, Locked, false); FLUENT_SETTING_DEFAULT(NThreading::TFuture, Future, NThreading::MakeFuture()); + + std::optional CachedMaxSeqNo; }; struct TMessageInfo { @@ -63,8 +65,13 @@ class TProducer : public IProducer, WriteSessionPtr Session; const std::uint32_t Partition; std::shared_ptr IdleSession = nullptr; + bool DirectToPartition = true; + bool Closed = false; + + // This field is used only when DirectToPartition is false. + size_t NonDirectToPartitionOwnership = 1; - TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition); + TWriteSessionWrapper(WriteSessionPtr session, std::uint32_t partition, bool directToPartition = true); }; using WrappedWriteSessionPtr = std::shared_ptr; @@ -118,18 +125,19 @@ class TProducer : public IProducer, struct TSessionsWorker { TSessionsWorker(TProducer* producer); + WrappedWriteSessionPtr GetOrCreateWriteSession(std::uint32_t partition, bool directToPartition = true); WrappedWriteSessionPtr GetWriteSession(std::uint32_t partition, bool directToPartition = true); void AddIdleSession(std::uint32_t partition); - void RemoveIdleSession(std::uint32_t partition); void DoWork(); size_t GetSessionsCount() const; size_t GetIdleSessionsCount() const; + void DestroyWriteSession(std::uint32_t partition); + void RemoveIdleSession(std::uint32_t partition); private: WrappedWriteSessionPtr CreateWriteSession(std::uint32_t partition, bool directToPartition = true); using TSessionsIndexIterator = std::unordered_map::iterator; - void DestroyWriteSession(TSessionsIndexIterator& it, TDuration closeTimeout); std::string GetProducerId(std::uint32_t partitionId); @@ -138,12 +146,11 @@ class TProducer : public IProducer, using IdlerSessionsIterator = std::set::iterator; std::unordered_map IdlerSessionsIndex; std::unordered_map SessionsIndex; - std::deque SessionsToRemove; - - static constexpr TDuration SESSION_REMOVE_DELAY = TDuration::Seconds(5); + std::unordered_set SessionsToRemove; + std::list ClosedSessionsToRemove; }; - struct TMessagesWorker { + struct TMessagesWorker : public std::enable_shared_from_this { TMessagesWorker(TProducer* producer); void DoWork(); @@ -160,6 +167,12 @@ class TProducer : public IProducer, void SetClosedStatusToFlushPromises(std::optional closedDescription); private: + enum class EState : std::uint8_t { + Init = 0, + PendingSeqNo = 1, + Ready = 2, + }; + using MessageIter = std::list::iterator; void PushInFlightMessage(std::uint32_t partition, TMessageInfo&& message); @@ -167,6 +180,10 @@ class TProducer : public IProducer, bool SendMessage(WrappedWriteSessionPtr wrappedSession, const TMessageInfo& message); std::optional GetContinuationToken(std::uint32_t partition); void RechoosePartitionIfNeeded(MessageIter message); + bool LazyInit(); + void MoveTo(EState state); + void HandleReadyInitSeqNoFutures(); + void FinishInit(); TProducer* Producer; @@ -177,6 +194,14 @@ class TProducer : public IProducer, std::unordered_map> ContinuationTokens; std::uint64_t MemoryUsage = 0; + std::uint64_t CurrentSeqNo = 0; + EState State = EState::Init; + + std::vector InitWriteSessions; + std::unordered_map> InitGetMaxSeqNoFutures; + + std::mutex InitLock; + std::vector GotInitSeqNoPartitions; friend class TProducer; }; @@ -193,7 +218,7 @@ class TProducer : public IProducer, }; void MoveTo(EState state); - void UpdateMaxSeqNo(uint64_t maxSeqNo); + void UpdateMaxSeqNo(std::uint32_t partitionId, std::uint64_t maxSeqNo); void LaunchGetMaxSeqNoFutures(std::unique_lock& lock); void HandleDescribeResult(); @@ -212,6 +237,7 @@ class TProducer : public IProducer, std::uint64_t MaxSeqNo = 0; std::vector WriteSessions; std::vector> GetMaxSeqNoFutures; + std::unordered_map CachedMaxSeqNos; std::mutex Lock; std::uint64_t NotReadyFutures = 0; size_t Retries = 0; @@ -239,11 +265,11 @@ class TProducer : public IProducer, std::list::iterator AckQueueEnd(std::uint32_t partition); std::optional GetContinuationToken(); std::optional GetSessionClosedEvent(); + bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition); private: void HandleSessionClosedEvent(TSessionClosedEvent&& event, std::uint32_t partition); void HandleReadyToAcceptEvent(std::uint32_t partition, TWriteSessionEvent::TReadyToAcceptEvent&& event); - bool RunEventLoop(WrappedWriteSessionPtr wrappedSession, std::uint32_t partition); bool TransferEventsToOutputQueue(); void AddContinuationToken(); bool AddSessionClosedIfNeeded(); @@ -356,6 +382,8 @@ class TProducer : public IProducer, TWriteResult WriteInternal(TContinuationToken&&, TWriteMessage&& message); + bool IsFederation(const std::string& endpoint); + public: TProducer(const TProducerSettings& settings, std::shared_ptr client, diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index f4c9367780d..682bcf21f63 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -68,6 +68,68 @@ std::string Serialize(const TExample& value) { return value.Payload; } +// Reads exactly expectedCount messages from the topic and asserts that within each partition +// messages are ordered by seqNo (strictly increasing). Uses provided client, topic path and consumer. +void ReadMessagesAndAssertOrderedBySeqNo(TTopicClient& client, + const std::string& topicPath, + const std::string& consumerName, + size_t expectedCount, + TDuration timeout = TDuration::Seconds(30)) { + struct TMessageInfo { + ui64 PartitionId; + TString ProducerId; + ui64 SeqNo; + }; + std::vector messages; + messages.reserve(expectedCount); + NThreading::TPromise donePromise = NThreading::NewPromise(); + + TTopicReadSettings topicSettings(topicPath); + topicSettings.ReadFromTimestamp(TInstant::Zero()); + + auto readSettings = TReadSessionSettings() + .ConsumerName(consumerName) + .AutoPartitioningSupport(true) + .AppendTopics(topicSettings); + + readSettings.EventHandlers_.SimpleDataHandlers([&](TReadSessionEvent::TDataReceivedEvent& ev) { + for (auto& msg : ev.GetMessages()) { + messages.push_back(TMessageInfo{ + msg.GetPartitionSession()->GetPartitionId(), + TString(msg.GetProducerId()), + msg.GetSeqNo(), + }); + } + if (messages.size() >= expectedCount) { + donePromise.SetValue(); + } + }, true); + + auto readSession = client.CreateReadSession(readSettings); + UNIT_ASSERT_C(donePromise.GetFuture().Wait(timeout), + "Expected to read " << expectedCount << " messages within " << timeout << ", got " << messages.size()); + readSession->Close(TDuration::Seconds(5)); + + UNIT_ASSERT_VALUES_EQUAL_C(messages.size(), expectedCount, + "Read message count mismatch: got " << messages.size() << ", expected " << expectedCount); + + // SeqNo ordering is guaranteed within one producer stream. + // Multiple producers can write into the same partition with independent seqNo sequences. + std::map, std::vector> byPartitionAndProducer; + for (const auto& m : messages) { + byPartitionAndProducer[{m.PartitionId, m.ProducerId}].push_back(m.SeqNo); + } + for (const auto& [key, seqNos] : byPartitionAndProducer) { + const auto& [partitionId, producerId] = key; + for (size_t i = 1; i < seqNos.size(); ++i) { + UNIT_ASSERT_C(seqNos[i] > seqNos[i - 1], + "Partition " << partitionId << ", producerId " << producerId + << ": expected seqNo strictly increasing, got " + << seqNos[i - 1] << " then " << seqNos[i] << " at index " << i); + } + } +} + // Write a message with binary (non-UTF8) producer ID using direct tablet communication // This bypasses gRPC string validation by sending directly to the PQ tablet // The SourceId field in TCmdWrite is defined as 'bytes' in protobuf, so it supports binary data @@ -164,6 +226,9 @@ static std::string FindKeyForBucket(size_t bucket, size_t bucketsCount) { void CreateTopicWithAutoPartitioning(TTopicClient& client) { TCreateTopicSettings createSettings; createSettings + .BeginAddConsumer() + .ConsumerName(TEST_CONSUMER) + .EndAddConsumer() .BeginConfigurePartitioningSettings() .MinActivePartitions(2) .MaxActivePartitions(100) @@ -1047,6 +1112,26 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_EXCEPTION(setup.MakeClient().CreateProducer(writeSettings), TContractViolation); } + Y_UNIT_TEST(Producer_IsNotSupportedForFederation) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + + auto config = setup.MakeDriverConfig(); + config.SetEndpoint("logbroker.yandex.net:2135"); + + TDriver driver(config); + TTopicClient federatedLikeClient(driver); + + TProducerSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); + + UNIT_ASSERT_EXCEPTION(federatedLikeClient.CreateProducer(writeSettings), TContractViolation); + } + Y_UNIT_TEST(Producer_SessionClosedDueToUserError) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); @@ -1223,7 +1308,40 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } } - Y_UNIT_TEST(Producer_EventLoop_Acks) { + Y_UNIT_TEST(Producer_WriteManyMessages) { + TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; + setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); + + auto client = setup.MakeClient(); + + TProducerSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath(TEST_TOPIC)) + .Codec(ECodec::RAW); + writeSettings.ProducerIdPrefix(CreateGuidAsString()); + writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.MaxBlock(TDuration::Seconds(30)); + + auto producer = client.CreateProducer(writeSettings); + + const ui64 count = 3000; + for (ui64 i = 1; i <= count; ++i) { + auto key = CreateGuidAsString(); + std::string payload = "data"; + TWriteMessage msg(payload); + msg.SeqNo(i); + msg.Key(key); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + } + + UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), count); + } + + Y_UNIT_TEST(Producer_AutoSeqNo) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 4); @@ -1252,6 +1370,19 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + + auto producer2 = client.CreateProducer(writeSettings); + for (ui64 i = 1; i <= count; ++i) { + auto key = CreateGuidAsString(); + std::string payload = "data"; + TWriteMessage msg(payload); + msg.Key(key); + UNIT_ASSERT_C(producer2->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + } + UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); + UNIT_ASSERT_C(producer2->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); + + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), count * 2); } Y_UNIT_TEST(Producer_WriteToClosedProducer) { @@ -1589,6 +1720,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(producer1->Close(TDuration::Seconds(30)).IsSuccess()); UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); UNIT_ASSERT(producer3->Close(TDuration::Seconds(30)).IsSuccess()); + + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), 14); } Y_UNIT_TEST(AutoPartitioning_Producer_SmallMessages) { @@ -1665,6 +1798,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(producer1->Close(TDuration::Seconds(30)).IsSuccess()); UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); + + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), totalMessages); } Y_UNIT_TEST(Producer_BasicWrite) { @@ -1701,6 +1836,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_EQUAL(messagesWritten, 100); UNIT_ASSERT_C(producer->Close(TDuration::Seconds(1)).IsSuccess(), "Failed to close producer"); + + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), 100); } Y_UNIT_TEST(TypedProducer_BasicWrite) { @@ -1787,20 +1924,27 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto producer = client.CreateProducer(writeSettings); auto producerRaw = dynamic_cast(producer.get()); auto msgData = TString(10_KB, 'a'); + for (const auto& partition : partitions) { + for (ui64 i = 0; i < 10; ++i) { + TWriteMessage msg(msgData); + msg.Partition(partition.GetPartitionId()); + UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); + } + } + UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); - for (ui64 i = 0; i < 3; ++i) { - for (const auto& partition : partitions) { - for (ui64 i = 0; i < 10; ++i) { - TWriteMessage msg(msgData); - msg.Partition(partition.GetPartitionId()); - UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); - } - UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); - UNIT_ASSERT((producerRaw->GetIdleSessionsCount() == 1 && producerRaw->GetSessionsCount() == 1) || - (producerRaw->GetIdleSessionsCount() == 0 && producerRaw->GetSessionsCount() == 0)); - Sleep(TDuration::Seconds(1)); + size_t idleSessionsCount = 0; + size_t sessionsCount = 0; + for (int i = 0; i < 5; ++i) { + idleSessionsCount = producerRaw->GetIdleSessionsCount(); + sessionsCount = producerRaw->GetSessionsCount(); + if (idleSessionsCount == 0 && sessionsCount == 0) { + break; } + + Sleep(TDuration::Seconds(1)); } + UNIT_ASSERT_C(idleSessionsCount == 0 && sessionsCount == 0, "Idle session count: " << idleSessionsCount << ", sessions count: " << sessionsCount); { auto describeResult = client.DescribeTopic(TEST_TOPIC, TDescribeTopicSettings().IncludeStats(true)).GetValueSync(); @@ -1810,7 +1954,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(stats); messagesWritten += stats->GetEndOffset() - stats->GetStartOffset(); } - UNIT_ASSERT_EQUAL(messagesWritten, 150); + UNIT_ASSERT_EQUAL(messagesWritten, 50); } UNIT_ASSERT(producer->Close(TDuration::Seconds(1)).IsSuccess()); } From 0ac09623a2e8be8bd00361092f2d0dbdef02f8da Mon Sep 17 00:00:00 2001 From: Vitaliy Filippov Date: Mon, 23 Mar 2026 10:00:45 +0000 Subject: [PATCH 72/93] Add JSON inverted index type (#35061) --- .github/last_commit.txt | 2 +- CHANGELOG.md | 2 ++ include/ydb-cpp-sdk/client/table/table.h | 3 +++ include/ydb-cpp-sdk/client/table/table_enum.h | 1 + src/api/protos/ydb_table.proto | 6 ++++++ src/client/table/table.cpp | 20 +++++++++++++++++++ 6 files changed, 33 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 6399adab2f6..e0c8b061335 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -748f926df551161c7b935fba9cdbf64ed5590fbb +b96222681ac5c4d55815698c9bc8a85fe7f44396 diff --git a/CHANGELOG.md b/CHANGELOG.md index 599b3e6c047..49353fbbda6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +* Added support for the new inverted index type: JSON, intended to speed-up queries on Json or JsonDocument columns. + * EXPERIMENTAL! Added `IProducer` interface to the SDK. This interface is used to write messages to a topic. Each message can be associated with a partitioning key, which is used to determine the partition to which the message will be written. diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index ab3466100df..26aca1647a8 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -815,6 +815,9 @@ class TTableDescription { // fulltext void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings); void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings); + // json + void AddJsonIndex(const std::string& indexName, const std::vector& indexColumns); + void AddJsonIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); // default void AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); diff --git a/include/ydb-cpp-sdk/client/table/table_enum.h b/include/ydb-cpp-sdk/client/table/table_enum.h index a53e9e1e86e..e9ed2d01123 100644 --- a/include/ydb-cpp-sdk/client/table/table_enum.h +++ b/include/ydb-cpp-sdk/client/table/table_enum.h @@ -37,6 +37,7 @@ enum class EIndexType { GlobalVectorKMeansTree, GlobalFulltextPlain, GlobalFulltextRelevance, + GlobalJson, Unknown = std::numeric_limits::max() }; diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index 86f806de757..05102c9e7ba 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -328,6 +328,10 @@ message LocalBloomNgramFilterIndex { optional bool case_sensitive = 5; } +message GlobalJsonIndex { + GlobalIndexSettings settings = 1; +} + // Represent table index message TableIndex { // Name of index @@ -344,6 +348,7 @@ message TableIndex { GlobalFulltextRelevanceIndex global_fulltext_relevance_index = 9; LocalBloomFilterIndex local_bloom_filter_index = 10; LocalBloomNgramFilterIndex local_bloom_ngram_filter_index = 11; + GlobalJsonIndex global_json_index = 12; } // list of columns content to be copied in to index table repeated string data_columns = 5; @@ -372,6 +377,7 @@ message TableIndexDescription { GlobalFulltextRelevanceIndex global_fulltext_relevance_index = 11; LocalBloomFilterIndex local_bloom_filter_index = 12; LocalBloomNgramFilterIndex local_bloom_ngram_filter_index = 13; + GlobalJsonIndex global_json_index = 14; } Status status = 4; // list of columns content to be copied in to index table diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 63bc32c6ce8..b24afacee6a 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -800,6 +800,14 @@ void TTableDescription::AddFulltextIndex(const std::string& indexName, EIndexTyp Impl_->AddFulltextIndex(indexName, indexType, indexColumns, dataColumns, indexSettings); } +void TTableDescription::AddJsonIndex(const std::string& indexName, const std::vector& indexColumns) { + AddSecondaryIndex(indexName, EIndexType::GlobalJson, indexColumns); +} + +void TTableDescription::AddJsonIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns) { + AddSecondaryIndex(indexName, EIndexType::GlobalJson, indexColumns, dataColumns); +} + void TTableDescription::AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns) { AddSyncSecondaryIndex(indexName, indexColumns); } @@ -2796,6 +2804,10 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { specializedIndexSettings = TFulltextIndexSettings::FromProto(fulltextProto.fulltext_settings()); break; } + case TProto::kGlobalJsonIndex: + type = EIndexType::GlobalJson; + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(proto.global_json_index().settings())); + break; default: // fallback to global sync type = EIndexType::GlobalSync; globalIndexSettings.resize(1); @@ -2886,6 +2898,13 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { } break; } + case EIndexType::GlobalJson: { + auto& settings = *proto.mutable_global_json_index()->mutable_settings(); + if (GlobalIndexSettings_.size() == 1) { + GlobalIndexSettings_.at(0).SerializeTo(settings); + } + break; + } case EIndexType::Unknown: break; } @@ -2911,6 +2930,7 @@ void TIndexDescription::Out(IOutputStream& o) const { case EIndexType::GlobalSync: case EIndexType::GlobalAsync: case EIndexType::GlobalUnique: + case EIndexType::GlobalJson: case EIndexType::Unknown: break; case EIndexType::GlobalVectorKMeansTree: From 74884de0c592abe9d16ab1a7d34658fc60b37a6d Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:00:52 +0000 Subject: [PATCH 73/93] LOGBROKER-10206 Fix include (#35825) --- .github/last_commit.txt | 2 +- src/client/topic/impl/producer.cpp | 6 ++---- src/client/topic/impl/producer.h | 2 -- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index e0c8b061335..38cf0e8ffac 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -b96222681ac5c4d55815698c9bc8a85fe7f44396 +4a175c015a727d34a540ffeebf10b5b4eced7c3d diff --git a/src/client/topic/impl/producer.cpp b/src/client/topic/impl/producer.cpp index f43eb71a3b7..adbb23fbc74 100644 --- a/src/client/topic/impl/producer.cpp +++ b/src/client/topic/impl/producer.cpp @@ -5,6 +5,8 @@ #include #include +#include + namespace NYdb::inline V3::NTopic { //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -2046,10 +2048,6 @@ void TProducer::HandleAutoPartitioning(std::uint32_t partition) { SplittedPartitionWorkers.try_emplace(partition, splittedPartitionWorker); } -std::string TProducer::GetProducerId(std::uint32_t partition) { - return std::format("{}_{}", Settings.ProducerIdPrefix_, partition); -} - TWriterCounters::TPtr TProducer::GetCounters() { return nullptr; } diff --git a/src/client/topic/impl/producer.h b/src/client/topic/impl/producer.h index 5760567aa9e..637fe4b2709 100644 --- a/src/client/topic/impl/producer.h +++ b/src/client/topic/impl/producer.h @@ -364,8 +364,6 @@ class TProducer : public IProducer, TDuration GetCloseTimeout(); - std::string GetProducerId(std::uint32_t partition); - void HandleAutoPartitioning(std::uint32_t partition); bool RunSplittedPartitionWorkers(); From 8dce6ad2406167259c7bb8362da5529abdf06222 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:00:59 +0000 Subject: [PATCH 74/93] LOGBROKER-10206 Add check in tests (#35861) --- .github/last_commit.txt | 2 +- src/client/topic/ut/basic_usage_ut.cpp | 22 +++++++++++++--------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 38cf0e8ffac..aa238a24c81 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -4a175c015a727d34a540ffeebf10b5b4eced7c3d +854ba0a535fa6b55e90c353f5ba524d2d16ad7f4 diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index 682bcf21f63..58eeb076a97 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -73,12 +73,14 @@ std::string Serialize(const TExample& value) { void ReadMessagesAndAssertOrderedBySeqNo(TTopicClient& client, const std::string& topicPath, const std::string& consumerName, + const std::string& expectedPayload, size_t expectedCount, TDuration timeout = TDuration::Seconds(30)) { struct TMessageInfo { ui64 PartitionId; - TString ProducerId; + std::string ProducerId; ui64 SeqNo; + std::string Data; }; std::vector messages; messages.reserve(expectedCount); @@ -98,6 +100,7 @@ void ReadMessagesAndAssertOrderedBySeqNo(TTopicClient& client, msg.GetPartitionSession()->GetPartitionId(), TString(msg.GetProducerId()), msg.GetSeqNo(), + TString(msg.GetData()), }); } if (messages.size() >= expectedCount) { @@ -115,8 +118,9 @@ void ReadMessagesAndAssertOrderedBySeqNo(TTopicClient& client, // SeqNo ordering is guaranteed within one producer stream. // Multiple producers can write into the same partition with independent seqNo sequences. - std::map, std::vector> byPartitionAndProducer; + std::map, std::vector> byPartitionAndProducer; for (const auto& m : messages) { + UNIT_ASSERT_VALUES_EQUAL(m.Data, expectedPayload); byPartitionAndProducer[{m.PartitionId, m.ProducerId}].push_back(m.SeqNo); } for (const auto& [key, seqNos] : byPartitionAndProducer) { @@ -1324,11 +1328,11 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.MaxBlock(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); + std::string payload = "data"; const ui64 count = 3000; for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); - std::string payload = "data"; TWriteMessage msg(payload); msg.SeqNo(i); msg.Key(key); @@ -1338,7 +1342,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); - ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), count); + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), payload, count); } Y_UNIT_TEST(Producer_AutoSeqNo) { @@ -1372,9 +1376,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); auto producer2 = client.CreateProducer(writeSettings); + std::string payload = "data"; for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); - std::string payload = "data"; TWriteMessage msg(payload); msg.Key(key); UNIT_ASSERT_C(producer2->Write(std::move(msg)).IsSuccess(), "Failed to write message"); @@ -1382,7 +1386,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); UNIT_ASSERT_C(producer2->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); - ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), count * 2); + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), payload, count * 2); } Y_UNIT_TEST(Producer_WriteToClosedProducer) { @@ -1721,7 +1725,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); UNIT_ASSERT(producer3->Close(TDuration::Seconds(30)).IsSuccess()); - ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), 14); + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), msgData, 14); } Y_UNIT_TEST(AutoPartitioning_Producer_SmallMessages) { @@ -1799,7 +1803,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(producer1->Close(TDuration::Seconds(30)).IsSuccess()); UNIT_ASSERT(producer2->Close(TDuration::Seconds(30)).IsSuccess()); - ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), totalMessages); + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), msgData, totalMessages); } Y_UNIT_TEST(Producer_BasicWrite) { @@ -1837,7 +1841,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_EQUAL(messagesWritten, 100); UNIT_ASSERT_C(producer->Close(TDuration::Seconds(1)).IsSuccess(), "Failed to close producer"); - ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), 100); + ReadMessagesAndAssertOrderedBySeqNo(client, setup.GetTopicPath(TEST_TOPIC), setup.GetConsumerName(), msgData, 100); } Y_UNIT_TEST(TypedProducer_BasicWrite) { From b2b29d59deea375c8eeb9b5d4be8fa0afb6fb662 Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Mon, 23 Mar 2026 10:01:14 +0000 Subject: [PATCH 75/93] Removed EnableTopicServiceTx, EnableTopicSplitMerge, EnableTopicMessageMeta, EnableTopicTransfer, EnableTopicAutopartitioningForCDC feature flags (#35881) --- .github/last_commit.txt | 2 +- src/client/topic/ut/local_partition_ut.cpp | 1 - src/client/topic/ut/ut_utils/txusage_fixture.cpp | 2 -- 3 files changed, 1 insertion(+), 4 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index aa238a24c81..3e9d75da56c 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -854ba0a535fa6b55e90c353f5ba524d2d16ad7f4 +30acea6972011a34d99f41ce1258b90dc65c3b47 diff --git a/src/client/topic/ut/local_partition_ut.cpp b/src/client/topic/ut/local_partition_ut.cpp index b0f94be2c8b..12f1c0dd180 100644 --- a/src/client/topic/ut/local_partition_ut.cpp +++ b/src/client/topic/ut/local_partition_ut.cpp @@ -34,7 +34,6 @@ namespace NYdb::inline V3::NTopic::NTests { TTopicSdkTestSetup CreateSetupForSplitMerge(const std::string& testCaseName) { NKikimrConfig::TFeatureFlags ff; - ff.SetEnableTopicSplitMerge(true); auto settings = TTopicSdkTestSetup::MakeServerSettings(); settings.SetFeatureFlags(ff); auto setup = TTopicSdkTestSetup(testCaseName, settings, false); diff --git a/src/client/topic/ut/ut_utils/txusage_fixture.cpp b/src/client/topic/ut/ut_utils/txusage_fixture.cpp index f92a261d439..3a601c0dcbb 100644 --- a/src/client/topic/ut/ut_utils/txusage_fixture.cpp +++ b/src/client/topic/ut/ut_utils/txusage_fixture.cpp @@ -36,8 +36,6 @@ TFixture::TTableRecord::TTableRecord(const std::string& key, const std::string& void TFixture::SetUp(NUnitTest::TTestContext&) { NKikimr::Tests::TServerSettings settings = TTopicSdkTestSetup::MakeServerSettings(); - settings.SetEnableTopicServiceTx(true); - settings.SetEnableTopicSplitMerge(true); settings.SetEnableHtapTx(GetEnableHtapTx()); settings.SetAllowOlapDataQuery(GetAllowOlapDataQuery()); From cc926ac52b45f05a1b74dc38ccc2183827bdf010 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:01:21 +0000 Subject: [PATCH 76/93] LOGBROKER-9648 Fix partition session id conflict (#35929) --- .github/last_commit.txt | 2 +- .../federated_topic/ut/basic_usage_ut.cpp | 78 +++++++++++++++++++ src/client/topic/impl/event_handlers.cpp | 40 ++++++---- 3 files changed, 103 insertions(+), 17 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 3e9d75da56c..36f20194ba5 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -30acea6972011a34d99f41ce1258b90dc65c3b47 +5e505b5bf44e9655b84f4690ad7431e6422d959b diff --git a/src/client/federated_topic/ut/basic_usage_ut.cpp b/src/client/federated_topic/ut/basic_usage_ut.cpp index ce87d140bdc..5d46c292998 100644 --- a/src/client/federated_topic/ut/basic_usage_ut.cpp +++ b/src/client/federated_topic/ut/basic_usage_ut.cpp @@ -20,6 +20,38 @@ namespace NYdb::NFederatedTopic::NTests { +void WriteMessages(std::shared_ptr setup, const TString& path, const TString& messageBase, size_t count) { + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(path).MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + IExecutor::TPtr executor = NPersQueue::CreateSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& pqClient = setup->GetPersQueueClient(); + auto writeSession = pqClient.CreateSimpleBlockingWriteSession(writeSettings); + + for (size_t i = 0; i < count; ++i) { + UNIT_ASSERT(writeSession->Write(messageBase + ToString(i))); + } + writeSession->Close(); +} + +void WriteMessages(std::shared_ptr setup, const TString& messageBase, size_t count) { + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic()).MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + IExecutor::TPtr executor = NPersQueue::CreateSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& pqClient = setup->GetPersQueueClient(); + auto writeSession = pqClient.CreateSimpleBlockingWriteSession(writeSettings); + + for (size_t i = 0; i < count; ++i) { + UNIT_ASSERT(writeSession->Write(messageBase + ToString(i))); + } + writeSession->Close(); +} + Y_UNIT_TEST_SUITE(BasicUsage) { Y_UNIT_TEST(GetAllStartPartitionSessions) { @@ -403,6 +435,52 @@ Y_UNIT_TEST_SUITE(BasicUsage) { ReadSession->Close(TDuration::MilliSeconds(10)); } + Y_UNIT_TEST(SimpleDataHandlersAndGetEvent) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + const TString topic1 = setup->GetTestTopic(); + const TString topic2 = setup->GetTestTopic() + "-second"; + setup->CreateTopic(topic2, setup->GetLocalCluster()); + + auto driverConfig = NYdb::TDriverConfig() + .SetEndpoint(TStringBuilder() << "localhost:" << setup->GetGrpcPort()) + .SetDatabase("/Root"); + NYdb::TDriver driver(driverConfig); + NYdb::NFederatedTopic::TFederatedTopicClient client(driver); + + TString messageBase = "hello-"; + WriteMessages(setup, topic1, messageBase, 5); + WriteMessages(setup, topic2, messageBase + "t2-", 3); + + TVector receivedMessages; + NYdb::NFederatedTopic::TFederatedReadSessionSettings settings; + settings + .ConsumerName("test-consumer") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(NYdb::NTopic::TTopicReadSettings(topic1)) + .AppendTopics(NYdb::NTopic::TTopicReadSettings(topic2)); + + settings.EventHandlers_.SimpleDataHandlers( + [&receivedMessages](NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent& event) { + for (const auto& message: event.GetMessages()) { + receivedMessages.push_back(TString(message.GetData())); + } + }, + true // commit on receive + ); + + auto session = client.CreateReadSession(settings); + std::jthread thread([&] { + Sleep(TDuration::Seconds(3)); + UNIT_ASSERT(session->Close(TDuration::MilliSeconds(10))); + }); + + auto event = session->GetEvent(/* block = */true); + UNIT_ASSERT(event.has_value()); + UNIT_ASSERT_VALUES_EQUAL(receivedMessages.size(), 8); + } + Y_UNIT_TEST(FallbackToSingleDbAfterBadRequest) { auto setup = std::make_shared(TEST_CASE_NAME, false); setup->Start(true, true); diff --git a/src/client/topic/impl/event_handlers.cpp b/src/client/topic/impl/event_handlers.cpp index 1a4fad20fa0..e91024f6775 100644 --- a/src/client/topic/impl/event_handlers.cpp +++ b/src/client/topic/impl/event_handlers.cpp @@ -1,6 +1,7 @@ #include #include +#include namespace NYdb::inline V3::NTopic { @@ -22,7 +23,8 @@ class TGracefulReleasingSimpleDataHandlers : public TThrRefBase { TDeferredCommit deferredCommit; { std::lock_guard guard(Lock); - auto& offsetSet = PartitionStreamToUncommittedOffsets[event.GetPartitionSession()->GetPartitionSessionId()]; + const std::string key = GetEventKey(event); + auto& offsetSet = PartitionStreamToUncommittedOffsets[key]; // Messages could contain holes in offset, but later commit ack will tell us right border. // So we can easily insert the whole interval with holes included. // It will be removed from set by specifying proper right border. @@ -41,16 +43,16 @@ class TGracefulReleasingSimpleDataHandlers : public TThrRefBase { void OnCommitAcknowledgement(TReadSessionEvent::TCommitOffsetAcknowledgementEvent& event) { std::lock_guard guard(Lock); - const ui64 partitionStreamId = event.GetPartitionSession()->GetPartitionSessionId(); - auto& offsetSet = PartitionStreamToUncommittedOffsets[partitionStreamId]; + const std::string key = GetEventKey(event); + auto& offsetSet = PartitionStreamToUncommittedOffsets[key]; if (offsetSet.EraseInterval(0, event.GetCommittedOffset() + 1)) { // Remove some offsets. if (offsetSet.Empty()) { // No offsets left. - auto unconfirmedDestroyIt = UnconfirmedDestroys.find(partitionStreamId); + auto unconfirmedDestroyIt = UnconfirmedDestroys.find(key); if (unconfirmedDestroyIt != UnconfirmedDestroys.end()) { // Confirm and forget about this partition stream. unconfirmedDestroyIt->second.Confirm(); UnconfirmedDestroys.erase(unconfirmedDestroyIt); - PartitionStreamToUncommittedOffsets.erase(partitionStreamId); + PartitionStreamToUncommittedOffsets.erase(key); } } } @@ -59,20 +61,21 @@ class TGracefulReleasingSimpleDataHandlers : public TThrRefBase { void OnCreatePartitionStream(TReadSessionEvent::TStartPartitionSessionEvent& event) { { std::lock_guard guard(Lock); - Y_ABORT_UNLESS(PartitionStreamToUncommittedOffsets[event.GetPartitionSession()->GetPartitionSessionId()].Empty()); + const std::string key = GetEventKey(event); + Y_ABORT_UNLESS(PartitionStreamToUncommittedOffsets[key].Empty()); } event.Confirm(); } void OnDestroyPartitionStream(TReadSessionEvent::TStopPartitionSessionEvent& event) { std::lock_guard guard(Lock); - const ui64 partitionStreamId = event.GetPartitionSession()->GetPartitionSessionId(); - Y_ABORT_UNLESS(UnconfirmedDestroys.find(partitionStreamId) == UnconfirmedDestroys.end()); - if (PartitionStreamToUncommittedOffsets[partitionStreamId].Empty()) { - PartitionStreamToUncommittedOffsets.erase(partitionStreamId); + auto key = GetEventKey(event); + Y_ABORT_UNLESS(UnconfirmedDestroys.find(key) == UnconfirmedDestroys.end()); + if (PartitionStreamToUncommittedOffsets[key].Empty()) { + PartitionStreamToUncommittedOffsets.erase(key); event.Confirm(); } else { - UnconfirmedDestroys.emplace(partitionStreamId, std::move(event)); + UnconfirmedDestroys.emplace(key, std::move(event)); } } @@ -82,17 +85,22 @@ class TGracefulReleasingSimpleDataHandlers : public TThrRefBase { void OnPartitionStreamClosed(TReadSessionEvent::TPartitionSessionClosedEvent& event) { std::lock_guard guard(Lock); - const ui64 partitionStreamId = event.GetPartitionSession()->GetPartitionSessionId(); - PartitionStreamToUncommittedOffsets.erase(partitionStreamId); - UnconfirmedDestroys.erase(partitionStreamId); + const std::string key = GetEventKey(event); + PartitionStreamToUncommittedOffsets.erase(key); + UnconfirmedDestroys.erase(key); } private: + template + std::string GetEventKey(const TEvent& event) { + return event.GetPartitionSession()->GetReadSessionId() + "_" + ToString(event.GetPartitionSession()->GetPartitionSessionId()); + } + TAdaptiveLock Lock; // For the case when user gave us multithreaded executor. const std::function DataHandler; const bool CommitAfterProcessing; - std::unordered_map> PartitionStreamToUncommittedOffsets; // Partition stream id -> set of offsets. - std::unordered_map UnconfirmedDestroys; // Partition stream id -> destroy events. + std::unordered_map> PartitionStreamToUncommittedOffsets; // Session id + Partition stream id -> set of offsets. + std::unordered_map UnconfirmedDestroys; // Session id + Partition stream id -> destroy events. }; TReadSessionSettings::TEventHandlers& TReadSessionSettings::TEventHandlers::SimpleDataHandlers(std::function dataHandler, From 6e20bd565f75a5c4e31fd26cf72ca807dbfa1234 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:01:28 +0000 Subject: [PATCH 77/93] LOGBROKER-10206 Better interface (changes for support) (#35932) --- .github/last_commit.txt | 2 +- .../producer/basic_write/main.cpp | 4 +- include/ydb-cpp-sdk/client/topic/producer.h | 15 +- src/client/topic/impl/producer.cpp | 31 +-- src/client/topic/impl/producer.h | 4 +- src/client/topic/ut/basic_usage_ut.cpp | 210 +++++------------- tests/integration/topic/basic_usage_it.cpp | 9 +- 7 files changed, 89 insertions(+), 186 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 36f20194ba5..13a850091f5 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -5e505b5bf44e9655b84f4690ad7431e6422d959b +59c5b07fc52bfc4e047331be4b65156c9303227a diff --git a/examples/topic_writer/producer/basic_write/main.cpp b/examples/topic_writer/producer/basic_write/main.cpp index 83af470330b..47629a00f06 100644 --- a/examples/topic_writer/producer/basic_write/main.cpp +++ b/examples/topic_writer/producer/basic_write/main.cpp @@ -11,7 +11,7 @@ std::shared_ptr CreateProducer(const std::string& topic producerSettings.ProducerIdPrefix("producer_basic"); producerSettings.PartitionChooserStrategy(NYdb::NTopic::TProducerSettings::EPartitionChooserStrategy::Bound); producerSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - producerSettings.MaxBlock(TDuration::Seconds(30)); + producerSettings.MaxBlockTimeout(TDuration::Seconds(30)); producerSettings.MaxMemoryUsage(100_MB); return topicClient.CreateProducer(producerSettings); } @@ -66,7 +66,7 @@ void WriteWithHandlingResult(std::shared_ptr producer, for (size_t retries = 0; retries < MAX_RETRIES; retries++) { auto writeResult = producer->Write(std::move(writeMessage)); - if (writeResult.IsSuccess()) { + if (writeResult.IsQueued()) { // if write was successful, we can continue writing messages continue; } diff --git a/include/ydb-cpp-sdk/client/topic/producer.h b/include/ydb-cpp-sdk/client/topic/producer.h index de48becd82e..3efb9c6aa23 100644 --- a/include/ydb-cpp-sdk/client/topic/producer.h +++ b/include/ydb-cpp-sdk/client/topic/producer.h @@ -12,7 +12,7 @@ struct TProducerSettings : public TWriteSessionSettings { enum class EPartitionChooserStrategy { Bound, - Hash, + KafkaHash, }; TProducerSettings() = default; @@ -39,14 +39,9 @@ struct TProducerSettings : public TWriteSessionSettings { //! ProducerId is generated as ProducerIdPrefix + partition id. FLUENT_SETTING(std::string, ProducerIdPrefix); - //! SessionID to use. - FLUENT_SETTING_DEFAULT(std::string, SessionId, ""); - - //! Maximum block time for write. If set, write will block for up to MaxBlockMs when the buffer is overloaded. - FLUENT_SETTING_DEFAULT(TDuration, MaxBlock, TDuration::Zero()); - - //! Key producer function. - FLUENT_SETTING_OPTIONAL(std::function, KeyProducer); + //! Maximum block timeout for write. If set, write will block for up to MaxBlockTimeout when the buffer is overloaded. + //! If not set, Write will block until the message is written to the buffer. + FLUENT_SETTING_DEFAULT(TDuration, MaxBlockTimeout, TDuration::Max()); private: using TWriteSessionSettings::ProducerId; @@ -88,7 +83,7 @@ struct TWriteResult { //! Value is std::nullopt if the session is not closed. std::optional ClosedDescription; - bool IsSuccess() const { + bool IsQueued() const { return Status == EWriteStatus::Queued; } diff --git a/src/client/topic/impl/producer.cpp b/src/client/topic/impl/producer.cpp index adbb23fbc74..766f71abe4e 100644 --- a/src/client/topic/impl/producer.cpp +++ b/src/client/topic/impl/producer.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include @@ -1428,7 +1429,8 @@ TProducer::TProducer( std::shared_ptr client, std::shared_ptr connections, TDbDriverStatePtr dbDriverState) - : Connections(connections), + : Id(CreateGuidAsString()), + Connections(connections), Client(client), DbDriverState(dbDriverState), Metrics(this), @@ -1457,7 +1459,7 @@ TProducer::TProducer( return a.GetPartitionId() < b.GetPartitionId(); }); - PartitionChooserStrategy = settings.PartitionChooserStrategy_; + auto partitionChooserStrategy = settings.PartitionChooserStrategy_; auto strategy = topicConfig.GetTopicDescription().GetPartitioningSettings().GetAutoPartitioningSettings().GetStrategy(); auto autoPartitioningEnabled = (strategy != EAutoPartitioningStrategy::Disabled && strategy != EAutoPartitioningStrategy::Unspecified); @@ -1496,7 +1498,7 @@ TProducer::TProducer( } } - switch (PartitionChooserStrategy) { + switch (partitionChooserStrategy) { case TProducerSettings::EPartitionChooserStrategy::Bound: PartitioningKeyHasher = settings.PartitioningKeyHasher_; PartitionChooser = std::make_unique(this); @@ -1513,9 +1515,9 @@ TProducer::TProducer( PartitionsIndex[partition.GetFromBound().value_or("")] = partition.GetPartitionId(); } break; - case TProducerSettings::EPartitionChooserStrategy::Hash: + case TProducerSettings::EPartitionChooserStrategy::KafkaHash: if (autoPartitioningEnabled) { - throw TContractViolation("Hash partition chooser strategy is not supported with auto partitioning enabled"); + throw TContractViolation("KafkaHash partition chooser strategy is not supported with auto partitioning enabled"); } std::vector partitionsIds; @@ -1628,6 +1630,11 @@ void TProducer::SetCloseDeadline(const TDuration& closeTimeout) { TProducer::~TProducer() { auto _ = Close(TDuration::Zero()); // Ignore the result, because we are destroying the producer Settings.EventHandlers_.HandlersExecutor_->Stop(); + + if (MainWorkerState.load() == 0) { + ShutdownPromise.TrySetValue(); + } + ShutdownFuture.Wait(); } @@ -1791,7 +1798,7 @@ void TProducer::GetSessionClosedEventAndDie(WrappedWriteSessionPtr wrappedSessio } TStringBuilder TProducer::LogPrefix() { - return TStringBuilder() << " SessionId: " << Settings.SessionId_ << " Epoch: " << Epoch.load() << " "; + return TStringBuilder() << " Id: " << Id << " Epoch: " << Epoch.load() << " "; } void TProducer::NextEpoch() { @@ -1942,14 +1949,8 @@ TWriteResult TProducer::WriteInternal(TContinuationToken&&, TWriteMessage&& mess chosenPartition = message.Partition_.value(); } else if (!message.Key_.has_value()) { - std::string key; - if (Settings.KeyProducer_) { - key = (*Settings.KeyProducer_)(message); - } else { - key = Settings.ProducerIdPrefix_; - } - message.Key(key); - chosenPartition = PartitionChooser->ChoosePartition(key); + message.Key(Settings.ProducerIdPrefix_); + chosenPartition = PartitionChooser->ChoosePartition(Settings.ProducerIdPrefix_); } else { chosenPartition = PartitionChooser->ChoosePartition(*message.Key_); } @@ -1970,7 +1971,7 @@ TWriteResult TProducer::WriteInternal(TContinuationToken&&, TWriteMessage&& mess } TWriteResult TProducer::Write(TWriteMessage&& message) { - auto remainingTimeout = Settings.MaxBlock_; + auto remainingTimeout = Settings.MaxBlockTimeout_; auto sleepTimeMs = DEFAULT_START_BLOCK_TIMEOUT; for (;;) { if (Closed.load()) { diff --git a/src/client/topic/impl/producer.h b/src/client/topic/impl/producer.h index 637fe4b2709..dca301c72a0 100644 --- a/src/client/topic/impl/producer.h +++ b/src/client/topic/impl/producer.h @@ -419,6 +419,9 @@ class TProducer : public IProducer, ~TProducer(); private: + // for logging + std::string Id; + std::shared_ptr Connections; std::shared_ptr Client; TDbDriverStatePtr DbDriverState; @@ -430,7 +433,6 @@ class TProducer : public IProducer, TProducerSettings Settings; ESeqNoStrategy SeqNoStrategy = ESeqNoStrategy::NotInitialized; - TProducerSettings::EPartitionChooserStrategy PartitionChooserStrategy = TProducerSettings::EPartitionChooserStrategy::Hash; NThreading::TPromise ClosePromise; NThreading::TFuture CloseFuture; diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index 58eeb076a97..aed9af75b63 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -1064,9 +1064,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); std::atomic acksCount{0}; std::atomic closedCount{0}; @@ -1093,7 +1093,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i + 1); msg.Key("key-" + ToString(i)); - UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(session->Close(TDuration::Seconds(30)).IsSuccess(), "Failed to close keyed write session"); @@ -1110,7 +1110,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); UNIT_ASSERT_EXCEPTION(setup.MakeClient().CreateProducer(writeSettings), TContractViolation); @@ -1130,7 +1130,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); UNIT_ASSERT_EXCEPTION(federatedLikeClient.CreateProducer(writeSettings), TContractViolation); @@ -1146,9 +1146,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto session = publicClient.CreateProducer(writeSettings); @@ -1156,7 +1156,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(0); msg.Key("key"); - UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); auto flushResult = session->Flush().GetValueSync(); UNIT_ASSERT_C(flushResult.IsClosed(), "Failed to flush producer"); UNIT_ASSERT_C(flushResult.ClosedDescription->GetStatus() == EStatus::BAD_REQUEST, "Status is not BAD_REQUEST"); @@ -1183,9 +1183,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto session = publicClient.CreateProducer(writeSettings); @@ -1201,14 +1201,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(seqNo++); msg.Key(key0); - UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } for (ui64 i = 0; i < count1; ++i) { std::string payload = "msg1"; TWriteMessage msg(payload); msg.SeqNo(seqNo++); msg.Key(key1); - UNIT_ASSERT_C(session->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(session->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close keyed write session"); @@ -1266,7 +1266,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = publicClient.CreateProducer(writeSettings); auto rawProducer = std::dynamic_pointer_cast(producer); @@ -1290,7 +1290,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i + 1); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); @@ -1324,8 +1324,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); std::string payload = "data"; @@ -1336,7 +1336,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); @@ -1357,8 +1357,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1369,7 +1369,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); @@ -1381,7 +1381,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto key = CreateGuidAsString(); TWriteMessage msg(payload); msg.Key(key); - UNIT_ASSERT_C(producer2->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer2->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); UNIT_ASSERT_C(producer2->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); @@ -1401,8 +1401,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(10)); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1413,7 +1413,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); @@ -1441,8 +1441,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1464,7 +1464,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { msg.Key(key); auto writeResult = producer->Write(std::move(msg)); UNIT_ASSERT_C( - writeResult.IsSuccess(), + writeResult.IsQueued(), "Failed to write message" ); } @@ -1491,8 +1491,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(5)); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1505,7 +1505,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(seqNo++); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); @@ -1517,14 +1517,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(seqNo++); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); } - Y_UNIT_TEST(KeyedWriteSession_BoundPartitionChooser_SplitPartition_MultiThreadedAcksOrder) { + Y_UNIT_TEST(Producer_BoundPartitionChooser_SplitPartition_MultiThreadedAcksOrder) { NKikimr::NPQ::NTest::TTopicSdkTestSetup setup = NKikimr::NPQ::NTest::CreateSetup(); setup.CreateTopicWithAutoscale(TEST_TOPIC, TEST_CONSUMER, 1, 100); @@ -1540,7 +1540,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1553,7 +1553,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i); msg.Key(key); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } }); @@ -1570,7 +1570,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer->Close(TDuration::Seconds(30)).IsSuccess(), "Failed to close producer"); } - Y_UNIT_TEST(KeyedWriteSession_CloseTimeout) { + Y_UNIT_TEST(Producer_CloseTimeout) { TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 3); @@ -1581,9 +1581,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { .Path(setup.GetTopicPath(TEST_TOPIC)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); @@ -1592,7 +1592,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { TWriteMessage msg(payload); msg.SeqNo(i + 1); msg.Key("key1"); - UNIT_ASSERT_C(producer->Write(std::move(msg)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } // Test Close timeout @@ -1604,7 +1604,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { // block longer than the timeout and that the session eventually closes successfully. UNIT_ASSERT_C( result.IsSuccess() || result.IsTimeout(), - TStringBuilder() << "Failed to close keyed write session, status: " << static_cast(result.Status) + TStringBuilder() << "Failed to close producer, status: " << static_cast(result.Status) ); const TDuration actualDuration = TInstant::Now() - startTime; @@ -1634,7 +1634,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings1.ProducerIdPrefix("autopartitioning_keyed_1"); writeSettings1.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings1.MaxBlock(TDuration::Seconds(30)); + writeSettings1.MaxBlockTimeout(TDuration::Seconds(30)); TProducerSettings writeSettings2 = writeSettings1; writeSettings2.ProducerIdPrefix("autopartitioning_keyed_2"); @@ -1653,7 +1653,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { if (key.empty()) { key = "lalala"; } - UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsQueued(), "Failed to write message"); }; { @@ -1746,7 +1746,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings1.ProducerIdPrefix("autopartitioning_keyed_small_1"); writeSettings1.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Bound); writeSettings1.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings1.MaxBlock(TDuration::Seconds(30)); + writeSettings1.MaxBlockTimeout(TDuration::Seconds(30)); TProducerSettings writeSettings2 = writeSettings1; writeSettings2.ProducerIdPrefix("autopartitioning_keyed_small_2"); @@ -1765,7 +1765,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto writeMessage = [&](std::shared_ptr s, std::string_view payload, ui64 seqNo) { auto key = keys[seqNo % keys.size()]; if (key.empty()) key = "a"; - UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsSuccess(), "Failed to write message"); + UNIT_ASSERT_C(s->Write(CreateMessage(payload, key, seqNo)).IsQueued(), "Failed to write message"); }; { @@ -1817,13 +1817,13 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); writeSettings.Codec(ECodec::RAW); writeSettings.ProducerIdPrefix("producer_basic_write"); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); auto producer = client.CreateProducer(writeSettings); auto msgData = TString(10_KB, 'a'); for (ui64 i = 0; i < 100; ++i) { - UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsQueued()); } UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); @@ -1857,11 +1857,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); writeSettings.Codec(ECodec::RAW); writeSettings.ProducerIdPrefix("producer_basic_write"); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.KeyProducer([](const TWriteMessage& message) -> std::string { - return ToString(MurmurHash(message.Data.data(), message.Data.size())); - }); - writeSettings.MaxBlock(TDuration::Seconds(1)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(1)); auto producer = client.CreateTypedProducer(writeSettings); @@ -1870,7 +1867,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < messageCount; ++i) { auto payload = CreateGuidAsString(); sentPayloads.push_back(payload); - UNIT_ASSERT(producer->Write(TExample{.Payload = payload}).IsSuccess()); + UNIT_ASSERT(producer->Write(TExample{.Payload = payload}).IsQueued()); } UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); @@ -1917,9 +1914,9 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); writeSettings.Codec(ECodec::RAW); writeSettings.ProducerIdPrefix("producer_basic_write"); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::MilliSeconds(500)); - writeSettings.MaxBlock(TDuration::Seconds(1)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(1)); auto describeResult = client.DescribeTopic(TEST_TOPIC).GetValueSync(); const auto& partitions = describeResult.GetTopicDescription().GetPartitions(); @@ -1932,7 +1929,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < 10; ++i) { TWriteMessage msg(msgData); msg.Partition(partition.GetPartitionId()); - UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); + UNIT_ASSERT(producer->Write(std::move(msg)).IsQueued()); } } UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); @@ -1963,99 +1960,6 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT(producer->Close(TDuration::Seconds(1)).IsSuccess()); } - Y_UNIT_TEST(Producer_CustomKeyProducerFunction) { - TTopicSdkTestSetup setup{TEST_CASE_NAME, TTopicSdkTestSetup::MakeServerSettings(), false}; - setup.CreateTopic(TEST_TOPIC, TEST_CONSUMER, 2); - - // Capture partition ids in the same order as DescribeTopic returns them - // (the keyed session uses the same DescribeTopic ordering to map hash bucket -> partition id). - auto publicClient = setup.MakeClient(); - auto describeTopicSettings = TDescribeTopicSettings().IncludeStats(true); - auto before = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); - UNIT_ASSERT_C(before.IsSuccess(), before.GetIssues().ToOneLineString()); - const auto& beforePartitions = before.GetTopicDescription().GetPartitions(); - UNIT_ASSERT_VALUES_EQUAL(beforePartitions.size(), 2); - const ui64 partitionId0 = beforePartitions[0].GetPartitionId(); - const ui64 partitionId1 = beforePartitions[1].GetPartitionId(); - - constexpr auto keyAttributeName = "__key"; - - TProducerSettings writeSettings; - writeSettings - .Path(setup.GetTopicPath(TEST_TOPIC)) - .Codec(ECodec::RAW); - writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); - writeSettings.KeyProducer([](const TWriteMessage& message) -> std::string { - for (const auto& [attributeName, attributeValue] : message.MessageMeta_) { - if (attributeName == keyAttributeName) { - return attributeValue; - } - } - return ""; - }); - writeSettings.MaxBlock(TDuration::Seconds(1)); - - auto producer = publicClient.CreateProducer(writeSettings); - - const std::string key0 = FindKeyForBucket(0, 2); - const std::string key1 = FindKeyForBucket(1, 2); - - const ui64 count0 = 7; - const ui64 count1 = 11; - - auto seqNo = 1; - for (ui64 i = 0; i < count0; ++i) { - std::string payload = "msg0"; - TWriteMessage msg(payload); - msg.SeqNo(seqNo++); - msg.MessageMeta_.emplace_back(keyAttributeName, key0); - UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); - } - for (ui64 i = 0; i < count1; ++i) { - std::string payload = "msg1"; - TWriteMessage msg(payload); - msg.SeqNo(seqNo++); - msg.MessageMeta_.emplace_back(keyAttributeName, key1); - UNIT_ASSERT(producer->Write(std::move(msg)).IsSuccess()); - } - - UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); - UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().MessagesWritten, count0 + count1); - UNIT_ASSERT_VALUES_EQUAL(producer->GetWriteStats().LastWrittenSeqNo, seqNo - 1); - - auto after = publicClient.DescribeTopic(setup.GetTopicPath(TEST_TOPIC), describeTopicSettings).GetValueSync(); - UNIT_ASSERT_C(after.IsSuccess(), after.GetIssues().ToOneLineString()); - const auto& afterPartitions = after.GetTopicDescription().GetPartitions(); - UNIT_ASSERT_VALUES_EQUAL(afterPartitions.size(), 2); - - std::unordered_map endOffsets; - for (const auto& p : afterPartitions) { - auto stats = p.GetPartitionStats(); - UNIT_ASSERT(stats.has_value()); - endOffsets[p.GetPartitionId()] = stats->GetEndOffset(); - } - - auto it0 = endOffsets.find(partitionId0); - auto it1 = endOffsets.find(partitionId1); - UNIT_ASSERT(it0 != endOffsets.end()); - UNIT_ASSERT(it1 != endOffsets.end()); - - const ui64 endOffset0 = it0->second; - const ui64 endOffset1 = it1->second; - - // Partition ordering in DescribeTopic is not a part of public API contract, so allow swapping. - UNIT_ASSERT_VALUES_EQUAL(endOffset0 + endOffset1, count0 + count1); - UNIT_ASSERT_C( - (endOffset0 == count0 && endOffset1 == count1) || (endOffset0 == count1 && endOffset1 == count0), - TStringBuilder() << "Unexpected end offsets distribution: " - << "partitionId0=" << partitionId0 << " endOffset0=" << endOffset0 << ", " - << "partitionId1=" << partitionId1 << " endOffset1=" << endOffset1 << ", " - << "expected (" << count0 << "," << count1 << ") in any order" - ); - } - Y_UNIT_TEST(Producer_BlockingWrite) { auto settings = TTopicSdkTestSetup::MakeServerSettings(); settings.PQConfig.SetUseSrcIdMetaMappingInFirstClass(true); @@ -2067,14 +1971,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); writeSettings.Codec(ECodec::RAW); writeSettings.ProducerIdPrefix("simple_blocking_producer_basic_write"); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); - writeSettings.MaxBlock(TDuration::Seconds(1)); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); + writeSettings.MaxBlockTimeout(TDuration::Seconds(1)); auto producer = client.CreateProducer(writeSettings); auto msgData = TString(10_KB, 'a'); for (ui64 i = 0; i < 100; ++i) { - UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsQueued()); } UNIT_ASSERT(producer->Flush().GetValueSync().IsSuccess()); @@ -2093,14 +1997,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { writeSettings.Path(setup.GetTopicPath(TEST_TOPIC)); writeSettings.Codec(ECodec::RAW); writeSettings.ProducerIdPrefix("simple_blocking_producer_basic_write"); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.MaxMemoryUsage(100_KB); - writeSettings.MaxBlock(TDuration::MilliSeconds(1)); + writeSettings.MaxBlockTimeout(TDuration::MilliSeconds(1)); auto producer = client.CreateProducer(writeSettings); auto msgData = TString(1_MB, 'a'); - UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsSuccess()); + UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsQueued()); UNIT_ASSERT(producer->Write(TWriteMessage(msgData)).IsTimeout()); UNIT_ASSERT(producer->Close(TDuration::Seconds(10)).IsSuccess()); } diff --git a/tests/integration/topic/basic_usage_it.cpp b/tests/integration/topic/basic_usage_it.cpp index d8f63572e1b..6999e9b7da8 100644 --- a/tests/integration/topic/basic_usage_it.cpp +++ b/tests/integration/topic/basic_usage_it.cpp @@ -875,22 +875,23 @@ TEST_F(BasicUsage, TEST_NAME(TProducerBasicWrite_NoAutoPartitioning)) { .Path(GetTopicPath(TOPIC_NAME)) .Codec(ECodec::RAW); writeSettings.ProducerIdPrefix(CreateGuidAsString()); - writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::Hash); + writeSettings.PartitionChooserStrategy(TProducerSettings::EPartitionChooserStrategy::KafkaHash); writeSettings.SubSessionIdleTimeout(TDuration::Seconds(30)); writeSettings.PartitioningKeyHasher([](const std::string_view key) -> std::string { return std::string{key}; }); - writeSettings.MaxBlock(TDuration::Seconds(30)); + writeSettings.MaxBlockTimeout(TDuration::Seconds(30)); auto producer = client.CreateProducer(writeSettings); auto keyedSession = std::dynamic_pointer_cast(producer); + std::string data = "message"; for (size_t i = 0; i < 100; ++i) { auto key = CreateGuidAsString(); - TWriteMessage msg("msg"); + TWriteMessage msg(data); msg.SeqNo(i + 1); msg.Key(key); - ASSERT_TRUE(producer->Write(std::move(msg)).IsSuccess()); + ASSERT_TRUE(producer->Write(std::move(msg)).IsQueued()); } ASSERT_TRUE(producer->Close(TDuration::Seconds(10)).IsSuccess()); From cb2e5804b2dbbe36dee1179a51d068275141e662 Mon Sep 17 00:00:00 2001 From: Kuzin Roman Date: Mon, 23 Mar 2026 10:01:49 +0000 Subject: [PATCH 78/93] LOGBROKER-10206 Fix after support (#36161) --- .github/last_commit.txt | 2 +- .../producer/basic_write/main.cpp | 3 +- .../ydb-cpp-sdk/client/topic/write_session.h | 44 +++++++++++----- src/client/topic/impl/producer.cpp | 16 +++--- src/client/topic/ut/basic_usage_ut.cpp | 51 +++++++------------ tests/integration/topic/basic_usage_it.cpp | 3 +- 6 files changed, 59 insertions(+), 60 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 13a850091f5..6f341d95f9e 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -59c5b07fc52bfc4e047331be4b65156c9303227a +8c62bbd77946de4417ed6d1f9093f5e12c5d04f5 diff --git a/examples/topic_writer/producer/basic_write/main.cpp b/examples/topic_writer/producer/basic_write/main.cpp index 47629a00f06..b5f8e9ddfb1 100644 --- a/examples/topic_writer/producer/basic_write/main.cpp +++ b/examples/topic_writer/producer/basic_write/main.cpp @@ -115,8 +115,7 @@ int main() { auto messageData = std::string(1_KB, 'a'); for (int i = 0; i < 10; i++) { - NYdb::NTopic::TWriteMessage writeMessage(messageData); - writeMessage.Key("key" + ToString(i)); + NYdb::NTopic::TWriteMessage writeMessage("key" + ToString(i), messageData); WriteWithHandlingResult(producer, std::move(writeMessage)); } return 0; diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index 89752c6feb6..d5b6b16d70f 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -203,6 +203,16 @@ struct TWriteMessage { : Data(data) {} + TWriteMessage(const std::string& key, std::string_view data) + : Data(data) + , Key(key) + {} + + TWriteMessage(uint32_t partition, std::string_view data) + : Data(data) + , Partition(partition) + {} + TWriteMessage(const TWriteMessage& other) : DataHolder(other.DataHolder) , Data(other.DataHolder ? std::string_view(*DataHolder) : other.Data) @@ -211,9 +221,9 @@ struct TWriteMessage { , SeqNo_(other.SeqNo_) , CreateTimestamp_(other.CreateTimestamp_) , MessageMeta_(other.MessageMeta_) - , Key_(other.Key_) - , Partition_(other.Partition_) , Tx_(other.Tx_) + , Key(other.Key) + , Partition(other.Partition) {} TWriteMessage(TWriteMessage&& other) noexcept @@ -224,9 +234,9 @@ struct TWriteMessage { , SeqNo_(std::move(other.SeqNo_)) , CreateTimestamp_(std::move(other.CreateTimestamp_)) , MessageMeta_(std::move(other.MessageMeta_)) - , Key_(std::move(other.Key_)) - , Partition_(std::move(other.Partition_)) , Tx_(std::move(other.Tx_)) + , Key(std::move(other.Key)) + , Partition(std::move(other.Partition)) {} TWriteMessage& operator=(const TWriteMessage& other) { @@ -241,8 +251,8 @@ struct TWriteMessage { SeqNo_ = other.SeqNo_; CreateTimestamp_ = other.CreateTimestamp_; MessageMeta_ = other.MessageMeta_; - Key_ = other.Key_; - Partition_ = other.Partition_; + Key = other.Key; + Partition = other.Partition; Tx_ = other.Tx_; return *this; @@ -260,8 +270,8 @@ struct TWriteMessage { SeqNo_ = std::move(other.SeqNo_); CreateTimestamp_ = std::move(other.CreateTimestamp_); MessageMeta_ = std::move(other.MessageMeta_); - Key_ = std::move(other.Key_); - Partition_ = std::move(other.Partition_); + Key = std::move(other.Key); + Partition = std::move(other.Partition); Tx_ = std::move(other.Tx_); return *this; @@ -304,12 +314,6 @@ struct TWriteMessage { //! Message metadata. Limited to 4096 characters overall (all keys and values combined). FLUENT_SETTING(TMessageMeta, MessageMeta); - //! Message key. It will be used to route message to the partition. - FLUENT_SETTING_OPTIONAL(std::string, Key); - - //! Partition to write to. It is not recommended to use this option, use Key instead. - FLUENT_SETTING_OPTIONAL(std::uint32_t, Partition); - //! Transaction id FLUENT_SETTING_OPTIONAL(std::reference_wrapper, Tx); @@ -317,6 +321,18 @@ struct TWriteMessage { { return Tx_ ? &Tx_->get() : nullptr; } + + const std::optional& GetKey() const { + return Key; + } + + const std::optional& GetPartition() const { + return Partition; + } + +private: + std::optional Key; + std::optional Partition; }; //! Simple write session. Does not need event handlers. Does not provide Events, ContinuationTokens, write Acks. diff --git a/src/client/topic/impl/producer.cpp b/src/client/topic/impl/producer.cpp index 766f71abe4e..18b8e642213 100644 --- a/src/client/topic/impl/producer.cpp +++ b/src/client/topic/impl/producer.cpp @@ -1939,23 +1939,25 @@ TWriteResult TProducer::WriteInternal(TContinuationToken&&, TWriteMessage&& mess } std::uint32_t chosenPartition; - if (message.Partition_.has_value()) { - if (!Partitions[message.Partition_.value()].Children_.empty()) { + std::string key; + if (message.GetPartition().has_value()) { + if (!Partitions[message.GetPartition().value()].Children_.empty()) { return TWriteResult{ .Status = EWriteStatus::Error, .ErrorMessage = "Partition was split", }; } - chosenPartition = message.Partition_.value(); - } else if (!message.Key_.has_value()) { - message.Key(Settings.ProducerIdPrefix_); + chosenPartition = message.GetPartition().value(); + } else if (!message.GetKey().has_value()) { + key = Settings.ProducerIdPrefix_; chosenPartition = PartitionChooser->ChoosePartition(Settings.ProducerIdPrefix_); } else { - chosenPartition = PartitionChooser->ChoosePartition(*message.Key_); + chosenPartition = PartitionChooser->ChoosePartition(*message.GetKey()); + key = *message.GetKey(); } - MessagesWorker->AddMessage(message.Key_.value_or(""), std::move(message), chosenPartition); + MessagesWorker->AddMessage(key, std::move(message), chosenPartition); eventsPromise = EventsWorker->HandleNewMessage(); RunUserEventLoop(); } diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index aed9af75b63..d7d278b7fe2 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -53,9 +53,8 @@ TString SerializeDataChunk(ui64 seqNo, const TString& payload) { } TWriteMessage CreateMessage(std::string_view payload, const std::string& key, ui64 seqNo) { - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(seqNo); - msg.Key(key); return msg; } @@ -1090,9 +1089,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const ui64 messages = 5; for (ui64 i = 0; i < messages; ++i) { std::string payload = "payload"; - TWriteMessage msg(payload); + TWriteMessage msg("key-" + ToString(i), payload); msg.SeqNo(i + 1); - msg.Key("key-" + ToString(i)); UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1153,9 +1151,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto session = publicClient.CreateProducer(writeSettings); std::string payload = "msg0"; - TWriteMessage msg(payload); + TWriteMessage msg("key", payload); msg.SeqNo(0); - msg.Key("key"); UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); auto flushResult = session->Flush().GetValueSync(); UNIT_ASSERT_C(flushResult.IsClosed(), "Failed to flush producer"); @@ -1198,16 +1195,14 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto seqNo = 1; for (ui64 i = 0; i < count0; ++i) { std::string payload = "msg0"; - TWriteMessage msg(payload); + TWriteMessage msg(key0, payload); msg.SeqNo(seqNo++); - msg.Key(key0); UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } for (ui64 i = 0; i < count1; ++i) { std::string payload = "msg1"; - TWriteMessage msg(payload); + TWriteMessage msg(key1, payload); msg.SeqNo(seqNo++); - msg.Key(key1); UNIT_ASSERT_C(session->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1287,9 +1282,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { } std::string payload = "msg"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(i + 1); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1333,9 +1327,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { const ui64 count = 3000; for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(i); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1366,9 +1359,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(i); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1379,8 +1371,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { std::string payload = "data"; for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); - TWriteMessage msg(payload); - msg.Key(key); + TWriteMessage msg(key, payload); UNIT_ASSERT_C(producer2->Write(std::move(msg)).IsQueued(), "Failed to write message"); } UNIT_ASSERT_C(producer2->Flush().GetValueSync().IsSuccess(), "Failed to flush producer"); @@ -1410,9 +1401,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 1; i <= count; ++i) { auto key = CreateGuidAsString(); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(i); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1420,9 +1410,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { UNIT_ASSERT_C(producer->Close(TDuration::Seconds(10)).IsSuccess(), "Failed to close producer"); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(CreateGuidAsString(), payload); msg.SeqNo(count + 1); - msg.Key(CreateGuidAsString()); auto writeResult = producer->Write(std::move(msg)); UNIT_ASSERT_C(writeResult.IsError(), "Failed to write message"); UNIT_ASSERT_C(writeResult.ErrorMessage == "producer is closed", "Error message is not correct"); @@ -1459,9 +1448,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < perThread; ++i) { const ui64 seqNo = nextSeqNo.fetch_add(1); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(seqNo); - msg.Key(key); auto writeResult = producer->Write(std::move(msg)); UNIT_ASSERT_C( writeResult.IsQueued(), @@ -1502,9 +1490,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < messages; ++i) { auto key = CreateGuidAsString(); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(seqNo++); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1514,9 +1501,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 0; i < messages; ++i) { auto key = CreateGuidAsString(); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(seqNo++); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1550,9 +1536,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (ui64 i = 1; i <= messages; ++i) { auto key = CreateGuidAsString(); std::string payload = "data"; - TWriteMessage msg(payload); + TWriteMessage msg(key, payload); msg.SeqNo(i); - msg.Key(key); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } }); @@ -1589,9 +1574,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) { for (int i = 0; i < 1000; ++i) { std::string payload = "message-" + ToString(i); - TWriteMessage msg(payload); + TWriteMessage msg("key1", payload); msg.SeqNo(i + 1); - msg.Key("key1"); UNIT_ASSERT_C(producer->Write(std::move(msg)).IsQueued(), "Failed to write message"); } @@ -1927,8 +1911,7 @@ Y_UNIT_TEST_SUITE(BasicUsage) { auto msgData = TString(10_KB, 'a'); for (const auto& partition : partitions) { for (ui64 i = 0; i < 10; ++i) { - TWriteMessage msg(msgData); - msg.Partition(partition.GetPartitionId()); + TWriteMessage msg(partition.GetPartitionId(), msgData); UNIT_ASSERT(producer->Write(std::move(msg)).IsQueued()); } } diff --git a/tests/integration/topic/basic_usage_it.cpp b/tests/integration/topic/basic_usage_it.cpp index 6999e9b7da8..ceef48a74af 100644 --- a/tests/integration/topic/basic_usage_it.cpp +++ b/tests/integration/topic/basic_usage_it.cpp @@ -888,9 +888,8 @@ TEST_F(BasicUsage, TEST_NAME(TProducerBasicWrite_NoAutoPartitioning)) { std::string data = "message"; for (size_t i = 0; i < 100; ++i) { auto key = CreateGuidAsString(); - TWriteMessage msg(data); + TWriteMessage msg(key, data); msg.SeqNo(i + 1); - msg.Key(key); ASSERT_TRUE(producer->Write(std::move(msg)).IsQueued()); } From 8d4325df15a06836f3afb7878be703b46af57762 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Mon, 23 Mar 2026 10:01:56 +0000 Subject: [PATCH 79/93] YQ-5187 fixed hanging in PQ read session (#36220) --- .github/last_commit.txt | 2 +- src/client/topic/impl/read_session_impl.h | 5 ++ src/client/topic/impl/read_session_impl.ipp | 15 +++- src/client/topic/ut/basic_usage_ut.cpp | 72 ++++++++++++++----- .../topic/utils/managed_executor.cpp | 21 ++++++ .../topic/utils/managed_executor.h | 1 + 6 files changed, 95 insertions(+), 21 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 6f341d95f9e..7f62950f082 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -8c62bbd77946de4417ed6d1f9093f5e12c5d04f5 +771638ae94ff15b68c6796c9f14f16dc1596c4d4 diff --git a/src/client/topic/impl/read_session_impl.h b/src/client/topic/impl/read_session_impl.h index c60558caab0..99e96bf241f 100644 --- a/src/client/topic/impl/read_session_impl.h +++ b/src/client/topic/impl/read_session_impl.h @@ -161,6 +161,7 @@ class TDeferredActions { void DeferReconnection(TCallbackContextPtr cbContext, TPlainStatus&& status); void DeferStartSession(TCallbackContextPtr cbContext); void DeferSignalWaiter(TWaiter&& waiter); + void DeferOnUserRetrievedEvent(TUserRetrievedEventsInfoAccumulator&& accumulator); void DeferDestroyDecompressionInfos(std::vector>&& infos); private: @@ -175,6 +176,7 @@ class TDeferredActions { void Reconnect(); void SignalWaiters(); void StartSessions(); + void OnUserRetrievedEvent(); void DestroyDecompressionInfos(); private: @@ -219,6 +221,9 @@ class TDeferredActions { // Contexts for sessions to start std::vector> CbContexts; + // User retrieved events info accumulator. + std::vector> UserRetrievedEventsInfoAccumulator; + std::vector> DecompressionInfos; }; diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index 78c3b6bbee5..770dee20e1e 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -237,7 +237,7 @@ void TRawPartitionStreamEventQueue::DeleteNotReadyTail(TDe } deferred.DeferDestroyDecompressionInfos(std::move(infos)); - accumulator.OnUserRetrievedEvent(); + deferred.DeferOnUserRetrievedEvent(std::move(accumulator)); swap(ready, NotReady); } @@ -3375,6 +3375,11 @@ void TDeferredActions::DeferSignalWaiter(TWaiter&& waiter) Waiters.emplace_back(std::move(waiter)); } +template +void TDeferredActions::DeferOnUserRetrievedEvent(TUserRetrievedEventsInfoAccumulator&& accumulator) { + UserRetrievedEventsInfoAccumulator.push_back(std::move(accumulator)); +} + template void TDeferredActions::DeferDestroyDecompressionInfos(std::vector>&& infos) { @@ -3392,6 +3397,7 @@ void TDeferredActions::DoActions() { Reconnect(); SignalWaiters(); StartSessions(); + OnUserRetrievedEvent(); DestroyDecompressionInfos(); } @@ -3478,6 +3484,13 @@ void TDeferredActions::SignalWaiters() { } } +template +void TDeferredActions::OnUserRetrievedEvent() { + for (auto& accumulator : UserRetrievedEventsInfoAccumulator) { + accumulator.OnUserRetrievedEvent(); + } +} + template void TDeferredActions::DestroyDecompressionInfos() { for (const auto& info : DecompressionInfos) { diff --git a/src/client/topic/ut/basic_usage_ut.cpp b/src/client/topic/ut/basic_usage_ut.cpp index d7d278b7fe2..6d34cd7c1f2 100644 --- a/src/client/topic/ut/basic_usage_ut.cpp +++ b/src/client/topic/ut/basic_usage_ut.cpp @@ -246,7 +246,7 @@ void CreateTopicWithAutoPartitioning(TTopicClient& client) { } void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSessionSettings writeSettings, const std::string& message, std::uint32_t count, - TTopicSdkTestSetup& setup, std::shared_ptr decompressor, ui32 restartPeriod = 7, ui32 maxRestartsCount = 10) + TTopicSdkTestSetup& setup, std::shared_ptr decompressor, ui32 restartPeriod = 7, ui32 maxRestartsCount = 10, ui64 shuffleRatio = 1, TDuration shuffleDelay = TDuration::MilliSeconds(10)) { auto client = setup.MakeClient(); auto session = client.CreateSimpleBlockingWriteSession(writeSettings); @@ -262,9 +262,21 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess TTopicClient topicClient = setup.MakeClient(); - auto WaitTasks = [&, timeout = TInstant::Now() + TDuration::Seconds(60)](auto f, size_t c) { - while (f() < c) { - UNIT_ASSERT(timeout > TInstant::Now()); + auto WaitTasks = [&](auto f, size_t c) { + const auto hardTimeout = TInstant::Now() + TDuration::Seconds(60); + const auto shuffleTimeout = TInstant::Now() + shuffleDelay; + while (true) { + const auto fVal = f(); + if (fVal >= c * shuffleRatio) { + return; + } + + const auto now = TInstant::Now(); + if (fVal >= c && now > shuffleTimeout) { + return; + } + + UNIT_ASSERT_GE(hardTimeout, now); ReadSession->WaitEvent(); std::this_thread::sleep_for(100ms); }; @@ -276,28 +288,23 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess WaitTasks([&]() { return e->GetExecutedCount(); }, count); }; - auto RunTasks = [&](auto e, const std::vector& tasks) { - size_t n = tasks.size(); - WaitPlannedTasks(e, n); + auto RunTask = [&](auto e) { + WaitPlannedTasks(e, 1); size_t completed = e->GetExecutedCount(); - e->StartFuncs(tasks); - WaitExecutedTasks(e, completed + n); + e->StartRandomFunc(); + WaitExecutedTasks(e, completed + 1); }; - Y_UNUSED(RunTasks); - auto PlanTasksAndRestart = [&](auto e, const std::vector& tasks) { - size_t n = tasks.size(); - WaitPlannedTasks(e, n); + auto PlanTaskAndRestart = [&](auto e) { + WaitPlannedTasks(e, 1); size_t completed = e->GetExecutedCount(); setup.GetServer().KillTopicPqrbTablet(JoinPath({TString(setup.MakeDriverConfig().GetDatabase()), TString(setup.GetTopicPath())})); std::this_thread::sleep_for(100ms); - e->StartFuncs(tasks); - WaitExecutedTasks(e, completed + n); + e->StartRandomFunc(); + WaitExecutedTasks(e, completed + 1); }; - Y_UNUSED(PlanTasksAndRestart); - NThreading::TPromise checkedPromise = NThreading::NewPromise(); TAtomic lastOffset = 0u; @@ -318,11 +325,12 @@ void WriteAndReadToEndWithRestarts(TReadSessionSettings readSettings, TWriteSess ui32 restartCount = 0; while (AtomicGet(lastOffset) + 1 < count) { if (restartCount < maxRestartsCount && i % restartPeriod == 1) { - PlanTasksAndRestart(decompressor, {i++}); + PlanTaskAndRestart(decompressor); restartCount++; } else { - RunTasks(decompressor, {i++}); + RunTask(decompressor); } + i++; } ReadSession->Close(TDuration::MilliSeconds(10)); @@ -1013,6 +1021,32 @@ Y_UNIT_TEST_SUITE(BasicUsage) { WriteAndReadToEndWithRestarts(readSettings, writeSettings, message, count, setup, decompressor); } + Y_UNIT_TEST(ReadWithRestartsAndLargeDataAndShuffle) { + TTopicSdkTestSetup setup(TEST_CASE_NAME); + auto compressor = std::make_shared(); + auto decompressor = CreateThreadPoolManagedExecutor(1); + + TReadSessionSettings readSettings; + readSettings + .ConsumerName(setup.GetConsumerName()) + .MaxMemoryUsageBytes(10_MB) + .DecompressionExecutor(decompressor) + .AppendTopics(setup.GetTopicPath()) + // .DirectRead(EnableDirectRead) + ; + + TWriteSessionSettings writeSettings; + writeSettings + .Path(setup.GetTopicPath()).MessageGroupId(TEST_MESSAGE_GROUP_ID) + .Codec(ECodec::RAW) + .CompressionExecutor(compressor); + + std::uint32_t count = 3000; + std::string message(8'000, 'x'); + + WriteAndReadToEndWithRestarts(readSettings, writeSettings, message, count, setup, decompressor, 7, 10, 10); + } + Y_UNIT_TEST(ConflictingWrites) { TTopicSdkTestSetup setup(TEST_CASE_NAME); diff --git a/tests/integration/topic/utils/managed_executor.cpp b/tests/integration/topic/utils/managed_executor.cpp index c96bc1d0716..fa519539327 100644 --- a/tests/integration/topic/utils/managed_executor.cpp +++ b/tests/integration/topic/utils/managed_executor.cpp @@ -1,5 +1,6 @@ #include "managed_executor.h" +#include namespace NYdb::inline V3::NTopic::NTests { @@ -50,6 +51,26 @@ void TManagedExecutor::RunTask(TFunction&& func) Executor->Post(MakeTask(std::move(func))); } +void TManagedExecutor::StartRandomFunc() { + std::lock_guard lock(Mutex); + + Y_ABORT_UNLESS(Planned > 0); + size_t index = RandomNumber(Planned); + + for (size_t i = 0; i < Funcs.size(); ++i) { + if (Funcs[i] != nullptr) { + if (index == 0) { + RunTask(std::move(Funcs[i])); + Funcs[i] = nullptr; + return; + } + --index; + } + } + + Y_ABORT("No functions to start"); +} + void TManagedExecutor::StartFuncs(const std::vector& indicies) { std::lock_guard lock(Mutex); diff --git a/tests/integration/topic/utils/managed_executor.h b/tests/integration/topic/utils/managed_executor.h index d379f469a3d..3076ae4044e 100644 --- a/tests/integration/topic/utils/managed_executor.h +++ b/tests/integration/topic/utils/managed_executor.h @@ -19,6 +19,7 @@ class TManagedExecutor : public IExecutor { void Stop() override; + void StartRandomFunc(); void StartFuncs(const std::vector& indicies); size_t GetFuncsCount() const; From a3cb3d7db7bbed1e54a34075fe0c1da0ae24622b Mon Sep 17 00:00:00 2001 From: Anely Date: Mon, 23 Mar 2026 10:02:03 +0000 Subject: [PATCH 80/93] Add session stop reason hint (#35434) --- .github/last_commit.txt | 2 +- src/api/protos/ydb_query.proto | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 7f62950f082..7b3760d4b24 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -771638ae94ff15b68c6796c9f14f16dc1596c4d4 +688edef85e4e4e55828630ef3943a91ca9306799 diff --git a/src/api/protos/ydb_query.proto b/src/api/protos/ydb_query.proto index afc60032e6e..c795eeaafa4 100644 --- a/src/api/protos/ydb_query.proto +++ b/src/api/protos/ydb_query.proto @@ -45,9 +45,21 @@ message AttachSessionRequest { string session_id = 1 [(Ydb.length).le = 1024]; } +message SessionShutdownHint { +} + +message NodeShutdownHint { +} + message SessionState { StatusIds.StatusCode status = 1; repeated Ydb.Issue.IssueMessage issues = 2; + + // The reason the session is ending, for SDK-side handling + oneof session_hint { + SessionShutdownHint session_shutdown = 3; + NodeShutdownHint node_shutdown = 4; + } } message SerializableModeSettings { From 624141442899aa36f30e43a91be981c714385957 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 23 Mar 2026 10:02:04 +0000 Subject: [PATCH 81/93] Update import generation: 37 --- .github/import_generation.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/import_generation.txt b/.github/import_generation.txt index 81b5c5d06cc..e522732c77e 100644 --- a/.github/import_generation.txt +++ b/.github/import_generation.txt @@ -1 +1 @@ -37 +38 From 6c625f2c87c22e7e1025d7fd404cb17811fb17a1 Mon Sep 17 00:00:00 2001 From: maladetska Date: Sun, 1 Mar 2026 02:20:51 +0300 Subject: [PATCH 82/93] fixes --- include/ydb-cpp-sdk/client/driver/driver.h | 7 +- include/ydb-cpp-sdk/client/metrics/metrics.h | 40 ++++- plugins/CMakeLists.txt | 6 + plugins/open_telemetry/CMakeLists.txt | 36 +++++ .../ydb-cpp-sdk/open_telemetry/extension.h | 16 ++ .../ydb-cpp-sdk/open_telemetry/metrics.h | 36 +++++ .../ydb-cpp-sdk/open_telemetry/trace.h | 29 ++++ plugins/open_telemetry/src/metrics.cpp | 140 ++++++++++++++++++ plugins/open_telemetry/src/trace.cpp | 70 +++++++++ src/client/driver/driver.cpp | 36 +++++ .../grpc_connections/grpc_connections.cpp | 14 ++ .../grpc_connections/grpc_connections.h | 21 +++ .../impl/internal/grpc_connections/params.h | 5 + src/client/metrics/CMakeLists.txt | 4 + src/client/query/client.cpp | 15 ++ src/client/query/impl/query_spans.cpp | 27 ++++ src/client/query/impl/query_spans.h | 3 + 17 files changed, 496 insertions(+), 9 deletions(-) create mode 100644 plugins/open_telemetry/CMakeLists.txt create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h create mode 100644 plugins/open_telemetry/src/metrics.cpp create mode 100644 plugins/open_telemetry/src/trace.cpp diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index 199459c58f7..d6880f73c0b 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -168,12 +169,6 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); - //! Set external metrics registry implementation. - TDriverConfig& SetMetricRegistry(std::shared_ptr registry); - - //! Set external trace provider implementation. - TDriverConfig& SetTraceProvider(std::shared_ptr provider); - private: class TImpl; std::shared_ptr Impl_; diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index 5faa930ed50..309ec9a9517 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -1,10 +1,10 @@ #pragma once +<<<<<<< HEAD #include +======= +>>>>>>> 1b2bf4fa5 (fixes) #include -#include -#include -#include namespace NYdb::inline V3::NMetrics { @@ -33,6 +33,7 @@ class IMetricRegistry { public: virtual ~IMetricRegistry() = default; +<<<<<<< HEAD virtual std::shared_ptr Counter( const std::string& name, const TLabels& labels = {}, @@ -52,6 +53,39 @@ class IMetricRegistry { const std::string& description = {}, const std::string& unit = {} ) = 0; +======= + virtual std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) = 0; + virtual std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) = 0; + virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; +}; + +enum class ESpanKind { + INTERNAL, + SERVER, + CLIENT, + PRODUCER, + CONSUMER +}; + +class ISpan { +public: + virtual ~ISpan() = default; + virtual void End() = 0; + virtual void SetAttribute(const std::string& key, const std::string& value) = 0; + virtual void SetAttribute(const std::string& key, int64_t value) = 0; +}; + +class ITracer { +public: + virtual ~ITracer() = default; + virtual std::shared_ptr StartSpan(const std::string& name, ESpanKind kind = ESpanKind::INTERNAL) = 0; +}; + +class ITraceProvider { +public: + virtual ~ITraceProvider() = default; + virtual std::shared_ptr GetTracer(const std::string& name) = 0; +>>>>>>> 1b2bf4fa5 (fixes) }; } // namespace NYdb::NMetrics diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 0d232800455..d28a25bf376 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -1,2 +1,8 @@ +<<<<<<< HEAD add_subdirectory(metrics) add_subdirectory(trace) +======= +if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) + add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) +endif() +>>>>>>> 1b2bf4fa5 (fixes) diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt new file mode 100644 index 00000000000..d005708d422 --- /dev/null +++ b/plugins/open_telemetry/CMakeLists.txt @@ -0,0 +1,36 @@ +if (YDB_SDK_ENABLE_OTEL_METRICS) + _ydb_sdk_add_library(open_telemetry_metrics) + target_sources(open_telemetry_metrics PRIVATE + src/metrics.cpp + ) + target_include_directories(open_telemetry_metrics PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_metrics PUBLIC + client-metrics + client-resources + opentelemetry-cpp::api + opentelemetry-cpp::metrics + ) + _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) +endif() + +if (YDB_SDK_ENABLE_OTEL_TRACE) + _ydb_sdk_add_library(open_telemetry_trace) + target_sources(open_telemetry_trace PRIVATE + src/trace.cpp + ) + target_include_directories(open_telemetry_trace PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_trace PUBLIC + client-metrics + opentelemetry-cpp::api + opentelemetry-cpp::trace + ) + _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) +endif() + +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h new file mode 100644 index 00000000000..b0db9ea7d7c --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace NYdb::inline V3::NMetrics { + +inline void AddOpenTelemetry(TDriverConfig& config + , opentelemetry::nostd::shared_ptr meterProvider + , opentelemetry::nostd::shared_ptr tracerProvider +) { + AddOpenTelemetryMetrics(config, std::move(meterProvider)); + AddOpenTelemetryTrace(config, std::move(tracerProvider)); +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h new file mode 100644 index 00000000000..5e9e9e77dea --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelMetricRegistry : public IMetricRegistry { +public: + TOtelMetricRegistry(opentelemetry::nostd::shared_ptr meterProvider); + + std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) override; + std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) override; + std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) override; + +private: + void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets); + + opentelemetry::nostd::shared_ptr MeterProvider_; + opentelemetry::nostd::shared_ptr Meter_; + std::mutex HistogramViewsLock_; + std::unordered_set HistogramViews_; +}; + +inline void AddOpenTelemetryMetrics( + TDriverConfig& config, + opentelemetry::nostd::shared_ptr meterProvider) +{ + if (meterProvider) { + config.SetMetricExporter(std::make_shared(std::move(meterProvider))); + } +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h new file mode 100644 index 00000000000..3ba2e146fd9 --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -0,0 +1,29 @@ +#pragma once + +#include +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelTraceProvider : public ITraceProvider { +public: + TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); + + std::shared_ptr GetTracer(const std::string& name) override; + +private: + opentelemetry::nostd::shared_ptr TracerProvider_; +}; + +inline void AddOpenTelemetryTrace( + TDriverConfig& config, + opentelemetry::nostd::shared_ptr tracerProvider) +{ + if (tracerProvider) { + config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); + } +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/metrics.cpp b/plugins/open_telemetry/src/metrics.cpp new file mode 100644 index 00000000000..65850fd08b3 --- /dev/null +++ b/plugins/open_telemetry/src/metrics.cpp @@ -0,0 +1,140 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +common::KeyValueIterableView MakeAttributes(const TLabels& labels) { + return common::KeyValueIterableView(labels); +} + +class TOtelCounter : public ICounter { +public: + TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Inc() override { + Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; +}; + +class TOtelUpDownCounterGauge : public IGauge { +public: + TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) + : Counter_(std::move(counter)) + , Labels_(labels) + {} + + void Add(double delta) override { + Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ += delta; + } + + void Set(double value) override { + Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + Value_ = value; + } + +private: + nostd::shared_ptr> Counter_; + TLabels Labels_; + double Value_ = 0; +}; + +class TOtelHistogram : public IHistogram { +public: + TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) + : Histogram_(std::move(histogram)) + , Labels_(labels) + {} + + void Record(double value) override { + Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Histogram_; + TLabels Labels_; +}; + +} // namespace + +TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) + : MeterProvider_(std::move(meterProvider)) + , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) +{} + +void TOtelMetricRegistry::ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { + if (buckets.empty()) { + return; + } + + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); + if (!sdkProvider) { + return; + } + + { + std::lock_guard lock(HistogramViewsLock_); + if (!HistogramViews_.insert(name).second) { + return; + } + } + + auto selector = std::make_unique( + sdk::metrics::InstrumentType::kHistogram, + name, + "" + ); + auto meterSelector = std::make_unique( + "ydb-cpp-sdk", + GetSdkSemver(), + {} + ); + + auto histogramConfig = std::make_shared(); + histogramConfig->boundaries_ = buckets; + + auto view = std::make_unique( + {}, + {}, + sdk::metrics::AggregationType::kHistogram, + histogramConfig + ); + + sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); +} + +std::shared_ptr TOtelMetricRegistry::Counter(const std::string& name, const TLabels& labels) { + auto counter = Meter_->CreateUInt64Counter(name); + return std::make_shared(std::move(counter), labels); +} + +std::shared_ptr TOtelMetricRegistry::Gauge(const std::string& name, const TLabels& labels) { + auto counter = Meter_->CreateDoubleUpDownCounter(name); + return std::make_shared(std::move(counter), labels); +} + +std::shared_ptr TOtelMetricRegistry::Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) { + ConfigureHistogramBuckets(name, buckets); + auto histogram = Meter_->CreateDoubleHistogram(name); + return std::make_shared(std::move(histogram), labels); +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/open_telemetry/src/trace.cpp new file mode 100644 index 00000000000..54f04cb84df --- /dev/null +++ b/plugins/open_telemetry/src/trace.cpp @@ -0,0 +1,70 @@ +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +trace::SpanKind MapSpanKind(ESpanKind kind) { + switch (kind) { + case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; + case ESpanKind::SERVER: return trace::SpanKind::kServer; + case ESpanKind::CLIENT: return trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + } + return trace::SpanKind::kInternal; +} + +class TOtelSpan : public ISpan { +public: + TOtelSpan(nostd::shared_ptr span) + : Span_(std::move(span)) + {} + + void End() override { + Span_->End(); + } + + void SetAttribute(const std::string& key, const std::string& value) override { + Span_->SetAttribute(key, value); + } + + void SetAttribute(const std::string& key, int64_t value) override { + Span_->SetAttribute(key, value); + } + +private: + nostd::shared_ptr Span_; +}; + +class TOtelTracer : public ITracer { +public: + TOtelTracer(nostd::shared_ptr tracer) + : Tracer_(std::move(tracer)) + {} + + std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { + trace::StartSpanOptions options; + options.kind = MapSpanKind(kind); + return std::make_shared(Tracer_->StartSpan(name, options)); + } + +private: + nostd::shared_ptr Tracer_; +}; + +} // namespace + +TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) + : TracerProvider_(std::move(tracerProvider)) +{} + +std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { + return std::make_shared(TracerProvider_->GetTracer(name)); +} + +} // namespace NYdb::NMetrics diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 9bfc5f80560..ff831986ddb 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -53,8 +53,13 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t GetMaxMessageSize() const override { return MaxMessageSize; } const TLog& GetLog() const override { return Log; } std::shared_ptr GetExecutor() const override { return Executor; } +<<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } std::shared_ptr GetTraceProvider() const override { return TraceProvider; } +======= + std::shared_ptr GetMetricExporter() const override { return MetricExporter; } + std::shared_ptr GetTraceExporter() const override { return TraceExporter; } +>>>>>>> 1b2bf4fa5 (fixes) std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -86,8 +91,13 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t MaxMessageSize = 0; TLog Log; // Null by default. std::shared_ptr Executor; +<<<<<<< HEAD std::shared_ptr MetricRegistry; std::shared_ptr TraceProvider; +======= + std::shared_ptr MetricExporter; + std::shared_ptr TraceExporter; +>>>>>>> 1b2bf4fa5 (fixes) }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -247,6 +257,7 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { return *this; } +<<<<<<< HEAD TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { Impl_->MetricRegistry = std::move(registry); return *this; @@ -257,6 +268,26 @@ TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr exporter) { + Impl_->MetricExporter = std::move(exporter); + return *this; +} + +TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { + Impl_->TraceExporter = std::move(exporter); + return *this; +} + +std::shared_ptr TDriverConfig::GetMetricExporter() const { + return Impl_->MetricExporter; +} + +std::shared_ptr TDriverConfig::GetTraceExporter() const { + return Impl_->TraceExporter; +} + +>>>>>>> 1b2bf4fa5 (fixes) //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -310,8 +341,13 @@ TDriverConfig TDriver::GetConfig() const { config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); config.SetMaxMessageSize(Impl_->MaxMessageSize_); config.Impl_->Log = Impl_->Log; +<<<<<<< HEAD config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); config.SetTraceProvider(Impl_->GetTraceProvider()); +======= + config.SetMetricExporter(Impl_->GetMetricExporter()); + config.SetTraceExporter(Impl_->GetTraceExporter()); +>>>>>>> 1b2bf4fa5 (fixes) return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index 625358f0023..c1690a9af17 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -169,8 +169,13 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout(), TcpNoDelay_) #endif +<<<<<<< HEAD , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) +======= + , MetricExporter_(params->GetMetricExporter()) + , TraceExporter_(params->GetTraceExporter()) +>>>>>>> 1b2bf4fa5 (fixes) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -438,12 +443,21 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { ExtensionApis_.emplace_back(api); } +<<<<<<< HEAD std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { return MetricRegistry_; } std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { return TraceProvider_; +======= +std::shared_ptr TGRpcConnectionsImpl::GetMetricExporter() const { + return MetricExporter_; +} + +std::shared_ptr TGRpcConnectionsImpl::GetTraceExporter() const { + return TraceExporter_; +>>>>>>> 1b2bf4fa5 (fixes) } void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 99555ffc24c..6c3d7533e58 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -586,8 +586,24 @@ class TGRpcConnectionsImpl ::NMonitoring::TMetricRegistry* GetMetricRegistry() override; void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); +<<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const; std::shared_ptr GetTraceProvider() const; +======= + std::shared_ptr GetMetricExporter() const; + std::shared_ptr GetTraceExporter() const; + + template + T* GetExtensionApi() { + std::lock_guard lock(ExtensionsLock_); + for (const auto& api : ExtensionApis_) { + if (auto ptr = dynamic_cast(api.get())) { + return ptr; + } + } + return nullptr; + } +>>>>>>> 1b2bf4fa5 (fixes) void SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb); const TLog& GetLog() const override; @@ -724,8 +740,13 @@ class TGRpcConnectionsImpl std::vector> Extensions_; std::vector> ExtensionApis_; +<<<<<<< HEAD std::shared_ptr MetricRegistry_; std::shared_ptr TraceProvider_; +======= + std::shared_ptr MetricExporter_; + std::shared_ptr TraceExporter_; +>>>>>>> 1b2bf4fa5 (fixes) IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 7f23bcceffb..61a2e0f5a2e 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -43,8 +43,13 @@ class IConnectionsParams { virtual uint64_t GetMaxOutboundMessageSize() const = 0; virtual uint64_t GetMaxMessageSize() const = 0; virtual std::shared_ptr GetExecutor() const = 0; +<<<<<<< HEAD virtual std::shared_ptr GetExternalMetricRegistry() const = 0; virtual std::shared_ptr GetTraceProvider() const = 0; +======= + virtual std::shared_ptr GetMetricExporter() const = 0; + virtual std::shared_ptr GetTraceExporter() const = 0; +>>>>>>> 1b2bf4fa5 (fixes) }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index e681a846b26..7d376c5cac3 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -1,7 +1,11 @@ _ydb_sdk_add_library(client-metrics) target_sources(client-metrics PRIVATE +<<<<<<< HEAD metrics.cpp +======= + metrics.cpp +>>>>>>> 1b2bf4fa5 (fixes) ) _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 8aefee460e2..a5c3ec09242 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -71,7 +71,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); +<<<<<<< HEAD if (auto traceProvider = Connections_->GetTraceProvider()) { +======= + if (auto traceProvider = Connections_->GetTraceExporter()) { +>>>>>>> 1b2bf4fa5 (fixes) Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } MetricRegistry_ = Connections_->GetExternalMetricRegistry(); @@ -490,9 +494,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(status.GetStatus()); } +<<<<<<< HEAD if (Metrics) { Metrics->End(status.GetStatus()); } +======= +>>>>>>> 1b2bf4fa5 (fixes) ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -508,9 +515,12 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(EStatus::SUCCESS); } +<<<<<<< HEAD if (Metrics) { Metrics->End(EStatus::SUCCESS); } +======= +>>>>>>> 1b2bf4fa5 (fixes) ScheduleReply(std::move(val)); } @@ -520,12 +530,17 @@ class TQueryClient::TImpl: public TClientImplCommon, public { auto val = future.ExtractValue(); if (span) { +<<<<<<< HEAD span->SetPeerEndpoint(val.GetEndpoint()); span->End(val.GetStatus()); } if (metrics) { metrics->End(val.GetStatus()); } +======= + span->End(val.GetStatus()); + } +>>>>>>> 1b2bf4fa5 (fixes) promise.SetValue(std::move(val)); }); } diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index 72d3e2eec0c..b18881559e7 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -53,6 +53,18 @@ void SafeLogSpanError(const char* message) noexcept { } } +void SafeLogSpanError(const char* message) noexcept { + try { + try { + Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + return; + } catch (...) { + } + Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + } catch (...) { + } +} + } // namespace TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { @@ -65,12 +77,20 @@ TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::str ParseEndpoint(endpoint, host, port); try { +<<<<<<< HEAD Span_ = tracer->StartSpan(operationName, NMetrics::ESpanKind::CLIENT); if (!Span_) { return; } Span_->SetAttribute("db.system.name", "other_sql"); Span_->SetAttribute("db.operation.name", operationName); +======= + Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); + if (!Span_) { + return; + } + Span_->SetAttribute("db.system.name", "ydb"); +>>>>>>> 1b2bf4fa5 (fixes) Span_->SetAttribute("server.address", host); Span_->SetAttribute("server.port", static_cast(port)); } catch (...) { @@ -89,6 +109,7 @@ TQuerySpan::~TQuerySpan() noexcept { } } +<<<<<<< HEAD void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { if (!Span_ || endpoint.empty()) { return; @@ -131,6 +152,12 @@ void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { Span_->SetAttribute("db.response.status_code", ToString(status)); +======= +void TQuerySpan::End(EStatus status) noexcept { + if (Span_) { + try { + Span_->SetAttribute("db.response.status_code", static_cast(status)); +>>>>>>> 1b2bf4fa5 (fixes) if (status != EStatus::SUCCESS) { Span_->SetAttribute("error.type", ToString(status)); } diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index 2aeb5d5b79b..168f5592843 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -15,10 +15,13 @@ class TQuerySpan { TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); ~TQuerySpan() noexcept; +<<<<<<< HEAD void SetPeerEndpoint(const std::string& endpoint) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; std::unique_ptr Activate() noexcept; +======= +>>>>>>> 1b2bf4fa5 (fixes) void End(EStatus status) noexcept; private: From e194ba88e26c92264a8b9640ee7eb0001de8a813 Mon Sep 17 00:00:00 2001 From: maladetska Date: Mon, 16 Mar 2026 01:55:33 +0800 Subject: [PATCH 83/93] fixes and add metric tests --- include/ydb-cpp-sdk/client/driver/driver.h | 10 +- include/ydb-cpp-sdk/client/metrics/metrics.h | 3 + include/ydb-cpp-sdk/client/trace/trace.h | 6 + plugins/CMakeLists.txt | 5 + .../ydb-cpp-sdk/open_telemetry/metrics.h | 7 + plugins/metrics/otel/src/metrics.cpp | 104 ++++++++++ plugins/open_telemetry/CMakeLists.txt | 36 ---- .../ydb-cpp-sdk/open_telemetry/extension.h | 16 -- .../ydb-cpp-sdk/open_telemetry/metrics.h | 36 ---- .../ydb-cpp-sdk/open_telemetry/trace.h | 29 --- plugins/open_telemetry/src/metrics.cpp | 140 ------------- plugins/open_telemetry/src/trace.cpp | 70 ------- .../ydb-cpp-sdk/open_telemetry/trace.h | 7 + plugins/trace/otel/src/trace.cpp | 54 +++++ src/client/driver/driver.cpp | 27 ++- .../grpc_connections/grpc_connections.cpp | 14 ++ .../grpc_connections/grpc_connections.h | 10 + .../impl/internal/grpc_connections/params.h | 5 + src/client/metrics/CMakeLists.txt | 4 + src/client/query/client.cpp | 26 +++ src/client/query/impl/CMakeLists.txt | 1 + src/client/query/impl/query_metrics.cpp | 71 +++++++ src/client/query/impl/query_metrics.h | 26 +++ src/client/query/impl/query_spans.cpp | 27 ++- src/client/query/impl/query_spans.h | 7 + tests/common/fake_metric_registry.h | 12 ++ tests/integration/metrics/main.cpp | 106 ++++++++++ tests/unit/client/CMakeLists.txt | 14 ++ tests/unit/client/query/query_metrics_ut.cpp | 190 ++++++++++++++++++ 29 files changed, 731 insertions(+), 332 deletions(-) delete mode 100644 plugins/open_telemetry/CMakeLists.txt delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h delete mode 100644 plugins/open_telemetry/src/metrics.cpp delete mode 100644 plugins/open_telemetry/src/trace.cpp create mode 100644 src/client/query/impl/query_metrics.cpp create mode 100644 tests/unit/client/query/query_metrics_ut.cpp diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index d6880f73c0b..4471fda9277 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -8,7 +8,6 @@ #include #include #include -#include #include #include #include @@ -169,6 +168,15 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); +<<<<<<< HEAD +======= + //! Set external metrics registry implementation. + TDriverConfig& SetMetricRegistry(std::shared_ptr registry); + + //! Set external trace provider implementation. + TDriverConfig& SetTraceProvider(std::shared_ptr provider); + +>>>>>>> 1ca4253b5 (fixes and add metric tests) private: class TImpl; std::shared_ptr Impl_; diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index 309ec9a9517..1430027fc55 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -59,6 +59,7 @@ class IMetricRegistry { virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; }; +<<<<<<< HEAD enum class ESpanKind { INTERNAL, SERVER, @@ -88,4 +89,6 @@ class ITraceProvider { >>>>>>> 1b2bf4fa5 (fixes) }; +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) } // namespace NYdb::NMetrics diff --git a/include/ydb-cpp-sdk/client/trace/trace.h b/include/ydb-cpp-sdk/client/trace/trace.h index 054fa258488..5644720cd92 100644 --- a/include/ydb-cpp-sdk/client/trace/trace.h +++ b/include/ydb-cpp-sdk/client/trace/trace.h @@ -15,11 +15,14 @@ enum class ESpanKind { CONSUMER }; +<<<<<<< HEAD class IScope { public: virtual ~IScope() = default; }; +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) class ISpan { public: virtual ~ISpan() = default; @@ -27,7 +30,10 @@ class ISpan { virtual void SetAttribute(const std::string& key, const std::string& value) = 0; virtual void SetAttribute(const std::string& key, int64_t value) = 0; virtual void AddEvent(const std::string& name, const std::map& attributes = {}) = 0; +<<<<<<< HEAD virtual std::unique_ptr Activate() = 0; +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; class ITracer { diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index d28a25bf376..6f38e537e6b 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -1,4 +1,5 @@ <<<<<<< HEAD +<<<<<<< HEAD add_subdirectory(metrics) add_subdirectory(trace) ======= @@ -6,3 +7,7 @@ if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) endif() >>>>>>> 1b2bf4fa5 (fixes) +======= +add_subdirectory(metrics) +add_subdirectory(trace) +>>>>>>> 1ca4253b5 (fixes and add metric tests) diff --git a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h index f992c577bf6..698e2efeb92 100644 --- a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h +++ b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -3,7 +3,14 @@ #include #include +<<<<<<< HEAD #include +======= + +namespace opentelemetry::metrics { +class MeterProvider; +} +>>>>>>> 1ca4253b5 (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { diff --git a/plugins/metrics/otel/src/metrics.cpp b/plugins/metrics/otel/src/metrics.cpp index f883fae7946..b78b4efacb7 100644 --- a/plugins/metrics/otel/src/metrics.cpp +++ b/plugins/metrics/otel/src/metrics.cpp @@ -9,12 +9,16 @@ #include #include +<<<<<<< HEAD #include +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { namespace { +<<<<<<< HEAD namespace otel_metrics = opentelemetry::metrics; namespace otel_nostd = opentelemetry::nostd; namespace otel_common = opentelemetry::common; @@ -23,70 +27,121 @@ namespace otel_sdk_metrics = opentelemetry::sdk::metrics; otel_common::KeyValueIterableView MakeAttributes(const TLabels& labels) { return otel_common::KeyValueIterableView(labels); +======= +using namespace opentelemetry; + +common::KeyValueIterableView MakeAttributes(const TLabels& labels) { + return common::KeyValueIterableView(labels); +>>>>>>> 1ca4253b5 (fixes and add metric tests) } class TOtelCounter : public ICounter { public: +<<<<<<< HEAD TOtelCounter(otel_nostd::shared_ptr> counter, const TLabels& labels) +======= + TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : Counter_(std::move(counter)) , Labels_(labels) {} void Inc() override { +<<<<<<< HEAD Counter_->Add(1, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } private: otel_nostd::shared_ptr> Counter_; +======= + Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Counter_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) TLabels Labels_; }; class TOtelUpDownCounterGauge : public IGauge { public: +<<<<<<< HEAD TOtelUpDownCounterGauge(otel_nostd::shared_ptr> counter, const TLabels& labels) +======= + TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : Counter_(std::move(counter)) , Labels_(labels) {} void Add(double delta) override { +<<<<<<< HEAD Counter_->Add(delta, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); +======= + Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); +>>>>>>> 1ca4253b5 (fixes and add metric tests) Value_ += delta; } void Set(double value) override { +<<<<<<< HEAD Counter_->Add(value - Value_, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); +======= + Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); +>>>>>>> 1ca4253b5 (fixes and add metric tests) Value_ = value; } private: +<<<<<<< HEAD otel_nostd::shared_ptr> Counter_; +======= + nostd::shared_ptr> Counter_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) TLabels Labels_; double Value_ = 0; }; class TOtelHistogram : public IHistogram { public: +<<<<<<< HEAD TOtelHistogram(otel_nostd::shared_ptr> histogram, const TLabels& labels) +======= + TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : Histogram_(std::move(histogram)) , Labels_(labels) {} void Record(double value) override { +<<<<<<< HEAD Histogram_->Record(value, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } private: otel_nostd::shared_ptr> Histogram_; +======= + Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); + } + +private: + nostd::shared_ptr> Histogram_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) TLabels Labels_; }; class TOtelMetricRegistry : public IMetricRegistry { public: +<<<<<<< HEAD TOtelMetricRegistry(otel_nostd::shared_ptr meterProvider) +======= + TOtelMetricRegistry(nostd::shared_ptr meterProvider) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : MeterProvider_(std::move(meterProvider)) , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) {} +<<<<<<< HEAD std::shared_ptr Counter(const std::string& name , const TLabels& labels , const std::string& description @@ -113,16 +168,39 @@ class TOtelMetricRegistry : public IMetricRegistry { ) override { ConfigureHistogramBuckets(name, unit, buckets); auto histogram = Meter_->CreateDoubleHistogram(name, description, unit); +======= + std::shared_ptr Counter(const std::string& name, const TLabels& labels) override { + auto counter = Meter_->CreateUInt64Counter(name); + return std::make_shared(std::move(counter), labels); + } + + std::shared_ptr Gauge(const std::string& name, const TLabels& labels) override { + auto counter = Meter_->CreateDoubleUpDownCounter(name); + return std::make_shared(std::move(counter), labels); + } + + std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) override { + ConfigureHistogramBuckets(name, buckets); + auto histogram = Meter_->CreateDoubleHistogram(name); +>>>>>>> 1ca4253b5 (fixes and add metric tests) return std::make_shared(std::move(histogram), labels); } private: +<<<<<<< HEAD void ConfigureHistogramBuckets(const std::string& name, const std::string& unit, const std::vector& buckets) { +======= + void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { +>>>>>>> 1ca4253b5 (fixes and add metric tests) if (buckets.empty()) { return; } +<<<<<<< HEAD auto* sdkProvider = dynamic_cast(MeterProvider_.get()); +======= + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); +>>>>>>> 1ca4253b5 (fixes and add metric tests) if (!sdkProvider) { return; } @@ -134,6 +212,7 @@ class TOtelMetricRegistry : public IMetricRegistry { } } +<<<<<<< HEAD auto selector = std::make_unique( otel_sdk_metrics::InstrumentType::kHistogram, name, @@ -152,14 +231,39 @@ class TOtelMetricRegistry : public IMetricRegistry { std::string(), std::string(), otel_sdk_metrics::AggregationType::kHistogram, +======= + auto selector = std::make_unique( + sdk::metrics::InstrumentType::kHistogram, + name, + "" + ); + auto meterSelector = std::make_unique( + "ydb-cpp-sdk", + GetSdkSemver(), + {} + ); + + auto histogramConfig = std::make_shared(); + histogramConfig->boundaries_ = buckets; + + auto view = std::make_unique( + {}, + {}, + sdk::metrics::AggregationType::kHistogram, +>>>>>>> 1ca4253b5 (fixes and add metric tests) histogramConfig ); sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); } +<<<<<<< HEAD otel_nostd::shared_ptr MeterProvider_; otel_nostd::shared_ptr Meter_; +======= + nostd::shared_ptr MeterProvider_; + nostd::shared_ptr Meter_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) std::mutex HistogramViewsLock_; std::unordered_set HistogramViews_; }; diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt deleted file mode 100644 index d005708d422..00000000000 --- a/plugins/open_telemetry/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -if (YDB_SDK_ENABLE_OTEL_METRICS) - _ydb_sdk_add_library(open_telemetry_metrics) - target_sources(open_telemetry_metrics PRIVATE - src/metrics.cpp - ) - target_include_directories(open_telemetry_metrics PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_metrics PUBLIC - client-metrics - client-resources - opentelemetry-cpp::api - opentelemetry-cpp::metrics - ) - _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) -endif() - -if (YDB_SDK_ENABLE_OTEL_TRACE) - _ydb_sdk_add_library(open_telemetry_trace) - target_sources(open_telemetry_trace PRIVATE - src/trace.cpp - ) - target_include_directories(open_telemetry_trace PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_trace PUBLIC - client-metrics - opentelemetry-cpp::api - opentelemetry-cpp::trace - ) - _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) -endif() - -_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h deleted file mode 100644 index b0db9ea7d7c..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include -#include - -namespace NYdb::inline V3::NMetrics { - -inline void AddOpenTelemetry(TDriverConfig& config - , opentelemetry::nostd::shared_ptr meterProvider - , opentelemetry::nostd::shared_ptr tracerProvider -) { - AddOpenTelemetryMetrics(config, std::move(meterProvider)); - AddOpenTelemetryTrace(config, std::move(tracerProvider)); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h deleted file mode 100644 index 5e9e9e77dea..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/metrics.h +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelMetricRegistry : public IMetricRegistry { -public: - TOtelMetricRegistry(opentelemetry::nostd::shared_ptr meterProvider); - - std::shared_ptr Counter(const std::string& name, const TLabels& labels = {}) override; - std::shared_ptr Gauge(const std::string& name, const TLabels& labels = {}) override; - std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) override; - -private: - void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets); - - opentelemetry::nostd::shared_ptr MeterProvider_; - opentelemetry::nostd::shared_ptr Meter_; - std::mutex HistogramViewsLock_; - std::unordered_set HistogramViews_; -}; - -inline void AddOpenTelemetryMetrics( - TDriverConfig& config, - opentelemetry::nostd::shared_ptr meterProvider) -{ - if (meterProvider) { - config.SetMetricExporter(std::make_shared(std::move(meterProvider))); - } -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h deleted file mode 100644 index 3ba2e146fd9..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelTraceProvider : public ITraceProvider { -public: - TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); - - std::shared_ptr GetTracer(const std::string& name) override; - -private: - opentelemetry::nostd::shared_ptr TracerProvider_; -}; - -inline void AddOpenTelemetryTrace( - TDriverConfig& config, - opentelemetry::nostd::shared_ptr tracerProvider) -{ - if (tracerProvider) { - config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); - } -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/metrics.cpp b/plugins/open_telemetry/src/metrics.cpp deleted file mode 100644 index 65850fd08b3..00000000000 --- a/plugins/open_telemetry/src/metrics.cpp +++ /dev/null @@ -1,140 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include -#include - -namespace NYdb::inline V3::NMetrics { - -namespace { - -using namespace opentelemetry; - -common::KeyValueIterableView MakeAttributes(const TLabels& labels) { - return common::KeyValueIterableView(labels); -} - -class TOtelCounter : public ICounter { -public: - TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) - : Counter_(std::move(counter)) - , Labels_(labels) - {} - - void Inc() override { - Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - } - -private: - nostd::shared_ptr> Counter_; - TLabels Labels_; -}; - -class TOtelUpDownCounterGauge : public IGauge { -public: - TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) - : Counter_(std::move(counter)) - , Labels_(labels) - {} - - void Add(double delta) override { - Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - Value_ += delta; - } - - void Set(double value) override { - Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - Value_ = value; - } - -private: - nostd::shared_ptr> Counter_; - TLabels Labels_; - double Value_ = 0; -}; - -class TOtelHistogram : public IHistogram { -public: - TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) - : Histogram_(std::move(histogram)) - , Labels_(labels) - {} - - void Record(double value) override { - Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); - } - -private: - nostd::shared_ptr> Histogram_; - TLabels Labels_; -}; - -} // namespace - -TOtelMetricRegistry::TOtelMetricRegistry(nostd::shared_ptr meterProvider) - : MeterProvider_(std::move(meterProvider)) - , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) -{} - -void TOtelMetricRegistry::ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { - if (buckets.empty()) { - return; - } - - auto* sdkProvider = dynamic_cast(MeterProvider_.get()); - if (!sdkProvider) { - return; - } - - { - std::lock_guard lock(HistogramViewsLock_); - if (!HistogramViews_.insert(name).second) { - return; - } - } - - auto selector = std::make_unique( - sdk::metrics::InstrumentType::kHistogram, - name, - "" - ); - auto meterSelector = std::make_unique( - "ydb-cpp-sdk", - GetSdkSemver(), - {} - ); - - auto histogramConfig = std::make_shared(); - histogramConfig->boundaries_ = buckets; - - auto view = std::make_unique( - {}, - {}, - sdk::metrics::AggregationType::kHistogram, - histogramConfig - ); - - sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); -} - -std::shared_ptr TOtelMetricRegistry::Counter(const std::string& name, const TLabels& labels) { - auto counter = Meter_->CreateUInt64Counter(name); - return std::make_shared(std::move(counter), labels); -} - -std::shared_ptr TOtelMetricRegistry::Gauge(const std::string& name, const TLabels& labels) { - auto counter = Meter_->CreateDoubleUpDownCounter(name); - return std::make_shared(std::move(counter), labels); -} - -std::shared_ptr TOtelMetricRegistry::Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) { - ConfigureHistogramBuckets(name, buckets); - auto histogram = Meter_->CreateDoubleHistogram(name); - return std::make_shared(std::move(histogram), labels); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/open_telemetry/src/trace.cpp deleted file mode 100644 index 54f04cb84df..00000000000 --- a/plugins/open_telemetry/src/trace.cpp +++ /dev/null @@ -1,70 +0,0 @@ -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -namespace { - -using namespace opentelemetry; - -trace::SpanKind MapSpanKind(ESpanKind kind) { - switch (kind) { - case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; - case ESpanKind::SERVER: return trace::SpanKind::kServer; - case ESpanKind::CLIENT: return trace::SpanKind::kClient; - case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; - case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; - } - return trace::SpanKind::kInternal; -} - -class TOtelSpan : public ISpan { -public: - TOtelSpan(nostd::shared_ptr span) - : Span_(std::move(span)) - {} - - void End() override { - Span_->End(); - } - - void SetAttribute(const std::string& key, const std::string& value) override { - Span_->SetAttribute(key, value); - } - - void SetAttribute(const std::string& key, int64_t value) override { - Span_->SetAttribute(key, value); - } - -private: - nostd::shared_ptr Span_; -}; - -class TOtelTracer : public ITracer { -public: - TOtelTracer(nostd::shared_ptr tracer) - : Tracer_(std::move(tracer)) - {} - - std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { - trace::StartSpanOptions options; - options.kind = MapSpanKind(kind); - return std::make_shared(Tracer_->StartSpan(name, options)); - } - -private: - nostd::shared_ptr Tracer_; -}; - -} // namespace - -TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) - : TracerProvider_(std::move(tracerProvider)) -{} - -std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { - return std::make_shared(TracerProvider_->GetTracer(name)); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h index 68b238d6a41..b4e9f09eab8 100644 --- a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h +++ b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -3,7 +3,14 @@ #include #include +<<<<<<< HEAD #include +======= + +namespace opentelemetry::trace { +class TracerProvider; +} +>>>>>>> 1ca4253b5 (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { diff --git a/plugins/trace/otel/src/trace.cpp b/plugins/trace/otel/src/trace.cpp index b5e6d8ae1ea..cc2c2d0d294 100644 --- a/plugins/trace/otel/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -1,7 +1,10 @@ #include #include +<<<<<<< HEAD #include +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) #include #include @@ -9,6 +12,7 @@ namespace NYdb::inline V3::NMetrics { namespace { +<<<<<<< HEAD namespace otel_trace = opentelemetry::trace; namespace otel_nostd = opentelemetry::nostd; namespace otel_common = opentelemetry::common; @@ -37,6 +41,24 @@ class TOtelScope : public IScope { class TOtelSpan : public ISpan { public: TOtelSpan(otel_nostd::shared_ptr span) +======= +using namespace opentelemetry; + +trace::SpanKind MapSpanKind(ESpanKind kind) { + switch (kind) { + case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; + case ESpanKind::SERVER: return trace::SpanKind::kServer; + case ESpanKind::CLIENT: return trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + } + return trace::SpanKind::kInternal; +} + +class TOtelSpan : public ISpan { +public: + TOtelSpan(nostd::shared_ptr span) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : Span_(std::move(span)) {} @@ -56,42 +78,70 @@ class TOtelSpan : public ISpan { if (attributes.empty()) { Span_->AddEvent(name); } else { +<<<<<<< HEAD std::vector> attrs; attrs.reserve(attributes.size()); for (const auto& [k, v] : attributes) { attrs.emplace_back(otel_nostd::string_view(k), otel_common::AttributeValue(otel_nostd::string_view(v))); +======= + std::vector> attrs; + attrs.reserve(attributes.size()); + for (const auto& [k, v] : attributes) { + attrs.emplace_back(nostd::string_view(k), common::AttributeValue(nostd::string_view(v))); +>>>>>>> 1ca4253b5 (fixes and add metric tests) } Span_->AddEvent(name, attrs); } } +<<<<<<< HEAD std::unique_ptr Activate() override { return std::make_unique(Span_); } private: otel_nostd::shared_ptr Span_; +======= +private: + nostd::shared_ptr Span_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; class TOtelTracer : public ITracer { public: +<<<<<<< HEAD TOtelTracer(otel_nostd::shared_ptr tracer) +======= + TOtelTracer(nostd::shared_ptr tracer) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : Tracer_(std::move(tracer)) {} std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { +<<<<<<< HEAD otel_trace::StartSpanOptions options; +======= + trace::StartSpanOptions options; +>>>>>>> 1ca4253b5 (fixes and add metric tests) options.kind = MapSpanKind(kind); return std::make_shared(Tracer_->StartSpan(name, options)); } private: +<<<<<<< HEAD otel_nostd::shared_ptr Tracer_; +======= + nostd::shared_ptr Tracer_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; class TOtelTraceProvider : public ITraceProvider { public: +<<<<<<< HEAD TOtelTraceProvider(otel_nostd::shared_ptr tracerProvider) +======= + TOtelTraceProvider(nostd::shared_ptr tracerProvider) +>>>>>>> 1ca4253b5 (fixes and add metric tests) : TracerProvider_(std::move(tracerProvider)) {} @@ -100,7 +150,11 @@ class TOtelTraceProvider : public ITraceProvider { } private: +<<<<<<< HEAD otel_nostd::shared_ptr TracerProvider_; +======= + nostd::shared_ptr TracerProvider_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; } // namespace diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index ff831986ddb..8284c4a2c04 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -53,6 +53,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t GetMaxMessageSize() const override { return MaxMessageSize; } const TLog& GetLog() const override { return Log; } std::shared_ptr GetExecutor() const override { return Executor; } +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } std::shared_ptr GetTraceProvider() const override { return TraceProvider; } @@ -60,6 +61,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr GetMetricExporter() const override { return MetricExporter; } std::shared_ptr GetTraceExporter() const override { return TraceExporter; } >>>>>>> 1b2bf4fa5 (fixes) +======= + std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } + std::shared_ptr GetTraceProvider() const override { return TraceProvider; } +>>>>>>> 1ca4253b5 (fixes and add metric tests) std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -91,6 +96,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { uint64_t MaxMessageSize = 0; TLog Log; // Null by default. std::shared_ptr Executor; +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry; std::shared_ptr TraceProvider; @@ -98,6 +104,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr MetricExporter; std::shared_ptr TraceExporter; >>>>>>> 1b2bf4fa5 (fixes) +======= + std::shared_ptr MetricRegistry; + std::shared_ptr TraceProvider; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -257,6 +267,7 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { return *this; } +<<<<<<< HEAD <<<<<<< HEAD TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { Impl_->MetricRegistry = std::move(registry); @@ -271,14 +282,19 @@ TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr exporter) { Impl_->MetricExporter = std::move(exporter); +======= +TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { + Impl_->MetricRegistry = std::move(registry); +>>>>>>> 1ca4253b5 (fixes and add metric tests) return *this; } -TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { - Impl_->TraceExporter = std::move(exporter); +TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr provider) { + Impl_->TraceProvider = std::move(provider); return *this; } +<<<<<<< HEAD std::shared_ptr TDriverConfig::GetMetricExporter() const { return Impl_->MetricExporter; } @@ -288,6 +304,8 @@ std::shared_ptr TDriverConfig::GetTraceExporter() cons } >>>>>>> 1b2bf4fa5 (fixes) +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -341,6 +359,7 @@ TDriverConfig TDriver::GetConfig() const { config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); config.SetMaxMessageSize(Impl_->MaxMessageSize_); config.Impl_->Log = Impl_->Log; +<<<<<<< HEAD <<<<<<< HEAD config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); config.SetTraceProvider(Impl_->GetTraceProvider()); @@ -348,6 +367,10 @@ TDriverConfig TDriver::GetConfig() const { config.SetMetricExporter(Impl_->GetMetricExporter()); config.SetTraceExporter(Impl_->GetTraceExporter()); >>>>>>> 1b2bf4fa5 (fixes) +======= + config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); + config.SetTraceProvider(Impl_->GetTraceProvider()); +>>>>>>> 1ca4253b5 (fixes and add metric tests) return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index c1690a9af17..d711d5a9435 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -169,6 +169,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout(), TcpNoDelay_) #endif +<<<<<<< HEAD <<<<<<< HEAD , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) @@ -176,6 +177,10 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MetricExporter_(params->GetMetricExporter()) , TraceExporter_(params->GetTraceExporter()) >>>>>>> 1b2bf4fa5 (fixes) +======= + , MetricRegistry_(params->GetExternalMetricRegistry()) + , TraceProvider_(params->GetTraceProvider()) +>>>>>>> 1ca4253b5 (fixes and add metric tests) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -443,6 +448,7 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { ExtensionApis_.emplace_back(api); } +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { return MetricRegistry_; @@ -458,6 +464,14 @@ std::shared_ptr TGRpcConnectionsImpl::GetMetricExport std::shared_ptr TGRpcConnectionsImpl::GetTraceExporter() const { return TraceExporter_; >>>>>>> 1b2bf4fa5 (fixes) +======= +std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { + return MetricRegistry_; +} + +std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { + return TraceProvider_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) } void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 6c3d7533e58..c1b1706e8a9 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -586,6 +586,7 @@ class TGRpcConnectionsImpl ::NMonitoring::TMetricRegistry* GetMetricRegistry() override; void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const; std::shared_ptr GetTraceProvider() const; @@ -604,6 +605,10 @@ class TGRpcConnectionsImpl return nullptr; } >>>>>>> 1b2bf4fa5 (fixes) +======= + std::shared_ptr GetExternalMetricRegistry() const; + std::shared_ptr GetTraceProvider() const; +>>>>>>> 1ca4253b5 (fixes and add metric tests) void SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb); const TLog& GetLog() const override; @@ -740,6 +745,7 @@ class TGRpcConnectionsImpl std::vector> Extensions_; std::vector> ExtensionApis_; +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry_; std::shared_ptr TraceProvider_; @@ -747,6 +753,10 @@ class TGRpcConnectionsImpl std::shared_ptr MetricExporter_; std::shared_ptr TraceExporter_; >>>>>>> 1b2bf4fa5 (fixes) +======= + std::shared_ptr MetricRegistry_; + std::shared_ptr TraceProvider_; +>>>>>>> 1ca4253b5 (fixes and add metric tests) IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 61a2e0f5a2e..81e134db107 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -43,6 +43,7 @@ class IConnectionsParams { virtual uint64_t GetMaxOutboundMessageSize() const = 0; virtual uint64_t GetMaxMessageSize() const = 0; virtual std::shared_ptr GetExecutor() const = 0; +<<<<<<< HEAD <<<<<<< HEAD virtual std::shared_ptr GetExternalMetricRegistry() const = 0; virtual std::shared_ptr GetTraceProvider() const = 0; @@ -50,6 +51,10 @@ class IConnectionsParams { virtual std::shared_ptr GetMetricExporter() const = 0; virtual std::shared_ptr GetTraceExporter() const = 0; >>>>>>> 1b2bf4fa5 (fixes) +======= + virtual std::shared_ptr GetExternalMetricRegistry() const = 0; + virtual std::shared_ptr GetTraceProvider() const = 0; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index 7d376c5cac3..f392b31a41c 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -1,11 +1,15 @@ _ydb_sdk_add_library(client-metrics) target_sources(client-metrics PRIVATE +<<<<<<< HEAD <<<<<<< HEAD metrics.cpp ======= metrics.cpp >>>>>>> 1b2bf4fa5 (fixes) +======= + metrics.cpp +>>>>>>> 1ca4253b5 (fixes and add metric tests) ) _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index a5c3ec09242..c7adaa855ae 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -71,11 +71,15 @@ class TQueryClient::TImpl: public TClientImplCommon, public SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); +<<<<<<< HEAD <<<<<<< HEAD if (auto traceProvider = Connections_->GetTraceProvider()) { ======= if (auto traceProvider = Connections_->GetTraceExporter()) { >>>>>>> 1b2bf4fa5 (fixes) +======= + if (auto traceProvider = Connections_->GetTraceProvider()) { +>>>>>>> 1ca4253b5 (fixes and add metric tests) Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } MetricRegistry_ = Connections_->GetExternalMetricRegistry(); @@ -108,6 +112,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public CollectParamsSize(params ? ¶ms->GetProtoMap() : nullptr); auto span = std::make_shared(Tracer_, "ExecuteQuery", DbDriverState_->DiscoveryEndpoint); +<<<<<<< HEAD +======= + span->SetQueryText(query); +>>>>>>> 1ca4253b5 (fixes and add metric tests) auto metrics = std::make_shared(MetricRegistry_, "ExecuteQuery"); return TExecQueryImpl::ExecuteQuery( @@ -494,12 +502,18 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(status.GetStatus()); } +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(status.GetStatus()); } ======= >>>>>>> 1b2bf4fa5 (fixes) +======= + if (Metrics) { + Metrics->End(status.GetStatus()); + } +>>>>>>> 1ca4253b5 (fixes and add metric tests) ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -515,12 +529,18 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (Span) { Span->End(EStatus::SUCCESS); } +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(EStatus::SUCCESS); } ======= >>>>>>> 1b2bf4fa5 (fixes) +======= + if (Metrics) { + Metrics->End(EStatus::SUCCESS); + } +>>>>>>> 1ca4253b5 (fixes and add metric tests) ScheduleReply(std::move(val)); } @@ -531,16 +551,22 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto val = future.ExtractValue(); if (span) { <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) span->SetPeerEndpoint(val.GetEndpoint()); span->End(val.GetStatus()); } if (metrics) { metrics->End(val.GetStatus()); } +<<<<<<< HEAD ======= span->End(val.GetStatus()); } >>>>>>> 1b2bf4fa5 (fixes) +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) promise.SetValue(std::move(val)); }); } diff --git a/src/client/query/impl/CMakeLists.txt b/src/client/query/impl/CMakeLists.txt index fdc46b50f95..9ab36b7c9f9 100644 --- a/src/client/query/impl/CMakeLists.txt +++ b/src/client/query/impl/CMakeLists.txt @@ -16,6 +16,7 @@ target_link_libraries(client-ydb_query-impl PUBLIC target_sources(client-ydb_query-impl PRIVATE exec_query.cpp client_session.cpp + query_metrics.cpp query_spans.cpp ) diff --git a/src/client/query/impl/query_metrics.cpp b/src/client/query/impl/query_metrics.cpp new file mode 100644 index 00000000000..f314f3d8b96 --- /dev/null +++ b/src/client/query/impl/query_metrics.cpp @@ -0,0 +1,71 @@ +#include "query_metrics.h" + +#include + +namespace NYdb::inline V3::NQuery { + +namespace { + +void SafeLogMetricsError(const char* message) noexcept { + try { + try { + std::cerr << "TQueryMetrics: " << message << ": " << CurrentExceptionMessage() << std::endl; + return; + } catch (...) { + } + std::cerr << "TQueryMetrics: " << message << ": (unknown)" << std::endl; + } catch (...) { + } +} + +} // namespace + +static const std::vector LatencyBuckets = { + 1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000 +}; + +TQueryMetrics::TQueryMetrics(std::shared_ptr registry, const std::string& operationName) { + if (!registry) { + return; + } + + try { + NMetrics::TLabels labels = {{"operation", operationName}}; + RequestCounter_ = registry->Counter("ydb.query.requests", labels); + ErrorCounter_ = registry->Counter("ydb.query.errors", labels); + LatencyHistogram_ = registry->Histogram("ydb.query.latency_ms", LatencyBuckets, labels); + + RequestCounter_->Inc(); + StartTime_ = TInstant::Now(); + } catch (...) { + SafeLogMetricsError("failed to initialize metrics"); + RequestCounter_.reset(); + ErrorCounter_.reset(); + LatencyHistogram_.reset(); + } +} + +TQueryMetrics::~TQueryMetrics() noexcept { + End(EStatus::CLIENT_INTERNAL_ERROR); +} + +void TQueryMetrics::End(EStatus status) noexcept { + if (Ended_) { + return; + } + Ended_ = true; + + try { + if (LatencyHistogram_) { + auto durationMs = (TInstant::Now() - StartTime_).MilliSeconds(); + LatencyHistogram_->Record(static_cast(durationMs)); + } + if (status != EStatus::SUCCESS && ErrorCounter_) { + ErrorCounter_->Inc(); + } + } catch (...) { + SafeLogMetricsError("failed to record metrics"); + } +} + +} // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_metrics.h b/src/client/query/impl/query_metrics.h index 2bd284f76f3..e62d7e014c2 100644 --- a/src/client/query/impl/query_metrics.h +++ b/src/client/query/impl/query_metrics.h @@ -1,5 +1,6 @@ #pragma once +<<<<<<< HEAD #include namespace NYdb::inline V3::NQuery { @@ -9,6 +10,31 @@ class TQueryMetrics : public NObservability::TClientMetrics { TQueryMetrics(std::shared_ptr registry, const std::string& operationName) : TClientMetrics(std::move(registry), operationName) {} +======= +#include +#include + +#include + +#include +#include + +namespace NYdb::inline V3::NQuery { + +class TQueryMetrics { +public: + TQueryMetrics(std::shared_ptr registry, const std::string& operationName); + ~TQueryMetrics() noexcept; + + void End(EStatus status) noexcept; + +private: + std::shared_ptr RequestCounter_; + std::shared_ptr ErrorCounter_; + std::shared_ptr LatencyHistogram_; + TInstant StartTime_; + bool Ended_ = false; +>>>>>>> 1ca4253b5 (fixes and add metric tests) }; } // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index b18881559e7..d6fe219885b 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -38,6 +38,7 @@ void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { } catch (...) {} } else { host = endpoint; +<<<<<<< HEAD } } @@ -50,17 +51,19 @@ void SafeLogSpanError(const char* message) noexcept { } std::cerr << "TQuerySpan: " << message << ": (unknown)" << std::endl; } catch (...) { +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) } } void SafeLogSpanError(const char* message) noexcept { try { try { - Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + std::cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << std::endl; return; } catch (...) { } - Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + std::cerr << "TQuerySpan: " << message << ": (unknown)" << std::endl; } catch (...) { } } @@ -110,6 +113,9 @@ TQuerySpan::~TQuerySpan() noexcept { } <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { if (!Span_ || endpoint.empty()) { return; @@ -125,6 +131,20 @@ void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { } } +<<<<<<< HEAD +======= +void TQuerySpan::SetQueryText(const std::string& query) noexcept { + if (!Span_ || query.empty()) { + return; + } + try { + Span_->SetAttribute("db.query.text", query); + } catch (...) { + SafeLogSpanError("failed to set query text"); + } +} + +>>>>>>> 1ca4253b5 (fixes and add metric tests) void TQuerySpan::AddEvent(const std::string& name, const std::map& attributes) noexcept { if (!Span_) { return; @@ -136,6 +156,7 @@ void TQuerySpan::AddEvent(const std::string& name, const std::map TQuerySpan::Activate() noexcept { if (!Span_) { return nullptr; @@ -153,6 +174,8 @@ void TQuerySpan::End(EStatus status) noexcept { try { Span_->SetAttribute("db.response.status_code", ToString(status)); ======= +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index 168f5592843..be8ea02a803 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -15,6 +15,7 @@ class TQuerySpan { TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); ~TQuerySpan() noexcept; +<<<<<<< HEAD <<<<<<< HEAD void SetPeerEndpoint(const std::string& endpoint) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; @@ -22,6 +23,12 @@ class TQuerySpan { ======= >>>>>>> 1b2bf4fa5 (fixes) +======= + void SetPeerEndpoint(const std::string& endpoint) noexcept; + void SetQueryText(const std::string& query) noexcept; + void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; + +>>>>>>> 1ca4253b5 (fixes and add metric tests) void End(EStatus status) noexcept; private: diff --git a/tests/common/fake_metric_registry.h b/tests/common/fake_metric_registry.h index 032234f080f..5f8f10572b4 100644 --- a/tests/common/fake_metric_registry.h +++ b/tests/common/fake_metric_registry.h @@ -68,11 +68,15 @@ struct TMetricKey { class TFakeMetricRegistry : public NMetrics::IMetricRegistry { public: +<<<<<<< HEAD std::shared_ptr Counter(const std::string& name , const NMetrics::TLabels& labels , const std::string& /*description*/ , const std::string& /*unit*/ ) override { +======= + std::shared_ptr Counter(const std::string& name, const NMetrics::TLabels& labels) override { +>>>>>>> 1ca4253b5 (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Counters_.find(key); @@ -84,11 +88,15 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return counter; } +<<<<<<< HEAD std::shared_ptr Gauge(const std::string& name , const NMetrics::TLabels& labels , const std::string& /*description*/ , const std::string& /*unit*/ ) override { +======= + std::shared_ptr Gauge(const std::string& name, const NMetrics::TLabels& labels) override { +>>>>>>> 1ca4253b5 (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto gauge = std::make_shared(); @@ -96,12 +104,16 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return gauge; } +<<<<<<< HEAD std::shared_ptr Histogram(const std::string& name , const std::vector& /*buckets*/ , const NMetrics::TLabels& labels , const std::string& /*description*/ , const std::string& /*unit*/ ) override { +======= + std::shared_ptr Histogram(const std::string& name, const std::vector& /*buckets*/, const NMetrics::TLabels& labels) override { +>>>>>>> 1ca4253b5 (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Histograms_.find(key); diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index bb58495d9bc..e36366c6fb7 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -1,7 +1,10 @@ #include #include #include +<<<<<<< HEAD #include +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) #include @@ -37,6 +40,7 @@ std::shared_ptr GetCounter( const std::string& name, const std::string& operation) { +<<<<<<< HEAD return registry->GetCounter(name, { {"db.system.name", "other_sql"}, {"db.operation.name", operation}, @@ -57,6 +61,17 @@ std::shared_ptr GetDuration( labels["error.type"] = ToString(status); } return registry->GetHistogram("db.client.operation.duration", labels); +======= + return registry->GetCounter(name, {{"operation", operation}}); +} + +std::shared_ptr GetHistogram( + const std::shared_ptr& registry, + const std::string& name, + const std::string& operation) +{ + return registry->GetHistogram(name, {{"operation", operation}}); +>>>>>>> 1ca4253b5 (fixes and add metric tests) } } // namespace @@ -74,6 +89,7 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); +<<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; EXPECT_GE(requests->Get(), 1); @@ -86,6 +102,20 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ASSERT_NE(duration, nullptr) << "ExecuteQuery duration histogram not created"; EXPECT_GE(duration->Count(), 1u); for (double v : duration->GetValues()) { +======= + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; + EXPECT_GE(requests->Get(), 1); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr) << "ExecuteQuery latency histogram not created"; + EXPECT_GE(latency->Count(), 1u); + for (double v : latency->GetValues()) { +>>>>>>> 1ca4253b5 (fixes and add metric tests) EXPECT_GE(v, 0.0); } @@ -105,6 +135,7 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ).ExtractValueSync(); EXPECT_NE(result.GetStatus(), EStatus::SUCCESS); +<<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_GE(requests->Get(), 1); @@ -116,6 +147,19 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { auto duration = GetDuration(registry, "ExecuteQuery", result.GetStatus()); ASSERT_NE(duration, nullptr); EXPECT_GE(duration->Count(), 1u); +======= + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_GE(requests->Get(), 1); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_GE(errors->Get(), 1); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + EXPECT_GE(latency->Count(), 1u); +>>>>>>> 1ca4253b5 (fixes and add metric tests) driver.Stop(true); } @@ -127,6 +171,7 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto session = client.GetSession().ExtractValueSync(); ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); +<<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "GetSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); @@ -134,6 +179,15 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto duration = GetDuration(registry, "GetSession", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; EXPECT_GE(duration->Count(), 1u); +======= + auto requests = GetCounter(registry, "ydb.query.requests", "CreateSession"); + ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; + EXPECT_GE(requests->Get(), 1); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "CreateSession"); + ASSERT_NE(latency, nullptr) << "CreateSession latency histogram not created"; + EXPECT_GE(latency->Count(), 1u); +>>>>>>> 1ca4253b5 (fixes and add metric tests) driver.Stop(true); } @@ -160,6 +214,7 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitResult = execResult.GetTransaction()->Commit().ExtractValueSync(); ASSERT_TRUE(commitResult.IsSuccess()) << commitResult.GetIssues().ToString(); +<<<<<<< HEAD auto commitRequests = GetCounter(registry, "db.client.operation.requests", "Commit"); ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; EXPECT_GE(commitRequests->Get(), 1); @@ -167,6 +222,15 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitDuration = GetDuration(registry, "Commit", EStatus::SUCCESS); ASSERT_NE(commitDuration, nullptr); EXPECT_GE(commitDuration->Count(), 1u); +======= + auto commitRequests = GetCounter(registry, "ydb.query.requests", "Commit"); + ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; + EXPECT_GE(commitRequests->Get(), 1); + + auto commitLatency = GetHistogram(registry, "ydb.query.latency_ms", "Commit"); + ASSERT_NE(commitLatency, nullptr); + EXPECT_GE(commitLatency->Count(), 1u); +>>>>>>> 1ca4253b5 (fixes and add metric tests) } driver.Stop(true); @@ -187,6 +251,7 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackResult = tx.Rollback().ExtractValueSync(); ASSERT_TRUE(rollbackResult.IsSuccess()) << rollbackResult.GetIssues().ToString(); +<<<<<<< HEAD auto rollbackRequests = GetCounter(registry, "db.client.operation.requests", "Rollback"); ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; EXPECT_GE(rollbackRequests->Get(), 1); @@ -198,6 +263,19 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackDuration = GetDuration(registry, "Rollback", EStatus::SUCCESS); ASSERT_NE(rollbackDuration, nullptr); EXPECT_GE(rollbackDuration->Count(), 1u); +======= + auto rollbackRequests = GetCounter(registry, "ydb.query.requests", "Rollback"); + ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; + EXPECT_GE(rollbackRequests->Get(), 1); + + auto rollbackErrors = GetCounter(registry, "ydb.query.errors", "Rollback"); + ASSERT_NE(rollbackErrors, nullptr); + EXPECT_EQ(rollbackErrors->Get(), 0); + + auto rollbackLatency = GetHistogram(registry, "ydb.query.latency_ms", "Rollback"); + ASSERT_NE(rollbackLatency, nullptr); + EXPECT_GE(rollbackLatency->Count(), 1u); +>>>>>>> 1ca4253b5 (fixes and add metric tests) driver.Stop(true); } @@ -219,6 +297,7 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); } +<<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_EQ(requests->Get(), numQueries); @@ -230,6 +309,19 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); EXPECT_EQ(duration->Count(), static_cast(numQueries)); +======= + auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), numQueries); + + auto errors = GetCounter(registry, "ydb.query.errors", "ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + EXPECT_EQ(latency->Count(), static_cast(numQueries)); +>>>>>>> 1ca4253b5 (fixes and add metric tests) driver.Stop(true); } @@ -258,7 +350,11 @@ TEST(QueryMetricsIntegration, NoRegistryDoesNotBreakOperations) { driver.Stop(true); } +<<<<<<< HEAD TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { +======= +TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { +>>>>>>> 1ca4253b5 (fixes and add metric tests) auto [driver, registry] = MakeRunArgs(); TQueryClient client(driver); @@ -272,6 +368,7 @@ TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); +<<<<<<< HEAD auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); ASSERT_GE(duration->Count(), 1u); @@ -279,6 +376,15 @@ TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { for (double v : duration->GetValues()) { EXPECT_GE(v, 0.0) << "Duration must be non-negative"; EXPECT_LT(v, 30.0) << "Duration > 30s is unrealistic for SELECT 1"; +======= + auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); + ASSERT_NE(latency, nullptr); + ASSERT_GE(latency->Count(), 1u); + + for (double v : latency->GetValues()) { + EXPECT_GE(v, 0.0) << "Latency must be non-negative"; + EXPECT_LT(v, 30000.0) << "Latency > 30s is unrealistic for SELECT 1"; +>>>>>>> 1ca4253b5 (fixes and add metric tests) } driver.Stop(true); diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index ee15a8cd2e8..99a3387c250 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -101,6 +101,7 @@ add_ydb_test(NAME client-ydb_value_ut GTEST unit ) +<<<<<<< HEAD add_ydb_test(NAME client-ydb_metrics_ut GTEST INCLUDE_DIRS ${YDB_SDK_SOURCE_DIR} @@ -111,10 +112,21 @@ add_ydb_test(NAME client-ydb_metrics_ut GTEST impl-observability client-ydb_query-impl client-ydb_table-impl +======= +add_ydb_test(NAME client-ydb_query_metrics_ut GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + query/query_metrics_ut.cpp + LINK_LIBRARIES + yutil + client-ydb_query-impl +>>>>>>> 1ca4253b5 (fixes and add metric tests) client-metrics LABELS unit ) +<<<<<<< HEAD add_ydb_test(NAME client-ydb_query_spans_ut GTEST INCLUDE_DIRS @@ -128,3 +140,5 @@ add_ydb_test(NAME client-ydb_query_spans_ut GTEST LABELS unit ) +======= +>>>>>>> 1ca4253b5 (fixes and add metric tests) diff --git a/tests/unit/client/query/query_metrics_ut.cpp b/tests/unit/client/query/query_metrics_ut.cpp new file mode 100644 index 00000000000..20c681b7eca --- /dev/null +++ b/tests/unit/client/query/query_metrics_ut.cpp @@ -0,0 +1,190 @@ +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NQuery; +using namespace NYdb::NMetrics; +using namespace NYdb::NTests; + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +class QueryMetricsTest : public ::testing::Test { +protected: + void SetUp() override { + Registry = std::make_shared(); + } + + std::shared_ptr RequestCounter(const std::string& op) { + return Registry->GetCounter("ydb.query.requests", {{"operation", op}}); + } + + std::shared_ptr ErrorCounter(const std::string& op) { + return Registry->GetCounter("ydb.query.errors", {{"operation", op}}); + } + + std::shared_ptr LatencyHistogram(const std::string& op) { + return Registry->GetHistogram("ydb.query.latency_ms", {{"operation", op}}); + } + + std::shared_ptr Registry; +}; + +TEST_F(QueryMetricsTest, RequestCounterIncrementedOnConstruction) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + + auto counter = RequestCounter("ExecuteQuery"); + ASSERT_NE(counter, nullptr); + EXPECT_EQ(counter->Get(), 1); +} + +TEST_F(QueryMetricsTest, SuccessDoesNotIncrementErrorCounter) { + { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + } + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); +} + +TEST_F(QueryMetricsTest, FailureIncrementsErrorCounter) { + { + TQueryMetrics metrics(Registry, "Commit"); + metrics.End(EStatus::UNAVAILABLE); + } + + auto errors = ErrorCounter("Commit"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); +} + +TEST_F(QueryMetricsTest, LatencyRecordedOnEnd) { + { + TQueryMetrics metrics(Registry, "Rollback"); + metrics.End(EStatus::SUCCESS); + } + + auto hist = LatencyHistogram("Rollback"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); + EXPECT_GE(hist->GetValues()[0], 0.0); +} + +TEST_F(QueryMetricsTest, DoubleEndIsIdempotent) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + metrics.End(EStatus::INTERNAL_ERROR); + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 0); + + auto hist = LatencyHistogram("ExecuteQuery"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(QueryMetricsTest, DestructorCallsEndWithClientInternalError) { + { + TQueryMetrics metrics(Registry, "CreateSession"); + } + + auto requests = RequestCounter("CreateSession"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), 1); + + auto errors = ErrorCounter("CreateSession"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 1); + + auto hist = LatencyHistogram("CreateSession"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 1u); +} + +TEST_F(QueryMetricsTest, NullRegistryDoesNotCrash) { + EXPECT_NO_THROW({ + TQueryMetrics metrics(nullptr, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + }); +} + +TEST_F(QueryMetricsTest, CorrectMetricNamesAndLabels) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(EStatus::SUCCESS); + + EXPECT_NE(Registry->GetCounter("ydb.query.requests", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(Registry->GetCounter("ydb.query.errors", {{"operation", "ExecuteQuery"}}), nullptr); + EXPECT_NE(Registry->GetHistogram("ydb.query.latency_ms", {{"operation", "ExecuteQuery"}}), nullptr); + + EXPECT_EQ(Registry->GetCounter("ydb.query.requests", {{"operation", "Commit"}}), nullptr); +} + +TEST_F(QueryMetricsTest, DifferentOperationsHaveSeparateMetrics) { + { + TQueryMetrics m1(Registry, "ExecuteQuery"); + m1.End(EStatus::SUCCESS); + } + { + TQueryMetrics m2(Registry, "Commit"); + m2.End(EStatus::OVERLOADED); + } + + auto execRequests = RequestCounter("ExecuteQuery"); + auto commitRequests = RequestCounter("Commit"); + ASSERT_NE(execRequests, nullptr); + ASSERT_NE(commitRequests, nullptr); + EXPECT_EQ(execRequests->Get(), 1); + EXPECT_EQ(commitRequests->Get(), 1); + + auto execErrors = ErrorCounter("ExecuteQuery"); + auto commitErrors = ErrorCounter("Commit"); + EXPECT_EQ(execErrors->Get(), 0); + EXPECT_EQ(commitErrors->Get(), 1); +} + +TEST_F(QueryMetricsTest, MultipleRequestsAccumulate) { + for (int i = 0; i < 5; ++i) { + TQueryMetrics metrics(Registry, "ExecuteQuery"); + metrics.End(i % 2 == 0 ? EStatus::SUCCESS : EStatus::TIMEOUT); + } + + auto requests = RequestCounter("ExecuteQuery"); + ASSERT_NE(requests, nullptr); + EXPECT_EQ(requests->Get(), 5); + + auto errors = ErrorCounter("ExecuteQuery"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), 2); + + auto hist = LatencyHistogram("ExecuteQuery"); + ASSERT_NE(hist, nullptr); + EXPECT_EQ(hist->Count(), 5u); +} + +TEST_F(QueryMetricsTest, AllErrorStatusesIncrementErrorCounter) { + std::vector errorStatuses = { + EStatus::BAD_REQUEST, + EStatus::UNAUTHORIZED, + EStatus::INTERNAL_ERROR, + EStatus::UNAVAILABLE, + EStatus::OVERLOADED, + EStatus::TIMEOUT, + EStatus::NOT_FOUND, + EStatus::CLIENT_INTERNAL_ERROR, + }; + + for (auto status : errorStatuses) { + TQueryMetrics metrics(Registry, "Rollback"); + metrics.End(status); + } + + auto errors = ErrorCounter("Rollback"); + ASSERT_NE(errors, nullptr); + EXPECT_EQ(errors->Get(), static_cast(errorStatuses.size())); +} From bb87881ae7fb543d78222aeaeda6a73a64914774 Mon Sep 17 00:00:00 2001 From: Bulat Date: Thu, 12 Feb 2026 13:39:18 +0000 Subject: [PATCH 84/93] [C++ SDK] Supported TCP_NODELAY option for grpc sockets (#32631) --- .github/last_commit.txt | 2 +- src/client/impl/internal/common/client_pid.cpp | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 7b3760d4b24..00bc1407094 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -688edef85e4e4e55828630ef3943a91ca9306799 +f64a82eff2c4bbe405178f8e71acada58dbc3c39 \ No newline at end of file diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index 4bc2136a99b..cb5cc8eb0a5 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,8 +3,6 @@ #include -#include - #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. From 086102936dc690f3cc4d7cad87584f94e98c10da Mon Sep 17 00:00:00 2001 From: maladetska Date: Sun, 1 Mar 2026 02:20:51 +0300 Subject: [PATCH 85/93] fixes --- include/ydb-cpp-sdk/client/driver/driver.h | 16 +++++ include/ydb-cpp-sdk/client/metrics/metrics.h | 6 ++ plugins/CMakeLists.txt | 6 ++ plugins/open_telemetry/CMakeLists.txt | 36 ++++++++++ .../ydb-cpp-sdk/open_telemetry/extension.h | 16 +++++ .../ydb-cpp-sdk/open_telemetry/trace.h | 29 ++++++++ plugins/open_telemetry/src/trace.cpp | 70 +++++++++++++++++++ src/client/driver/driver.cpp | 31 ++++++++ .../grpc_connections/grpc_connections.cpp | 11 +++ .../grpc_connections/grpc_connections.h | 8 +++ .../impl/internal/grpc_connections/params.h | 5 ++ src/client/metrics/CMakeLists.txt | 3 + src/client/query/client.cpp | 15 ++++ src/client/query/impl/query_spans.cpp | 24 +++++++ src/client/query/impl/query_spans.h | 3 + 15 files changed, 279 insertions(+) create mode 100644 plugins/open_telemetry/CMakeLists.txt create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h create mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h create mode 100644 plugins/open_telemetry/src/trace.cpp diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index 4471fda9277..d484aaeeb73 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -168,6 +169,7 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); +<<<<<<< HEAD <<<<<<< HEAD ======= //! Set external metrics registry implementation. @@ -177,6 +179,20 @@ class TDriverConfig { TDriverConfig& SetTraceProvider(std::shared_ptr provider); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + //! Set external metrics exporter implementation. + TDriverConfig& SetMetricExporter(std::shared_ptr exporter); + + //! Set external tracing exporter implementation. + TDriverConfig& SetTraceExporter(std::shared_ptr exporter); + + //! Get configured metrics exporter implementation. + std::shared_ptr GetMetricExporter() const; + + //! Get configured tracing exporter implementation. + std::shared_ptr GetTraceExporter() const; + +>>>>>>> a979e6bda (fixes) private: class TImpl; std::shared_ptr Impl_; diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index 1430027fc55..e90b0f50a34 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -1,9 +1,12 @@ #pragma once +<<<<<<< HEAD <<<<<<< HEAD #include ======= >>>>>>> 1b2bf4fa5 (fixes) +======= +>>>>>>> a979e6bda (fixes) #include namespace NYdb::inline V3::NMetrics { @@ -89,6 +92,9 @@ class ITraceProvider { >>>>>>> 1b2bf4fa5 (fixes) }; +<<<<<<< HEAD ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) } // namespace NYdb::NMetrics diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 6f38e537e6b..69d1a265ed5 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -1,5 +1,6 @@ <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD add_subdirectory(metrics) add_subdirectory(trace) ======= @@ -11,3 +12,8 @@ endif() add_subdirectory(metrics) add_subdirectory(trace) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) + add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) +endif() +>>>>>>> a979e6bda (fixes) diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt new file mode 100644 index 00000000000..d005708d422 --- /dev/null +++ b/plugins/open_telemetry/CMakeLists.txt @@ -0,0 +1,36 @@ +if (YDB_SDK_ENABLE_OTEL_METRICS) + _ydb_sdk_add_library(open_telemetry_metrics) + target_sources(open_telemetry_metrics PRIVATE + src/metrics.cpp + ) + target_include_directories(open_telemetry_metrics PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_metrics PUBLIC + client-metrics + client-resources + opentelemetry-cpp::api + opentelemetry-cpp::metrics + ) + _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) +endif() + +if (YDB_SDK_ENABLE_OTEL_TRACE) + _ydb_sdk_add_library(open_telemetry_trace) + target_sources(open_telemetry_trace PRIVATE + src/trace.cpp + ) + target_include_directories(open_telemetry_trace PUBLIC + $ + $ + ) + target_link_libraries(open_telemetry_trace PUBLIC + client-metrics + opentelemetry-cpp::api + opentelemetry-cpp::trace + ) + _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) +endif() + +_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h new file mode 100644 index 00000000000..b0db9ea7d7c --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace NYdb::inline V3::NMetrics { + +inline void AddOpenTelemetry(TDriverConfig& config + , opentelemetry::nostd::shared_ptr meterProvider + , opentelemetry::nostd::shared_ptr tracerProvider +) { + AddOpenTelemetryMetrics(config, std::move(meterProvider)); + AddOpenTelemetryTrace(config, std::move(tracerProvider)); +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h new file mode 100644 index 00000000000..3ba2e146fd9 --- /dev/null +++ b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -0,0 +1,29 @@ +#pragma once + +#include +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +class TOtelTraceProvider : public ITraceProvider { +public: + TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); + + std::shared_ptr GetTracer(const std::string& name) override; + +private: + opentelemetry::nostd::shared_ptr TracerProvider_; +}; + +inline void AddOpenTelemetryTrace( + TDriverConfig& config, + opentelemetry::nostd::shared_ptr tracerProvider) +{ + if (tracerProvider) { + config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); + } +} + +} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/open_telemetry/src/trace.cpp new file mode 100644 index 00000000000..54f04cb84df --- /dev/null +++ b/plugins/open_telemetry/src/trace.cpp @@ -0,0 +1,70 @@ +#include + +#include + +namespace NYdb::inline V3::NMetrics { + +namespace { + +using namespace opentelemetry; + +trace::SpanKind MapSpanKind(ESpanKind kind) { + switch (kind) { + case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; + case ESpanKind::SERVER: return trace::SpanKind::kServer; + case ESpanKind::CLIENT: return trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + } + return trace::SpanKind::kInternal; +} + +class TOtelSpan : public ISpan { +public: + TOtelSpan(nostd::shared_ptr span) + : Span_(std::move(span)) + {} + + void End() override { + Span_->End(); + } + + void SetAttribute(const std::string& key, const std::string& value) override { + Span_->SetAttribute(key, value); + } + + void SetAttribute(const std::string& key, int64_t value) override { + Span_->SetAttribute(key, value); + } + +private: + nostd::shared_ptr Span_; +}; + +class TOtelTracer : public ITracer { +public: + TOtelTracer(nostd::shared_ptr tracer) + : Tracer_(std::move(tracer)) + {} + + std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { + trace::StartSpanOptions options; + options.kind = MapSpanKind(kind); + return std::make_shared(Tracer_->StartSpan(name, options)); + } + +private: + nostd::shared_ptr Tracer_; +}; + +} // namespace + +TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) + : TracerProvider_(std::move(tracerProvider)) +{} + +std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { + return std::make_shared(TracerProvider_->GetTracer(name)); +} + +} // namespace NYdb::NMetrics diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 8284c4a2c04..397b79e7566 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -54,6 +54,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { const TLog& GetLog() const override { return Log; } std::shared_ptr GetExecutor() const override { return Executor; } <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } std::shared_ptr GetTraceProvider() const override { return TraceProvider; } @@ -65,6 +66,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } std::shared_ptr GetTraceProvider() const override { return TraceProvider; } >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr GetMetricExporter() const override { return MetricExporter; } + std::shared_ptr GetTraceExporter() const override { return TraceExporter; } +>>>>>>> a979e6bda (fixes) std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -97,6 +102,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { TLog Log; // Null by default. std::shared_ptr Executor; <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry; std::shared_ptr TraceProvider; @@ -108,6 +114,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr MetricRegistry; std::shared_ptr TraceProvider; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr MetricExporter; + std::shared_ptr TraceExporter; +>>>>>>> a979e6bda (fixes) }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -267,6 +277,7 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { return *this; } +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { @@ -295,6 +306,18 @@ TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr exporter) { + Impl_->MetricExporter = std::move(exporter); + return *this; +} + +TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { + Impl_->TraceExporter = std::move(exporter); + return *this; +} + +>>>>>>> a979e6bda (fixes) std::shared_ptr TDriverConfig::GetMetricExporter() const { return Impl_->MetricExporter; } @@ -303,9 +326,12 @@ std::shared_ptr TDriverConfig::GetTraceExporter() cons return Impl_->TraceExporter; } +<<<<<<< HEAD >>>>>>> 1b2bf4fa5 (fixes) ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -360,6 +386,7 @@ TDriverConfig TDriver::GetConfig() const { config.SetMaxMessageSize(Impl_->MaxMessageSize_); config.Impl_->Log = Impl_->Log; <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); config.SetTraceProvider(Impl_->GetTraceProvider()); @@ -371,6 +398,10 @@ TDriverConfig TDriver::GetConfig() const { config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); config.SetTraceProvider(Impl_->GetTraceProvider()); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + config.SetMetricExporter(Impl_->GetMetricExporter()); + config.SetTraceExporter(Impl_->GetTraceExporter()); +>>>>>>> a979e6bda (fixes) return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index d711d5a9435..e766741e152 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -170,6 +170,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout(), TcpNoDelay_) #endif <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) @@ -181,6 +182,10 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + , MetricExporter_(params->GetMetricExporter()) + , TraceExporter_(params->GetTraceExporter()) +>>>>>>> a979e6bda (fixes) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -448,6 +453,7 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { ExtensionApis_.emplace_back(api); } +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { @@ -457,12 +463,15 @@ std::shared_ptr TGRpcConnectionsImpl::GetExternalMetr std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { return TraceProvider_; ======= +======= +>>>>>>> a979e6bda (fixes) std::shared_ptr TGRpcConnectionsImpl::GetMetricExporter() const { return MetricExporter_; } std::shared_ptr TGRpcConnectionsImpl::GetTraceExporter() const { return TraceExporter_; +<<<<<<< HEAD >>>>>>> 1b2bf4fa5 (fixes) ======= std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { @@ -472,6 +481,8 @@ std::shared_ptr TGRpcConnectionsImpl::GetExternalMetr std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { return TraceProvider_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) } void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index c1b1706e8a9..3b57d085f1f 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -587,10 +587,13 @@ class TGRpcConnectionsImpl void RegisterExtension(IExtension* extension); void RegisterExtensionApi(IExtensionApi* api); <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const; std::shared_ptr GetTraceProvider() const; ======= +======= +>>>>>>> a979e6bda (fixes) std::shared_ptr GetMetricExporter() const; std::shared_ptr GetTraceExporter() const; @@ -746,6 +749,7 @@ class TGRpcConnectionsImpl std::vector> Extensions_; std::vector> ExtensionApis_; <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry_; std::shared_ptr TraceProvider_; @@ -757,6 +761,10 @@ class TGRpcConnectionsImpl std::shared_ptr MetricRegistry_; std::shared_ptr TraceProvider_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr MetricExporter_; + std::shared_ptr TraceExporter_; +>>>>>>> a979e6bda (fixes) IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 81e134db107..30554cff67d 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -44,6 +44,7 @@ class IConnectionsParams { virtual uint64_t GetMaxMessageSize() const = 0; virtual std::shared_ptr GetExecutor() const = 0; <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD virtual std::shared_ptr GetExternalMetricRegistry() const = 0; virtual std::shared_ptr GetTraceProvider() const = 0; @@ -55,6 +56,10 @@ class IConnectionsParams { virtual std::shared_ptr GetExternalMetricRegistry() const = 0; virtual std::shared_ptr GetTraceProvider() const = 0; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + virtual std::shared_ptr GetMetricExporter() const = 0; + virtual std::shared_ptr GetTraceExporter() const = 0; +>>>>>>> a979e6bda (fixes) }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index f392b31a41c..ef4841c5026 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -6,10 +6,13 @@ target_sources(client-metrics PRIVATE metrics.cpp ======= metrics.cpp +<<<<<<< HEAD >>>>>>> 1b2bf4fa5 (fixes) ======= metrics.cpp >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) ) _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index c7adaa855ae..b631534e6a1 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -71,6 +71,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public SetStatCollector(DbDriverState_->StatCollector.GetClientStatCollector("Query")); SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Query")); +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD if (auto traceProvider = Connections_->GetTraceProvider()) { @@ -80,6 +81,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public ======= if (auto traceProvider = Connections_->GetTraceProvider()) { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + if (auto traceProvider = Connections_->GetTraceExporter()) { +>>>>>>> a979e6bda (fixes) Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } MetricRegistry_ = Connections_->GetExternalMetricRegistry(); @@ -503,6 +507,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public Span->End(status.GetStatus()); } <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(status.GetStatus()); @@ -514,6 +519,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public Metrics->End(status.GetStatus()); } >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -530,6 +537,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public Span->End(EStatus::SUCCESS); } <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(EStatus::SUCCESS); @@ -541,6 +549,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public Metrics->End(EStatus::SUCCESS); } >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) ScheduleReply(std::move(val)); } @@ -552,6 +562,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public if (span) { <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) span->SetPeerEndpoint(val.GetEndpoint()); @@ -567,6 +578,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public >>>>>>> 1b2bf4fa5 (fixes) ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + span->End(val.GetStatus()); + } +>>>>>>> a979e6bda (fixes) promise.SetValue(std::move(val)); }); } diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index d6fe219885b..9facc7c02dd 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -68,6 +68,18 @@ void SafeLogSpanError(const char* message) noexcept { } } +void SafeLogSpanError(const char* message) noexcept { + try { + try { + Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + return; + } catch (...) { + } + Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + } catch (...) { + } +} + } // namespace TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { @@ -80,6 +92,7 @@ TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::str ParseEndpoint(endpoint, host, port); try { +<<<<<<< HEAD <<<<<<< HEAD Span_ = tracer->StartSpan(operationName, NMetrics::ESpanKind::CLIENT); if (!Span_) { @@ -88,12 +101,17 @@ TQuerySpan::TQuerySpan(std::shared_ptr tracer, const std::str Span_->SetAttribute("db.system.name", "other_sql"); Span_->SetAttribute("db.operation.name", operationName); ======= +======= +>>>>>>> a979e6bda (fixes) Span_ = tracer->StartSpan("ydb." + operationName, NMetrics::ESpanKind::CLIENT); if (!Span_) { return; } Span_->SetAttribute("db.system.name", "ydb"); +<<<<<<< HEAD >>>>>>> 1b2bf4fa5 (fixes) +======= +>>>>>>> a979e6bda (fixes) Span_->SetAttribute("server.address", host); Span_->SetAttribute("server.port", static_cast(port)); } catch (...) { @@ -112,6 +130,7 @@ TQuerySpan::~TQuerySpan() noexcept { } } +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD ======= @@ -176,11 +195,16 @@ void TQuerySpan::End(EStatus status) noexcept { ======= ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { Span_->SetAttribute("db.response.status_code", static_cast(status)); +<<<<<<< HEAD >>>>>>> 1b2bf4fa5 (fixes) +======= +>>>>>>> a979e6bda (fixes) if (status != EStatus::SUCCESS) { Span_->SetAttribute("error.type", ToString(status)); } diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index be8ea02a803..3c8d6443b2e 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -15,6 +15,7 @@ class TQuerySpan { TQuerySpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); ~TQuerySpan() noexcept; +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD void SetPeerEndpoint(const std::string& endpoint) noexcept; @@ -29,6 +30,8 @@ class TQuerySpan { void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> a979e6bda (fixes) void End(EStatus status) noexcept; private: From 76591da9e3ecacc6536f9a109dccd08fbb890242 Mon Sep 17 00:00:00 2001 From: maladetska Date: Mon, 16 Mar 2026 01:55:33 +0800 Subject: [PATCH 86/93] fixes and add metric tests --- include/ydb-cpp-sdk/client/driver/driver.h | 16 ++-- include/ydb-cpp-sdk/client/metrics/metrics.h | 3 + include/ydb-cpp-sdk/client/trace/trace.h | 6 ++ plugins/CMakeLists.txt | 5 ++ .../ydb-cpp-sdk/open_telemetry/metrics.h | 6 ++ plugins/metrics/otel/src/metrics.cpp | 74 +++++++++++++++++++ plugins/open_telemetry/CMakeLists.txt | 36 --------- .../ydb-cpp-sdk/open_telemetry/extension.h | 16 ---- .../ydb-cpp-sdk/open_telemetry/trace.h | 29 -------- plugins/open_telemetry/src/trace.cpp | 70 ------------------ .../ydb-cpp-sdk/open_telemetry/trace.h | 6 ++ plugins/trace/otel/src/trace.cpp | 40 ++++++++++ src/client/driver/driver.cpp | 27 ++++++- .../grpc_connections/grpc_connections.cpp | 14 ++++ .../grpc_connections/grpc_connections.h | 10 +++ .../impl/internal/grpc_connections/params.h | 5 ++ src/client/metrics/CMakeLists.txt | 4 + src/client/query/client.cpp | 26 +++++++ src/client/query/impl/query_metrics.h | 6 ++ src/client/query/impl/query_spans.cpp | 19 ++++- src/client/query/impl/query_spans.h | 6 ++ tests/common/fake_metric_registry.h | 12 +++ tests/integration/metrics/main.cpp | 55 ++++++++++++++ tests/unit/client/CMakeLists.txt | 9 +++ 24 files changed, 336 insertions(+), 164 deletions(-) delete mode 100644 plugins/open_telemetry/CMakeLists.txt delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h delete mode 100644 plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h delete mode 100644 plugins/open_telemetry/src/trace.cpp diff --git a/include/ydb-cpp-sdk/client/driver/driver.h b/include/ydb-cpp-sdk/client/driver/driver.h index d484aaeeb73..3ab0f79f953 100644 --- a/include/ydb-cpp-sdk/client/driver/driver.h +++ b/include/ydb-cpp-sdk/client/driver/driver.h @@ -8,7 +8,6 @@ #include #include #include -#include #include #include #include @@ -169,6 +168,7 @@ class TDriverConfig { //! If not set, default executor will be used. TDriverConfig& SetExecutor(std::shared_ptr executor); +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD ======= @@ -182,15 +182,13 @@ class TDriverConfig { ======= //! Set external metrics exporter implementation. TDriverConfig& SetMetricExporter(std::shared_ptr exporter); +======= + //! Set external metrics registry implementation. + TDriverConfig& SetMetricRegistry(std::shared_ptr registry); +>>>>>>> dcae6d69e (fixes and add metric tests) - //! Set external tracing exporter implementation. - TDriverConfig& SetTraceExporter(std::shared_ptr exporter); - - //! Get configured metrics exporter implementation. - std::shared_ptr GetMetricExporter() const; - - //! Get configured tracing exporter implementation. - std::shared_ptr GetTraceExporter() const; + //! Set external trace provider implementation. + TDriverConfig& SetTraceProvider(std::shared_ptr provider); >>>>>>> a979e6bda (fixes) private: diff --git a/include/ydb-cpp-sdk/client/metrics/metrics.h b/include/ydb-cpp-sdk/client/metrics/metrics.h index e90b0f50a34..3a33ca33750 100644 --- a/include/ydb-cpp-sdk/client/metrics/metrics.h +++ b/include/ydb-cpp-sdk/client/metrics/metrics.h @@ -62,6 +62,7 @@ class IMetricRegistry { virtual std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels = {}) = 0; }; +<<<<<<< HEAD <<<<<<< HEAD enum class ESpanKind { INTERNAL, @@ -97,4 +98,6 @@ class ITraceProvider { >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } // namespace NYdb::NMetrics diff --git a/include/ydb-cpp-sdk/client/trace/trace.h b/include/ydb-cpp-sdk/client/trace/trace.h index 5644720cd92..98002cdb636 100644 --- a/include/ydb-cpp-sdk/client/trace/trace.h +++ b/include/ydb-cpp-sdk/client/trace/trace.h @@ -15,6 +15,7 @@ enum class ESpanKind { CONSUMER }; +<<<<<<< HEAD <<<<<<< HEAD class IScope { public: @@ -23,6 +24,8 @@ class IScope { ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) class ISpan { public: virtual ~ISpan() = default; @@ -30,10 +33,13 @@ class ISpan { virtual void SetAttribute(const std::string& key, const std::string& value) = 0; virtual void SetAttribute(const std::string& key, int64_t value) = 0; virtual void AddEvent(const std::string& name, const std::map& attributes = {}) = 0; +<<<<<<< HEAD <<<<<<< HEAD virtual std::unique_ptr Activate() = 0; ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) }; class ITracer { diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 69d1a265ed5..50ac60ce0f8 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -1,6 +1,7 @@ <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD add_subdirectory(metrics) add_subdirectory(trace) ======= @@ -17,3 +18,7 @@ if (YDB_SDK_ENABLE_OTEL_METRICS OR YDB_SDK_ENABLE_OTEL_TRACE) add_subdirectory(open_telemetry EXCLUDE_FROM_ALL) endif() >>>>>>> a979e6bda (fixes) +======= +add_subdirectory(metrics) +add_subdirectory(trace) +>>>>>>> dcae6d69e (fixes and add metric tests) diff --git a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h index 698e2efeb92..9152c8ccff3 100644 --- a/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h +++ b/plugins/metrics/otel/include/ydb-cpp-sdk/open_telemetry/metrics.h @@ -4,13 +4,19 @@ #include <<<<<<< HEAD +<<<<<<< HEAD #include ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) namespace opentelemetry::metrics { class MeterProvider; } +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { diff --git a/plugins/metrics/otel/src/metrics.cpp b/plugins/metrics/otel/src/metrics.cpp index b78b4efacb7..078961ab1cb 100644 --- a/plugins/metrics/otel/src/metrics.cpp +++ b/plugins/metrics/otel/src/metrics.cpp @@ -9,15 +9,19 @@ #include #include +<<<<<<< HEAD <<<<<<< HEAD #include ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { namespace { +<<<<<<< HEAD <<<<<<< HEAD namespace otel_metrics = opentelemetry::metrics; namespace otel_nostd = opentelemetry::nostd; @@ -28,25 +32,35 @@ namespace otel_sdk_metrics = opentelemetry::sdk::metrics; otel_common::KeyValueIterableView MakeAttributes(const TLabels& labels) { return otel_common::KeyValueIterableView(labels); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) using namespace opentelemetry; common::KeyValueIterableView MakeAttributes(const TLabels& labels) { return common::KeyValueIterableView(labels); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } class TOtelCounter : public ICounter { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelCounter(otel_nostd::shared_ptr> counter, const TLabels& labels) ======= TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelCounter(nostd::shared_ptr> counter, const TLabels& labels) +>>>>>>> dcae6d69e (fixes and add metric tests) : Counter_(std::move(counter)) , Labels_(labels) {} void Inc() override { +<<<<<<< HEAD <<<<<<< HEAD Counter_->Add(1, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } @@ -54,66 +68,92 @@ class TOtelCounter : public ICounter { private: otel_nostd::shared_ptr> Counter_; ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) Counter_->Add(1, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); } private: nostd::shared_ptr> Counter_; +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) TLabels Labels_; }; class TOtelUpDownCounterGauge : public IGauge { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelUpDownCounterGauge(otel_nostd::shared_ptr> counter, const TLabels& labels) ======= TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelUpDownCounterGauge(nostd::shared_ptr> counter, const TLabels& labels) +>>>>>>> dcae6d69e (fixes and add metric tests) : Counter_(std::move(counter)) , Labels_(labels) {} void Add(double delta) override { +<<<<<<< HEAD <<<<<<< HEAD Counter_->Add(delta, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); ======= Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + Counter_->Add(delta, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); +>>>>>>> dcae6d69e (fixes and add metric tests) Value_ += delta; } void Set(double value) override { +<<<<<<< HEAD <<<<<<< HEAD Counter_->Add(value - Value_, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); ======= Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + Counter_->Add(value - Value_, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); +>>>>>>> dcae6d69e (fixes and add metric tests) Value_ = value; } private: +<<<<<<< HEAD <<<<<<< HEAD otel_nostd::shared_ptr> Counter_; ======= nostd::shared_ptr> Counter_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + nostd::shared_ptr> Counter_; +>>>>>>> dcae6d69e (fixes and add metric tests) TLabels Labels_; double Value_ = 0; }; class TOtelHistogram : public IHistogram { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelHistogram(otel_nostd::shared_ptr> histogram, const TLabels& labels) ======= TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelHistogram(nostd::shared_ptr> histogram, const TLabels& labels) +>>>>>>> dcae6d69e (fixes and add metric tests) : Histogram_(std::move(histogram)) , Labels_(labels) {} void Record(double value) override { +<<<<<<< HEAD <<<<<<< HEAD Histogram_->Record(value, MakeAttributes(Labels_), otel_context::RuntimeContext::GetCurrent()); } @@ -121,26 +161,36 @@ class TOtelHistogram : public IHistogram { private: otel_nostd::shared_ptr> Histogram_; ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) Histogram_->Record(value, MakeAttributes(Labels_), context::RuntimeContext::GetCurrent()); } private: nostd::shared_ptr> Histogram_; +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) TLabels Labels_; }; class TOtelMetricRegistry : public IMetricRegistry { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelMetricRegistry(otel_nostd::shared_ptr meterProvider) ======= TOtelMetricRegistry(nostd::shared_ptr meterProvider) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelMetricRegistry(nostd::shared_ptr meterProvider) +>>>>>>> dcae6d69e (fixes and add metric tests) : MeterProvider_(std::move(meterProvider)) , Meter_(MeterProvider_->GetMeter("ydb-cpp-sdk", GetSdkSemver())) {} +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr Counter(const std::string& name , const TLabels& labels @@ -169,6 +219,8 @@ class TOtelMetricRegistry : public IMetricRegistry { ConfigureHistogramBuckets(name, unit, buckets); auto histogram = Meter_->CreateDoubleHistogram(name, description, unit); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) std::shared_ptr Counter(const std::string& name, const TLabels& labels) override { auto counter = Meter_->CreateUInt64Counter(name); return std::make_shared(std::move(counter), labels); @@ -182,25 +234,36 @@ class TOtelMetricRegistry : public IMetricRegistry { std::shared_ptr Histogram(const std::string& name, const std::vector& buckets, const TLabels& labels) override { ConfigureHistogramBuckets(name, buckets); auto histogram = Meter_->CreateDoubleHistogram(name); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) return std::make_shared(std::move(histogram), labels); } private: +<<<<<<< HEAD <<<<<<< HEAD void ConfigureHistogramBuckets(const std::string& name, const std::string& unit, const std::vector& buckets) { ======= void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + void ConfigureHistogramBuckets(const std::string& name, const std::vector& buckets) { +>>>>>>> dcae6d69e (fixes and add metric tests) if (buckets.empty()) { return; } +<<<<<<< HEAD <<<<<<< HEAD auto* sdkProvider = dynamic_cast(MeterProvider_.get()); ======= auto* sdkProvider = dynamic_cast(MeterProvider_.get()); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + auto* sdkProvider = dynamic_cast(MeterProvider_.get()); +>>>>>>> dcae6d69e (fixes and add metric tests) if (!sdkProvider) { return; } @@ -212,6 +275,7 @@ class TOtelMetricRegistry : public IMetricRegistry { } } +<<<<<<< HEAD <<<<<<< HEAD auto selector = std::make_unique( otel_sdk_metrics::InstrumentType::kHistogram, @@ -232,6 +296,8 @@ class TOtelMetricRegistry : public IMetricRegistry { std::string(), otel_sdk_metrics::AggregationType::kHistogram, ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto selector = std::make_unique( sdk::metrics::InstrumentType::kHistogram, name, @@ -250,13 +316,17 @@ class TOtelMetricRegistry : public IMetricRegistry { {}, {}, sdk::metrics::AggregationType::kHistogram, +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) histogramConfig ); sdkProvider->AddView(std::move(selector), std::move(meterSelector), std::move(view)); } +<<<<<<< HEAD <<<<<<< HEAD otel_nostd::shared_ptr MeterProvider_; otel_nostd::shared_ptr Meter_; @@ -264,6 +334,10 @@ class TOtelMetricRegistry : public IMetricRegistry { nostd::shared_ptr MeterProvider_; nostd::shared_ptr Meter_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + nostd::shared_ptr MeterProvider_; + nostd::shared_ptr Meter_; +>>>>>>> dcae6d69e (fixes and add metric tests) std::mutex HistogramViewsLock_; std::unordered_set HistogramViews_; }; diff --git a/plugins/open_telemetry/CMakeLists.txt b/plugins/open_telemetry/CMakeLists.txt deleted file mode 100644 index d005708d422..00000000000 --- a/plugins/open_telemetry/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -if (YDB_SDK_ENABLE_OTEL_METRICS) - _ydb_sdk_add_library(open_telemetry_metrics) - target_sources(open_telemetry_metrics PRIVATE - src/metrics.cpp - ) - target_include_directories(open_telemetry_metrics PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_metrics PUBLIC - client-metrics - client-resources - opentelemetry-cpp::api - opentelemetry-cpp::metrics - ) - _ydb_sdk_make_client_component(OpenTelemetryMetrics open_telemetry_metrics) -endif() - -if (YDB_SDK_ENABLE_OTEL_TRACE) - _ydb_sdk_add_library(open_telemetry_trace) - target_sources(open_telemetry_trace PRIVATE - src/trace.cpp - ) - target_include_directories(open_telemetry_trace PUBLIC - $ - $ - ) - target_link_libraries(open_telemetry_trace PUBLIC - client-metrics - opentelemetry-cpp::api - opentelemetry-cpp::trace - ) - _ydb_sdk_make_client_component(OpenTelemetryTrace open_telemetry_trace) -endif() - -_ydb_sdk_install_headers(${CMAKE_INSTALL_INCLUDEDIR} DIRECTORY include/) diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h deleted file mode 100644 index b0db9ea7d7c..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/extension.h +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include -#include - -namespace NYdb::inline V3::NMetrics { - -inline void AddOpenTelemetry(TDriverConfig& config - , opentelemetry::nostd::shared_ptr meterProvider - , opentelemetry::nostd::shared_ptr tracerProvider -) { - AddOpenTelemetryMetrics(config, std::move(meterProvider)); - AddOpenTelemetryTrace(config, std::move(tracerProvider)); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h deleted file mode 100644 index 3ba2e146fd9..00000000000 --- a/plugins/open_telemetry/include/ydb-cpp-sdk/open_telemetry/trace.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -class TOtelTraceProvider : public ITraceProvider { -public: - TOtelTraceProvider(opentelemetry::nostd::shared_ptr tracerProvider); - - std::shared_ptr GetTracer(const std::string& name) override; - -private: - opentelemetry::nostd::shared_ptr TracerProvider_; -}; - -inline void AddOpenTelemetryTrace( - TDriverConfig& config, - opentelemetry::nostd::shared_ptr tracerProvider) -{ - if (tracerProvider) { - config.SetTraceExporter(std::make_shared(std::move(tracerProvider))); - } -} - -} // namespace NYdb::NMetrics diff --git a/plugins/open_telemetry/src/trace.cpp b/plugins/open_telemetry/src/trace.cpp deleted file mode 100644 index 54f04cb84df..00000000000 --- a/plugins/open_telemetry/src/trace.cpp +++ /dev/null @@ -1,70 +0,0 @@ -#include - -#include - -namespace NYdb::inline V3::NMetrics { - -namespace { - -using namespace opentelemetry; - -trace::SpanKind MapSpanKind(ESpanKind kind) { - switch (kind) { - case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; - case ESpanKind::SERVER: return trace::SpanKind::kServer; - case ESpanKind::CLIENT: return trace::SpanKind::kClient; - case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; - case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; - } - return trace::SpanKind::kInternal; -} - -class TOtelSpan : public ISpan { -public: - TOtelSpan(nostd::shared_ptr span) - : Span_(std::move(span)) - {} - - void End() override { - Span_->End(); - } - - void SetAttribute(const std::string& key, const std::string& value) override { - Span_->SetAttribute(key, value); - } - - void SetAttribute(const std::string& key, int64_t value) override { - Span_->SetAttribute(key, value); - } - -private: - nostd::shared_ptr Span_; -}; - -class TOtelTracer : public ITracer { -public: - TOtelTracer(nostd::shared_ptr tracer) - : Tracer_(std::move(tracer)) - {} - - std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { - trace::StartSpanOptions options; - options.kind = MapSpanKind(kind); - return std::make_shared(Tracer_->StartSpan(name, options)); - } - -private: - nostd::shared_ptr Tracer_; -}; - -} // namespace - -TOtelTraceProvider::TOtelTraceProvider(nostd::shared_ptr tracerProvider) - : TracerProvider_(std::move(tracerProvider)) -{} - -std::shared_ptr TOtelTraceProvider::GetTracer(const std::string& name) { - return std::make_shared(TracerProvider_->GetTracer(name)); -} - -} // namespace NYdb::NMetrics diff --git a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h index b4e9f09eab8..bc835141de4 100644 --- a/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h +++ b/plugins/trace/otel/include/ydb-cpp-sdk/open_telemetry/trace.h @@ -4,13 +4,19 @@ #include <<<<<<< HEAD +<<<<<<< HEAD #include ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) namespace opentelemetry::trace { class TracerProvider; } +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) namespace NYdb::inline V3::NMetrics { diff --git a/plugins/trace/otel/src/trace.cpp b/plugins/trace/otel/src/trace.cpp index cc2c2d0d294..f123243d8ee 100644 --- a/plugins/trace/otel/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -2,9 +2,12 @@ #include <<<<<<< HEAD +<<<<<<< HEAD #include ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) #include #include @@ -12,6 +15,7 @@ namespace NYdb::inline V3::NMetrics { namespace { +<<<<<<< HEAD <<<<<<< HEAD namespace otel_trace = opentelemetry::trace; namespace otel_nostd = opentelemetry::nostd; @@ -42,6 +46,8 @@ class TOtelSpan : public ISpan { public: TOtelSpan(otel_nostd::shared_ptr span) ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) using namespace opentelemetry; trace::SpanKind MapSpanKind(ESpanKind kind) { @@ -58,7 +64,10 @@ trace::SpanKind MapSpanKind(ESpanKind kind) { class TOtelSpan : public ISpan { public: TOtelSpan(nostd::shared_ptr span) +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) : Span_(std::move(span)) {} @@ -78,22 +87,29 @@ class TOtelSpan : public ISpan { if (attributes.empty()) { Span_->AddEvent(name); } else { +<<<<<<< HEAD <<<<<<< HEAD std::vector> attrs; attrs.reserve(attributes.size()); for (const auto& [k, v] : attributes) { attrs.emplace_back(otel_nostd::string_view(k), otel_common::AttributeValue(otel_nostd::string_view(v))); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) std::vector> attrs; attrs.reserve(attributes.size()); for (const auto& [k, v] : attributes) { attrs.emplace_back(nostd::string_view(k), common::AttributeValue(nostd::string_view(v))); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } Span_->AddEvent(name, attrs); } } +<<<<<<< HEAD <<<<<<< HEAD std::unique_ptr Activate() override { return std::make_unique(Span_); @@ -105,43 +121,63 @@ class TOtelSpan : public ISpan { private: nostd::shared_ptr Span_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +private: + nostd::shared_ptr Span_; +>>>>>>> dcae6d69e (fixes and add metric tests) }; class TOtelTracer : public ITracer { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelTracer(otel_nostd::shared_ptr tracer) ======= TOtelTracer(nostd::shared_ptr tracer) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelTracer(nostd::shared_ptr tracer) +>>>>>>> dcae6d69e (fixes and add metric tests) : Tracer_(std::move(tracer)) {} std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { +<<<<<<< HEAD <<<<<<< HEAD otel_trace::StartSpanOptions options; ======= trace::StartSpanOptions options; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + trace::StartSpanOptions options; +>>>>>>> dcae6d69e (fixes and add metric tests) options.kind = MapSpanKind(kind); return std::make_shared(Tracer_->StartSpan(name, options)); } private: +<<<<<<< HEAD <<<<<<< HEAD otel_nostd::shared_ptr Tracer_; ======= nostd::shared_ptr Tracer_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + nostd::shared_ptr Tracer_; +>>>>>>> dcae6d69e (fixes and add metric tests) }; class TOtelTraceProvider : public ITraceProvider { public: +<<<<<<< HEAD <<<<<<< HEAD TOtelTraceProvider(otel_nostd::shared_ptr tracerProvider) ======= TOtelTraceProvider(nostd::shared_ptr tracerProvider) >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + TOtelTraceProvider(nostd::shared_ptr tracerProvider) +>>>>>>> dcae6d69e (fixes and add metric tests) : TracerProvider_(std::move(tracerProvider)) {} @@ -150,11 +186,15 @@ class TOtelTraceProvider : public ITraceProvider { } private: +<<<<<<< HEAD <<<<<<< HEAD otel_nostd::shared_ptr TracerProvider_; ======= nostd::shared_ptr TracerProvider_; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + nostd::shared_ptr TracerProvider_; +>>>>>>> dcae6d69e (fixes and add metric tests) }; } // namespace diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index 397b79e7566..78312b18441 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -55,6 +55,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr GetExecutor() const override { return Executor; } <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } std::shared_ptr GetTraceProvider() const override { return TraceProvider; } @@ -70,6 +71,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr GetMetricExporter() const override { return MetricExporter; } std::shared_ptr GetTraceExporter() const override { return TraceExporter; } >>>>>>> a979e6bda (fixes) +======= + std::shared_ptr GetExternalMetricRegistry() const override { return MetricRegistry; } + std::shared_ptr GetTraceProvider() const override { return TraceProvider; } +>>>>>>> dcae6d69e (fixes and add metric tests) std::string Endpoint; size_t NetworkThreadsNum = 2; @@ -103,6 +108,7 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr Executor; <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry; std::shared_ptr TraceProvider; @@ -118,6 +124,10 @@ class TDriverConfig::TImpl : public IConnectionsParams { std::shared_ptr MetricExporter; std::shared_ptr TraceExporter; >>>>>>> a979e6bda (fixes) +======= + std::shared_ptr MetricRegistry; + std::shared_ptr TraceProvider; +>>>>>>> dcae6d69e (fixes and add metric tests) }; TDriverConfig::TDriverConfig(const std::string& connectionString) @@ -280,6 +290,7 @@ TDriverConfig& TDriverConfig::SetExecutor(std::shared_ptr executor) { <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { Impl_->MetricRegistry = std::move(registry); return *this; @@ -309,14 +320,19 @@ TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr exporter) { Impl_->MetricExporter = std::move(exporter); +======= +TDriverConfig& TDriverConfig::SetMetricRegistry(std::shared_ptr registry) { + Impl_->MetricRegistry = std::move(registry); +>>>>>>> dcae6d69e (fixes and add metric tests) return *this; } -TDriverConfig& TDriverConfig::SetTraceExporter(std::shared_ptr exporter) { - Impl_->TraceExporter = std::move(exporter); +TDriverConfig& TDriverConfig::SetTraceProvider(std::shared_ptr provider) { + Impl_->TraceProvider = std::move(provider); return *this; } +<<<<<<< HEAD >>>>>>> a979e6bda (fixes) std::shared_ptr TDriverConfig::GetMetricExporter() const { return Impl_->MetricExporter; @@ -332,6 +348,8 @@ std::shared_ptr TDriverConfig::GetTraceExporter() cons >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) //////////////////////////////////////////////////////////////////////////////// std::shared_ptr CreateInternalInterface(const TDriver connection) { @@ -387,6 +405,7 @@ TDriverConfig TDriver::GetConfig() const { config.Impl_->Log = Impl_->Log; <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); config.SetTraceProvider(Impl_->GetTraceProvider()); @@ -402,6 +421,10 @@ TDriverConfig TDriver::GetConfig() const { config.SetMetricExporter(Impl_->GetMetricExporter()); config.SetTraceExporter(Impl_->GetTraceExporter()); >>>>>>> a979e6bda (fixes) +======= + config.SetMetricRegistry(Impl_->GetExternalMetricRegistry()); + config.SetTraceProvider(Impl_->GetTraceProvider()); +>>>>>>> dcae6d69e (fixes and add metric tests) return config; } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index e766741e152..86d4dbb7f81 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -171,6 +171,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p #endif <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD , MetricRegistry_(params->GetExternalMetricRegistry()) , TraceProvider_(params->GetTraceProvider()) @@ -186,6 +187,10 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MetricExporter_(params->GetMetricExporter()) , TraceExporter_(params->GetTraceExporter()) >>>>>>> a979e6bda (fixes) +======= + , MetricRegistry_(params->GetExternalMetricRegistry()) + , TraceProvider_(params->GetTraceProvider()) +>>>>>>> dcae6d69e (fixes and add metric tests) , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) , GRpcClientLow_(NetworkThreadsNum_) @@ -456,6 +461,7 @@ void TGRpcConnectionsImpl::RegisterExtensionApi(IExtensionApi* api) { <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { return MetricRegistry_; } @@ -483,6 +489,14 @@ std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= +std::shared_ptr TGRpcConnectionsImpl::GetExternalMetricRegistry() const { + return MetricRegistry_; +} + +std::shared_ptr TGRpcConnectionsImpl::GetTraceProvider() const { + return TraceProvider_; +>>>>>>> dcae6d69e (fixes and add metric tests) } void TGRpcConnectionsImpl::SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb) { diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 3b57d085f1f..0b041699e4b 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -588,6 +588,7 @@ class TGRpcConnectionsImpl void RegisterExtensionApi(IExtensionApi* api); <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr GetExternalMetricRegistry() const; std::shared_ptr GetTraceProvider() const; @@ -612,6 +613,10 @@ class TGRpcConnectionsImpl std::shared_ptr GetExternalMetricRegistry() const; std::shared_ptr GetTraceProvider() const; >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr GetExternalMetricRegistry() const; + std::shared_ptr GetTraceProvider() const; +>>>>>>> dcae6d69e (fixes and add metric tests) void SetDiscoveryMutator(IDiscoveryMutatorApi::TMutatorCb&& cb); const TLog& GetLog() const override; @@ -750,6 +755,7 @@ class TGRpcConnectionsImpl std::vector> ExtensionApis_; <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr MetricRegistry_; std::shared_ptr TraceProvider_; @@ -765,6 +771,10 @@ class TGRpcConnectionsImpl std::shared_ptr MetricExporter_; std::shared_ptr TraceExporter_; >>>>>>> a979e6bda (fixes) +======= + std::shared_ptr MetricRegistry_; + std::shared_ptr TraceProvider_; +>>>>>>> dcae6d69e (fixes and add metric tests) IDiscoveryMutatorApi::TMutatorCb DiscoveryMutatorCb; diff --git a/src/client/impl/internal/grpc_connections/params.h b/src/client/impl/internal/grpc_connections/params.h index 30554cff67d..b0bdfd628d5 100644 --- a/src/client/impl/internal/grpc_connections/params.h +++ b/src/client/impl/internal/grpc_connections/params.h @@ -45,6 +45,7 @@ class IConnectionsParams { virtual std::shared_ptr GetExecutor() const = 0; <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD virtual std::shared_ptr GetExternalMetricRegistry() const = 0; virtual std::shared_ptr GetTraceProvider() const = 0; @@ -60,6 +61,10 @@ class IConnectionsParams { virtual std::shared_ptr GetMetricExporter() const = 0; virtual std::shared_ptr GetTraceExporter() const = 0; >>>>>>> a979e6bda (fixes) +======= + virtual std::shared_ptr GetExternalMetricRegistry() const = 0; + virtual std::shared_ptr GetTraceProvider() const = 0; +>>>>>>> dcae6d69e (fixes and add metric tests) }; } // namespace NYdb diff --git a/src/client/metrics/CMakeLists.txt b/src/client/metrics/CMakeLists.txt index ef4841c5026..ad9c06dac3e 100644 --- a/src/client/metrics/CMakeLists.txt +++ b/src/client/metrics/CMakeLists.txt @@ -2,6 +2,7 @@ _ydb_sdk_add_library(client-metrics) target_sources(client-metrics PRIVATE <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD metrics.cpp ======= @@ -13,6 +14,9 @@ target_sources(client-metrics PRIVATE >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= + metrics.cpp +>>>>>>> dcae6d69e (fixes and add metric tests) ) _ydb_sdk_make_client_component(Metrics client-metrics) diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index b631534e6a1..819d3b04bdb 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -73,6 +73,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if (auto traceProvider = Connections_->GetTraceProvider()) { ======= @@ -84,6 +85,9 @@ class TQueryClient::TImpl: public TClientImplCommon, public ======= if (auto traceProvider = Connections_->GetTraceExporter()) { >>>>>>> a979e6bda (fixes) +======= + if (auto traceProvider = Connections_->GetTraceProvider()) { +>>>>>>> dcae6d69e (fixes and add metric tests) Tracer_ = traceProvider->GetTracer("ydb-cpp-sdk-query"); } MetricRegistry_ = Connections_->GetExternalMetricRegistry(); @@ -117,9 +121,13 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto span = std::make_shared(Tracer_, "ExecuteQuery", DbDriverState_->DiscoveryEndpoint); <<<<<<< HEAD +<<<<<<< HEAD ======= span->SetQueryText(query); >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + span->SetQueryText(query); +>>>>>>> dcae6d69e (fixes and add metric tests) auto metrics = std::make_shared(MetricRegistry_, "ExecuteQuery"); return TExecQueryImpl::ExecuteQuery( @@ -508,6 +516,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public } <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(status.GetStatus()); @@ -521,6 +530,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= + if (Metrics) { + Metrics->End(status.GetStatus()); + } +>>>>>>> dcae6d69e (fixes and add metric tests) ScheduleReply(TCreateSessionResult(std::move(status), std::move(session))); } @@ -538,6 +552,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public } <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if (Metrics) { Metrics->End(EStatus::SUCCESS); @@ -551,6 +566,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= + if (Metrics) { + Metrics->End(EStatus::SUCCESS); + } +>>>>>>> dcae6d69e (fixes and add metric tests) ScheduleReply(std::move(val)); } @@ -563,8 +583,11 @@ class TQueryClient::TImpl: public TClientImplCommon, public <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) span->SetPeerEndpoint(val.GetEndpoint()); span->End(val.GetStatus()); } @@ -572,6 +595,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public metrics->End(val.GetStatus()); } <<<<<<< HEAD +<<<<<<< HEAD ======= span->End(val.GetStatus()); } @@ -582,6 +606,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public span->End(val.GetStatus()); } >>>>>>> a979e6bda (fixes) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) promise.SetValue(std::move(val)); }); } diff --git a/src/client/query/impl/query_metrics.h b/src/client/query/impl/query_metrics.h index e62d7e014c2..c2b99e9d0db 100644 --- a/src/client/query/impl/query_metrics.h +++ b/src/client/query/impl/query_metrics.h @@ -1,5 +1,6 @@ #pragma once +<<<<<<< HEAD <<<<<<< HEAD #include @@ -11,6 +12,8 @@ class TQueryMetrics : public NObservability::TClientMetrics { : TClientMetrics(std::move(registry), operationName) {} ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) #include #include @@ -34,7 +37,10 @@ class TQueryMetrics { std::shared_ptr LatencyHistogram_; TInstant StartTime_; bool Ended_ = false; +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) }; } // namespace NYdb::NQuery diff --git a/src/client/query/impl/query_spans.cpp b/src/client/query/impl/query_spans.cpp index 9facc7c02dd..81a58fbd15b 100644 --- a/src/client/query/impl/query_spans.cpp +++ b/src/client/query/impl/query_spans.cpp @@ -38,6 +38,7 @@ void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { } catch (...) {} } else { host = endpoint; +<<<<<<< HEAD <<<<<<< HEAD } } @@ -65,17 +66,19 @@ void SafeLogSpanError(const char* message) noexcept { } std::cerr << "TQuerySpan: " << message << ": (unknown)" << std::endl; } catch (...) { +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } } void SafeLogSpanError(const char* message) noexcept { try { try { - Cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << Endl; + std::cerr << "TQuerySpan: " << message << ": " << CurrentExceptionMessage() << std::endl; return; } catch (...) { } - Cerr << "TQuerySpan: " << message << ": (unknown)" << Endl; + std::cerr << "TQuerySpan: " << message << ": (unknown)" << std::endl; } catch (...) { } } @@ -133,8 +136,11 @@ TQuerySpan::~TQuerySpan() noexcept { <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { if (!Span_ || endpoint.empty()) { return; @@ -151,7 +157,10 @@ void TQuerySpan::SetPeerEndpoint(const std::string& endpoint) noexcept { } <<<<<<< HEAD +<<<<<<< HEAD +======= ======= +>>>>>>> dcae6d69e (fixes and add metric tests) void TQuerySpan::SetQueryText(const std::string& query) noexcept { if (!Span_ || query.empty()) { return; @@ -163,7 +172,10 @@ void TQuerySpan::SetQueryText(const std::string& query) noexcept { } } +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) void TQuerySpan::AddEvent(const std::string& name, const std::map& attributes) noexcept { if (!Span_) { return; @@ -175,6 +187,7 @@ void TQuerySpan::AddEvent(const std::string& name, const std::map TQuerySpan::Activate() noexcept { if (!Span_) { @@ -197,6 +210,8 @@ void TQuerySpan::End(EStatus status) noexcept { >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) void TQuerySpan::End(EStatus status) noexcept { if (Span_) { try { diff --git a/src/client/query/impl/query_spans.h b/src/client/query/impl/query_spans.h index 3c8d6443b2e..eaa2f3158bb 100644 --- a/src/client/query/impl/query_spans.h +++ b/src/client/query/impl/query_spans.h @@ -17,6 +17,7 @@ class TQuerySpan { <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD void SetPeerEndpoint(const std::string& endpoint) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; @@ -25,13 +26,18 @@ class TQuerySpan { ======= >>>>>>> 1b2bf4fa5 (fixes) ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) void SetPeerEndpoint(const std::string& endpoint) noexcept; void SetQueryText(const std::string& query) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) ======= >>>>>>> a979e6bda (fixes) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) void End(EStatus status) noexcept; private: diff --git a/tests/common/fake_metric_registry.h b/tests/common/fake_metric_registry.h index 5f8f10572b4..859ee264ac1 100644 --- a/tests/common/fake_metric_registry.h +++ b/tests/common/fake_metric_registry.h @@ -68,6 +68,7 @@ struct TMetricKey { class TFakeMetricRegistry : public NMetrics::IMetricRegistry { public: +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr Counter(const std::string& name , const NMetrics::TLabels& labels @@ -77,6 +78,9 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { ======= std::shared_ptr Counter(const std::string& name, const NMetrics::TLabels& labels) override { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr Counter(const std::string& name, const NMetrics::TLabels& labels) override { +>>>>>>> dcae6d69e (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Counters_.find(key); @@ -88,6 +92,7 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return counter; } +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr Gauge(const std::string& name , const NMetrics::TLabels& labels @@ -97,6 +102,9 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { ======= std::shared_ptr Gauge(const std::string& name, const NMetrics::TLabels& labels) override { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr Gauge(const std::string& name, const NMetrics::TLabels& labels) override { +>>>>>>> dcae6d69e (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto gauge = std::make_shared(); @@ -104,6 +112,7 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { return gauge; } +<<<<<<< HEAD <<<<<<< HEAD std::shared_ptr Histogram(const std::string& name , const std::vector& /*buckets*/ @@ -114,6 +123,9 @@ class TFakeMetricRegistry : public NMetrics::IMetricRegistry { ======= std::shared_ptr Histogram(const std::string& name, const std::vector& /*buckets*/, const NMetrics::TLabels& labels) override { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= + std::shared_ptr Histogram(const std::string& name, const std::vector& /*buckets*/, const NMetrics::TLabels& labels) override { +>>>>>>> dcae6d69e (fixes and add metric tests) std::lock_guard lock(Mutex_); auto key = TMetricKey{name, labels}; auto it = Histograms_.find(key); diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index e36366c6fb7..75f6572d1ce 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -2,9 +2,12 @@ #include #include <<<<<<< HEAD +<<<<<<< HEAD #include ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) #include @@ -40,6 +43,7 @@ std::shared_ptr GetCounter( const std::string& name, const std::string& operation) { +<<<<<<< HEAD <<<<<<< HEAD return registry->GetCounter(name, { {"db.system.name", "other_sql"}, @@ -62,6 +66,8 @@ std::shared_ptr GetDuration( } return registry->GetHistogram("db.client.operation.duration", labels); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) return registry->GetCounter(name, {{"operation", operation}}); } @@ -71,7 +77,10 @@ std::shared_ptr GetHistogram( const std::string& operation) { return registry->GetHistogram(name, {{"operation", operation}}); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } } // namespace @@ -89,6 +98,7 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); +<<<<<<< HEAD <<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; @@ -103,6 +113,8 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { EXPECT_GE(duration->Count(), 1u); for (double v : duration->GetValues()) { ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; EXPECT_GE(requests->Get(), 1); @@ -115,7 +127,10 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ASSERT_NE(latency, nullptr) << "ExecuteQuery latency histogram not created"; EXPECT_GE(latency->Count(), 1u); for (double v : latency->GetValues()) { +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) EXPECT_GE(v, 0.0); } @@ -135,6 +150,7 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ).ExtractValueSync(); EXPECT_NE(result.GetStatus(), EStatus::SUCCESS); +<<<<<<< HEAD <<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); @@ -148,6 +164,8 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ASSERT_NE(duration, nullptr); EXPECT_GE(duration->Count(), 1u); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_GE(requests->Get(), 1); @@ -159,7 +177,10 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); ASSERT_NE(latency, nullptr); EXPECT_GE(latency->Count(), 1u); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) driver.Stop(true); } @@ -171,6 +192,7 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto session = client.GetSession().ExtractValueSync(); ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); +<<<<<<< HEAD <<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "GetSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; @@ -180,6 +202,8 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; EXPECT_GE(duration->Count(), 1u); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto requests = GetCounter(registry, "ydb.query.requests", "CreateSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); @@ -187,7 +211,10 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto latency = GetHistogram(registry, "ydb.query.latency_ms", "CreateSession"); ASSERT_NE(latency, nullptr) << "CreateSession latency histogram not created"; EXPECT_GE(latency->Count(), 1u); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) driver.Stop(true); } @@ -214,6 +241,7 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitResult = execResult.GetTransaction()->Commit().ExtractValueSync(); ASSERT_TRUE(commitResult.IsSuccess()) << commitResult.GetIssues().ToString(); +<<<<<<< HEAD <<<<<<< HEAD auto commitRequests = GetCounter(registry, "db.client.operation.requests", "Commit"); ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; @@ -223,6 +251,8 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { ASSERT_NE(commitDuration, nullptr); EXPECT_GE(commitDuration->Count(), 1u); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto commitRequests = GetCounter(registry, "ydb.query.requests", "Commit"); ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; EXPECT_GE(commitRequests->Get(), 1); @@ -230,7 +260,10 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitLatency = GetHistogram(registry, "ydb.query.latency_ms", "Commit"); ASSERT_NE(commitLatency, nullptr); EXPECT_GE(commitLatency->Count(), 1u); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } driver.Stop(true); @@ -251,6 +284,7 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackResult = tx.Rollback().ExtractValueSync(); ASSERT_TRUE(rollbackResult.IsSuccess()) << rollbackResult.GetIssues().ToString(); +<<<<<<< HEAD <<<<<<< HEAD auto rollbackRequests = GetCounter(registry, "db.client.operation.requests", "Rollback"); ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; @@ -264,6 +298,8 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { ASSERT_NE(rollbackDuration, nullptr); EXPECT_GE(rollbackDuration->Count(), 1u); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto rollbackRequests = GetCounter(registry, "ydb.query.requests", "Rollback"); ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; EXPECT_GE(rollbackRequests->Get(), 1); @@ -275,7 +311,10 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackLatency = GetHistogram(registry, "ydb.query.latency_ms", "Rollback"); ASSERT_NE(rollbackLatency, nullptr); EXPECT_GE(rollbackLatency->Count(), 1u); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) driver.Stop(true); } @@ -297,6 +336,7 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); } +<<<<<<< HEAD <<<<<<< HEAD auto requests = GetCounter(registry, "db.client.operation.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); @@ -310,6 +350,8 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_NE(duration, nullptr); EXPECT_EQ(duration->Count(), static_cast(numQueries)); ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto requests = GetCounter(registry, "ydb.query.requests", "ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_EQ(requests->Get(), numQueries); @@ -321,7 +363,10 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); ASSERT_NE(latency, nullptr); EXPECT_EQ(latency->Count(), static_cast(numQueries)); +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) driver.Stop(true); } @@ -350,11 +395,15 @@ TEST(QueryMetricsIntegration, NoRegistryDoesNotBreakOperations) { driver.Stop(true); } +<<<<<<< HEAD <<<<<<< HEAD TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { ======= TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { +>>>>>>> dcae6d69e (fixes and add metric tests) auto [driver, registry] = MakeRunArgs(); TQueryClient client(driver); @@ -368,6 +417,7 @@ TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); +<<<<<<< HEAD <<<<<<< HEAD auto duration = GetDuration(registry, "ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); @@ -377,6 +427,8 @@ TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { EXPECT_GE(v, 0.0) << "Duration must be non-negative"; EXPECT_LT(v, 30.0) << "Duration > 30s is unrealistic for SELECT 1"; ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) auto latency = GetHistogram(registry, "ydb.query.latency_ms", "ExecuteQuery"); ASSERT_NE(latency, nullptr); ASSERT_GE(latency->Count(), 1u); @@ -384,7 +436,10 @@ TEST(QueryMetricsIntegration, LatencyValuesAreRealistic) { for (double v : latency->GetValues()) { EXPECT_GE(v, 0.0) << "Latency must be non-negative"; EXPECT_LT(v, 30000.0) << "Latency > 30s is unrealistic for SELECT 1"; +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) } driver.Stop(true); diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index 99a3387c250..3a3cb29f874 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -101,6 +101,7 @@ add_ydb_test(NAME client-ydb_value_ut GTEST unit ) +<<<<<<< HEAD <<<<<<< HEAD add_ydb_test(NAME client-ydb_metrics_ut GTEST INCLUDE_DIRS @@ -113,6 +114,8 @@ add_ydb_test(NAME client-ydb_metrics_ut GTEST client-ydb_query-impl client-ydb_table-impl ======= +======= +>>>>>>> dcae6d69e (fixes and add metric tests) add_ydb_test(NAME client-ydb_query_metrics_ut GTEST INCLUDE_DIRS ${YDB_SDK_SOURCE_DIR} @@ -121,12 +124,16 @@ add_ydb_test(NAME client-ydb_query_metrics_ut GTEST LINK_LIBRARIES yutil client-ydb_query-impl +<<<<<<< HEAD >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) client-metrics LABELS unit ) <<<<<<< HEAD +<<<<<<< HEAD add_ydb_test(NAME client-ydb_query_spans_ut GTEST INCLUDE_DIRS @@ -142,3 +149,5 @@ add_ydb_test(NAME client-ydb_query_spans_ut GTEST ) ======= >>>>>>> 1ca4253b5 (fixes and add metric tests) +======= +>>>>>>> dcae6d69e (fixes and add metric tests) From a3fc3110d9531548c772dd4155a66156dfc2a85a Mon Sep 17 00:00:00 2001 From: qyryq Date: Thu, 12 Feb 2026 13:39:24 +0000 Subject: [PATCH 87/93] Enable direct read without consumer (LOGBROKER-9364) (#32846) --- .github/last_commit.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index d769cd9b98e..bcd5f5861c2 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -fd4d02d6b3ebb3565fec58f5a585b44e09e6899e +fd4d02d6b3ebb3565fec58f5a585b44e09e6899 \ No newline at end of file From f5e0a28ade8d99ab599bb5fdb63f15de5a1c65df Mon Sep 17 00:00:00 2001 From: qyryq Date: Thu, 12 Feb 2026 13:39:30 +0000 Subject: [PATCH 88/93] Add TSimpleBlockingFederatedWriteSession (#32794) --- src/client/topic/impl/write_session.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index bd103aee855..c0b7c5331e8 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -2,6 +2,7 @@ #include #include +<<<<<<< HEAD #include #include @@ -12,6 +13,8 @@ #include #include +======= +>>>>>>> f3e35e369 (Add TSimpleBlockingFederatedWriteSession (#32794)) namespace NYdb::inline V3::NTopic { From 591fc554e98bbe17259a6aded932664c6111d50d Mon Sep 17 00:00:00 2001 From: azevaykin <145343289+azevaykin@users.noreply.github.com> Date: Thu, 12 Feb 2026 13:39:37 +0000 Subject: [PATCH 89/93] Access levels in whoami (#32907) --- .github/last_commit.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/last_commit.txt b/.github/last_commit.txt index bcd5f5861c2..62b939ab8b3 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -fd4d02d6b3ebb3565fec58f5a585b44e09e6899 \ No newline at end of file +fd4d02d6b3ebb3565fec58f5a585b44e09e6899e \ No newline at end of file From 6e7e43531b22c0e05e6312de1a36d98e315d8f23 Mon Sep 17 00:00:00 2001 From: maladetska Date: Tue, 28 Apr 2026 02:05:31 +0300 Subject: [PATCH 90/93] fix test --- examples/otel_tracing/main.cpp | 2 +- src/client/table/impl/CMakeLists.txt | 1 - src/client/table/impl/table_spans.cpp | 91 ------------------------- src/client/table/impl/table_spans.h | 24 ------- src/client/topic/impl/write_session.cpp | 13 ---- 5 files changed, 1 insertion(+), 130 deletions(-) delete mode 100644 src/client/table/impl/table_spans.cpp delete mode 100644 src/client/table/impl/table_spans.h diff --git a/examples/otel_tracing/main.cpp b/examples/otel_tracing/main.cpp index b06c224c2ca..3551836abd2 100644 --- a/examples/otel_tracing/main.cpp +++ b/examples/otel_tracing/main.cpp @@ -382,7 +382,7 @@ int main(int argc, char** argv) { auto tracerProvider = InitTracing(cfg); auto meterProvider = InitMetrics(cfg); - auto ydbTraceProvider = NMetrics::CreateOtelTraceProvider(tracerProvider); + auto ydbTraceProvider = NTrace::CreateOtelTraceProvider(tracerProvider); auto ydbMetricRegistry = NMetrics::CreateOtelMetricRegistry(meterProvider); std::cout << "Connecting to YDB at " << cfg.Endpoint << cfg.Database << std::endl; diff --git a/src/client/table/impl/CMakeLists.txt b/src/client/table/impl/CMakeLists.txt index 8e14331f062..8ecfe4ead87 100644 --- a/src/client/table/impl/CMakeLists.txt +++ b/src/client/table/impl/CMakeLists.txt @@ -22,7 +22,6 @@ target_sources(client-ydb_table-impl PRIVATE readers.cpp request_migrator.cpp table_client.cpp - table_spans.cpp transaction.cpp ) diff --git a/src/client/table/impl/table_spans.cpp b/src/client/table/impl/table_spans.cpp deleted file mode 100644 index 2593ef2a0b3..00000000000 --- a/src/client/table/impl/table_spans.cpp +++ /dev/null @@ -1,91 +0,0 @@ -#include "table_spans.h" - -#include - -namespace NYdb::inline V3::NTable { - -namespace { - -constexpr int DefaultGrpcPort = 2135; - -void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { - port = DefaultGrpcPort; - - if (endpoint.empty()) { - host = endpoint; - return; - } - - if (endpoint.front() == '[') { - auto bracketEnd = endpoint.find(']'); - if (bracketEnd != std::string::npos) { - host = endpoint.substr(1, bracketEnd - 1); - if (bracketEnd + 2 < endpoint.size() && endpoint[bracketEnd + 1] == ':') { - try { - port = std::stoi(endpoint.substr(bracketEnd + 2)); - } catch (...) {} - } - return; - } - } - - auto pos = endpoint.rfind(':'); - if (pos != std::string::npos) { - host = endpoint.substr(0, pos); - try { - port = std::stoi(endpoint.substr(pos + 1)); - } catch (...) {} - } else { - host = endpoint; - } -} - -} // namespace - -TTableSpan::TTableSpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint) { - if (!tracer) { - return; - } - - std::string host; - int port; - ParseEndpoint(endpoint, host, port); - - try { - Span_ = tracer->StartSpan(operationName, NMetrics::ESpanKind::CLIENT); - if (!Span_) { - return; - } - Span_->SetAttribute("db.system.name", "other_sql"); - Span_->SetAttribute("db.operation.name", operationName); - Span_->SetAttribute("server.address", host); - Span_->SetAttribute("server.port", static_cast(port)); - } catch (...) { - Span_.reset(); - } -} - -TTableSpan::~TTableSpan() noexcept { - if (Span_) { - try { - Span_->End(); - } catch (...) { - } - } -} - -void TTableSpan::End(EStatus status) noexcept { - if (Span_) { - try { - Span_->SetAttribute("db.response.status_code", ToString(status)); - if (status != EStatus::SUCCESS) { - Span_->SetAttribute("error.type", ToString(status)); - } - Span_->End(); - } catch (...) { - } - Span_.reset(); - } -} - -} // namespace NYdb::NTable diff --git a/src/client/table/impl/table_spans.h b/src/client/table/impl/table_spans.h deleted file mode 100644 index 283249b782e..00000000000 --- a/src/client/table/impl/table_spans.h +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include -#include -#include - -#include -#include -#include - -namespace NYdb::inline V3::NTable { - -class TTableSpan { -public: - TTableSpan(std::shared_ptr tracer, const std::string& operationName, const std::string& endpoint); - ~TTableSpan() noexcept; - - void End(EStatus status) noexcept; - -private: - std::shared_ptr Span_; -}; - -} // namespace NYdb::NTable diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index c0b7c5331e8..5df5c760912 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -2,19 +2,6 @@ #include #include -<<<<<<< HEAD -#include - -#include -#include - -#include -#include -#include - -#include -======= ->>>>>>> f3e35e369 (Add TSimpleBlockingFederatedWriteSession (#32794)) namespace NYdb::inline V3::NTopic { From bf693d734a068dc4a2df6cafb867046a3b55a465 Mon Sep 17 00:00:00 2001 From: maladetska Date: Tue, 28 Apr 2026 02:37:18 +0300 Subject: [PATCH 91/93] fix test --- src/client/impl/internal/common/client_pid.cpp | 5 +++++ src/client/topic/impl/write_session.cpp | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/src/client/impl/internal/common/client_pid.cpp b/src/client/impl/internal/common/client_pid.cpp index cb5cc8eb0a5..76347f3f740 100644 --- a/src/client/impl/internal/common/client_pid.cpp +++ b/src/client/impl/internal/common/client_pid.cpp @@ -3,6 +3,11 @@ #include +#include + +#include +#include + #ifdef _win_ // copied from util/system/getpid.cpp // to avoid extra util dep. diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 5df5c760912..bd103aee855 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -2,6 +2,16 @@ #include #include +#include + +#include +#include + +#include +#include +#include + +#include namespace NYdb::inline V3::NTopic { From dc1df0a44a130ff21e91aa4b8e53cb68aa02dbff Mon Sep 17 00:00:00 2001 From: maladetska Date: Wed, 29 Apr 2026 01:47:29 +0300 Subject: [PATCH 92/93] update --- examples/otel_tracing/README.md | 34 +- examples/otel_tracing/docker-compose.yml | 5 +- examples/otel_tracing/main.cpp | 10 + include/ydb-cpp-sdk/client/trace/trace.h | 21 +- plugins/trace/otel/src/trace.cpp | 77 ++- src/client/impl/internal/retry/CMakeLists.txt | 1 + src/client/impl/internal/retry/retry.cpp | 12 +- src/client/impl/internal/retry/retry.h | 4 +- src/client/impl/internal/retry/retry_async.h | 74 ++- src/client/impl/internal/retry/retry_sync.h | 90 +++- src/client/impl/observability/CMakeLists.txt | 6 + .../error_category/CMakeLists.txt | 11 + .../error_category/error_category.cpp | 13 + .../error_category/error_category.h | 14 + src/client/impl/observability/metrics.cpp | 6 +- src/client/impl/observability/observation.cpp | 28 +- src/client/impl/observability/observation.h | 7 +- .../impl/observability/operation_name.cpp | 17 + .../impl/observability/operation_name.h | 9 + src/client/impl/observability/span.cpp | 153 +++++- src/client/impl/observability/span.h | 51 +- src/client/impl/stats/CMakeLists.txt | 1 + src/client/impl/stats/stats.h | 5 +- src/client/query/client.cpp | 75 ++- src/client/table/impl/table_client.cpp | 43 +- src/client/table/impl/table_client.h | 15 +- src/client/table/table.cpp | 27 +- tests/common/fake_trace_provider.h | 163 ++++++ tests/integration/metrics/CMakeLists.txt | 1 + tests/integration/metrics/main.cpp | 39 +- tests/unit/client/CMakeLists.txt | 16 + .../unit/client/observability/metrics_ut.cpp | 48 +- tests/unit/client/observability/spans_ut.cpp | 498 ++++++++++++++++++ 33 files changed, 1373 insertions(+), 201 deletions(-) create mode 100644 src/client/impl/observability/error_category/CMakeLists.txt create mode 100644 src/client/impl/observability/error_category/error_category.cpp create mode 100644 src/client/impl/observability/error_category/error_category.h create mode 100644 src/client/impl/observability/operation_name.cpp create mode 100644 src/client/impl/observability/operation_name.h create mode 100644 tests/common/fake_trace_provider.h create mode 100644 tests/unit/client/observability/spans_ut.cpp diff --git a/examples/otel_tracing/README.md b/examples/otel_tracing/README.md index f3def443d96..58f60fe0c8e 100644 --- a/examples/otel_tracing/README.md +++ b/examples/otel_tracing/README.md @@ -61,7 +61,7 @@ cmake --build . --target otel_tracing_example -j$(nproc) ```bash ./examples/otel_tracing/otel_tracing_example \ - --endpoint grpc://localhost:2136 \ + --endpoint localhost:2136 \ --database /local \ --otlp http://localhost:4328 \ --iterations 20 \ @@ -91,15 +91,15 @@ cmake --build . --target otel_tracing_example -j$(nproc) В трейсах появятся: ``` -ydb.RunWithRetry (INTERNAL) -├── ydb.Try attempt=0 (INTERNAL, backoff_ms=0) +ydb.RunWithRetry (INTERNAL, ydb.retry.count=N) +├── ydb.Try (INTERNAL) # первая попытка: ydb.retry.attempt и ydb.retry.backoff_ms отсутствуют │ ├── ydb.CreateSession │ ├── ydb.ExecuteQuery -│ └── ydb.Commit status=ABORTED, error.type=ABORTED, exception event -├── ydb.Try attempt=1 (INTERNAL, backoff_ms=...) -│ └── ... status=ABORTED -└── ydb.Try attempt=N (INTERNAL) - └── ... status=SUCCESS +│ └── ydb.Commit db.response.status_code=ABORTED, error.type=ydb_error, exception event +├── ydb.Try ydb.retry.attempt=1 (INTERNAL, ydb.retry.backoff_ms=...) +│ └── ... db.response.status_code=ABORTED, error.type=ydb_error +└── ydb.Try ydb.retry.attempt=N (INTERNAL, ydb.retry.backoff_ms=...) + └── ... db.response.status_code=SUCCESS ``` Для усиления конфликтов поднимите воркеров и операций: @@ -152,9 +152,14 @@ ydb.RunWithRetry (INTERNAL) `ydb.CreateSession`, `ydb.ExecuteQuery`, `ydb.ExecuteDataQuery`, `ydb.BeginTransaction`, `ydb.Commit`, `ydb.Rollback`, `ydb.ExecuteSchemeQuery`, `ydb.BulkUpsert`. -- Retry-спаны (`SpanKind = INTERNAL`): `ydb.RunWithRetry`, - `ydb.Try` (по одному на каждую попытку, с атрибутами - `ydb.retry.attempt`, `ydb.retry.backoff_ms`). +- Retry-спаны (`SpanKind = INTERNAL`): + - `ydb.RunWithRetry` — обёртка над всей retryable-логикой. + При фактических повторах содержит атрибут `ydb.retry.count` (общее число + выполненных повторов, `>= 1`). + - `ydb.Try` — по одному на каждую попытку. На retry-попытках содержит + атрибуты `ydb.retry.attempt` (`1..N`) и `ydb.retry.backoff_ms` + (длительность sleep перед этой попыткой). На первой (не retry) попытке + эти атрибуты не выставляются. - Общие атрибуты на всех YDB-спанах: - `db.system.name = ydb` - `db.namespace` (имя базы) @@ -162,7 +167,8 @@ ydb.RunWithRetry (INTERNAL) - `network.peer.address`, `network.peer.port` (фактический узел кластера) - На ошибках добавляются: - `db.response.status_code` — строковый статус YDB (например, `ABORTED`) - - `error.type` — тот же строковый статус + - `error.type` — категория источника ошибки: `ydb_error` (ошибка, + возвращённая YDB) или `transport_error` (ошибка транспортного уровня) - событие `exception` с `exception.type` и `exception.message` #### В Prometheus: @@ -170,7 +176,9 @@ ydb.RunWithRetry (INTERNAL) (OTel Semantic Conventions). Лейблы: `db.system.name`, `db.namespace`, `db.operation.name` (с префиксом `ydb.`), `ydb.client.api` (`Query` / `Table`). Для ошибок добавляются `db.response.status_code` - и `error.type`. + (точный YDB-статус, например `ABORTED`) и `error.type` — + низкокардинальная категория источника ошибки: `ydb_error` (статусы YDB-сервера) + или `transport_error` (клиентские/транспортные статусы). - `db_client_operation_requests_total` — счётчик начатых операций (включая каждую попытку ретрая). - `db_client_operation_errors_total` — счётчик неуспешных попыток. diff --git a/examples/otel_tracing/docker-compose.yml b/examples/otel_tracing/docker-compose.yml index 9d01c8fa823..8af4a6194c9 100644 --- a/examples/otel_tracing/docker-compose.yml +++ b/examples/otel_tracing/docker-compose.yml @@ -2,6 +2,7 @@ services: ydb: image: cr.yandex/yc/yandex-docker-local-ydb:latest platform: linux/amd64 + hostname: localhost ports: - "2136:2136" - "8765:8765" @@ -38,7 +39,7 @@ services: - otel-collector otel-collector: - image: otel/opentelemetry-collector-contrib:latest + image: otel/opentelemetry-collector-contrib:0.110.0 ports: - "4327:4317" - "4328:4318" @@ -49,7 +50,7 @@ services: - jaeger grafana: - image: grafana/grafana:latest + image: grafana/grafana:11.1.0 ports: - "3000:3000" environment: diff --git a/examples/otel_tracing/main.cpp b/examples/otel_tracing/main.cpp index 3551836abd2..acaf4d51608 100644 --- a/examples/otel_tracing/main.cpp +++ b/examples/otel_tracing/main.cpp @@ -2,7 +2,9 @@ #include #include #include +#include #include +#include #include #include @@ -23,6 +25,8 @@ #include #include +#include + #include #include #include @@ -376,6 +380,12 @@ int main(int argc, char** argv) { NLastGetopt::TOptsParseResult parsedOpts(&opts, argc, argv); + if (cfg.Endpoint.rfind("grpc://", 0) == 0) { + cfg.Endpoint.erase(0, 7); + } else if (cfg.Endpoint.rfind("grpcs://", 0) == 0) { + cfg.Endpoint.erase(0, 8); + } + std::cout << "Initializing OpenTelemetry..." << std::endl; std::cout << " OTLP endpoint: " << cfg.OtlpEndpoint << std::endl; diff --git a/include/ydb-cpp-sdk/client/trace/trace.h b/include/ydb-cpp-sdk/client/trace/trace.h index 117f4220b39..54ace20e8fb 100644 --- a/include/ydb-cpp-sdk/client/trace/trace.h +++ b/include/ydb-cpp-sdk/client/trace/trace.h @@ -15,6 +15,17 @@ enum class ESpanKind { CONSUMER }; +enum class ESpanStatus { + Unset, + Ok, + Error +}; + +class IScope { +public: + virtual ~IScope() = default; +}; + class ISpan { public: virtual ~ISpan() = default; @@ -22,12 +33,20 @@ class ISpan { virtual void SetAttribute(const std::string& key, const std::string& value) = 0; virtual void SetAttribute(const std::string& key, int64_t value) = 0; virtual void AddEvent(const std::string& name, const std::map& attributes = {}) = 0; + virtual std::unique_ptr Activate() = 0; + + virtual void SetStatus(ESpanStatus /*status*/, const std::string& /*description*/ = {}) {} }; class ITracer { public: virtual ~ITracer() = default; - virtual std::shared_ptr StartSpan(const std::string& name, ESpanKind kind = ESpanKind::INTERNAL) = 0; + + virtual std::shared_ptr StartSpan( + const std::string& name + , ESpanKind kind = ESpanKind::INTERNAL + , ISpan* parent = nullptr + ) = 0; }; class ITraceProvider { diff --git a/plugins/trace/otel/src/trace.cpp b/plugins/trace/otel/src/trace.cpp index 41b1df64793..16d81a8d72a 100644 --- a/plugins/trace/otel/src/trace.cpp +++ b/plugins/trace/otel/src/trace.cpp @@ -1,6 +1,10 @@ #include #include +#include +#include +#include +#include #include #include @@ -8,25 +12,50 @@ namespace NYdb::inline V3::NTrace { namespace { -using namespace opentelemetry; +namespace otel_trace = opentelemetry::trace; +namespace otel_nostd = opentelemetry::nostd; +namespace otel_common = opentelemetry::common; -trace::SpanKind MapSpanKind(ESpanKind kind) { +otel_trace::SpanKind MapSpanKind(ESpanKind kind) { switch (kind) { - case ESpanKind::INTERNAL: return trace::SpanKind::kInternal; - case ESpanKind::SERVER: return trace::SpanKind::kServer; - case ESpanKind::CLIENT: return trace::SpanKind::kClient; - case ESpanKind::PRODUCER: return trace::SpanKind::kProducer; - case ESpanKind::CONSUMER: return trace::SpanKind::kConsumer; + case ESpanKind::INTERNAL: return otel_trace::SpanKind::kInternal; + case ESpanKind::SERVER: return otel_trace::SpanKind::kServer; + case ESpanKind::CLIENT: return otel_trace::SpanKind::kClient; + case ESpanKind::PRODUCER: return otel_trace::SpanKind::kProducer; + case ESpanKind::CONSUMER: return otel_trace::SpanKind::kConsumer; } - return trace::SpanKind::kInternal; + return otel_trace::SpanKind::kInternal; } +otel_trace::StatusCode MapSpanStatus(ESpanStatus status) { + switch (status) { + case ESpanStatus::Unset: return otel_trace::StatusCode::kUnset; + case ESpanStatus::Ok: return otel_trace::StatusCode::kOk; + case ESpanStatus::Error: return otel_trace::StatusCode::kError; + } + return otel_trace::StatusCode::kUnset; +} + +class TOtelScope : public IScope { +public: + TOtelScope(otel_nostd::shared_ptr span) + : Scope_(std::move(span)) + {} + +private: + otel_trace::Scope Scope_; +}; + class TOtelSpan : public ISpan { public: - TOtelSpan(nostd::shared_ptr span) + TOtelSpan(otel_nostd::shared_ptr span) : Span_(std::move(span)) {} + const otel_nostd::shared_ptr& RawSpan() const noexcept { + return Span_; + } + void End() override { Span_->End(); } @@ -43,38 +72,50 @@ class TOtelSpan : public ISpan { if (attributes.empty()) { Span_->AddEvent(name); } else { - std::vector> attrs; + std::vector> attrs; attrs.reserve(attributes.size()); for (const auto& [k, v] : attributes) { - attrs.emplace_back(nostd::string_view(k), common::AttributeValue(nostd::string_view(v))); + attrs.emplace_back(otel_nostd::string_view(k), otel_common::AttributeValue(otel_nostd::string_view(v))); } Span_->AddEvent(name, attrs); } } + std::unique_ptr Activate() override { + return std::make_unique(Span_); + } + + void SetStatus(ESpanStatus status, const std::string& description) override { + Span_->SetStatus(MapSpanStatus(status), description); + } + private: - nostd::shared_ptr Span_; + otel_nostd::shared_ptr Span_; }; class TOtelTracer : public ITracer { public: - TOtelTracer(nostd::shared_ptr tracer) + TOtelTracer(otel_nostd::shared_ptr tracer) : Tracer_(std::move(tracer)) {} - std::shared_ptr StartSpan(const std::string& name, ESpanKind kind) override { - trace::StartSpanOptions options; + std::shared_ptr StartSpan(const std::string& name, ESpanKind kind, ISpan* parent) override { + otel_trace::StartSpanOptions options; options.kind = MapSpanKind(kind); + if (auto* otelParent = dynamic_cast(parent)) { + auto context = opentelemetry::context::RuntimeContext::GetCurrent(); + options.parent = otel_trace::SetSpan(context, otelParent->RawSpan()); + } return std::make_shared(Tracer_->StartSpan(name, options)); } private: - nostd::shared_ptr Tracer_; + otel_nostd::shared_ptr Tracer_; }; class TOtelTraceProvider : public ITraceProvider { public: - TOtelTraceProvider(nostd::shared_ptr tracerProvider) + TOtelTraceProvider(otel_nostd::shared_ptr tracerProvider) : TracerProvider_(std::move(tracerProvider)) {} @@ -83,7 +124,7 @@ class TOtelTraceProvider : public ITraceProvider { } private: - nostd::shared_ptr TracerProvider_; + otel_nostd::shared_ptr TracerProvider_; }; } // namespace diff --git a/src/client/impl/internal/retry/CMakeLists.txt b/src/client/impl/internal/retry/CMakeLists.txt index bde0f9e0aeb..d8d58530000 100644 --- a/src/client/impl/internal/retry/CMakeLists.txt +++ b/src/client/impl/internal/retry/CMakeLists.txt @@ -3,6 +3,7 @@ _ydb_sdk_add_library(impl-internal-retry) target_link_libraries(impl-internal-retry PUBLIC yutil impl-internal-grpc_connections + impl-observability ) target_sources(impl-internal-retry PRIVATE diff --git a/src/client/impl/internal/retry/retry.cpp b/src/client/impl/internal/retry/retry.cpp index 73880d0e5c6..526175189b9 100644 --- a/src/client/impl/internal/retry/retry.cpp +++ b/src/client/impl/internal/retry/retry.cpp @@ -28,14 +28,18 @@ TBackoffDuration CalcBackoffTime(const TBackoffSettings& settings, std::uint32_t } -void Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber) { - std::this_thread::sleep_for(CalcBackoffTime(settings, retryNumber)); +std::chrono::microseconds Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber) { + const auto duration = CalcBackoffTime(settings, retryNumber); + std::this_thread::sleep_for(duration); + return std::chrono::duration_cast(duration); } -void AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, +std::chrono::microseconds AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, std::uint32_t retryNumber, const std::function& fn) { - client->ScheduleTask(fn, std::chrono::duration_cast(CalcBackoffTime(settings, retryNumber))); + const auto duration = CalcBackoffTime(settings, retryNumber); + client->ScheduleTask(fn, std::chrono::duration_cast(duration)); + return std::chrono::duration_cast(duration); } } diff --git a/src/client/impl/internal/retry/retry.h b/src/client/impl/internal/retry/retry.h index 6fe090409c8..fad48053579 100644 --- a/src/client/impl/internal/retry/retry.h +++ b/src/client/impl/internal/retry/retry.h @@ -21,8 +21,8 @@ class IClientImplCommon; namespace NYdb::inline V3::NRetry { -void Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber); -void AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, +std::chrono::microseconds Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber); +std::chrono::microseconds AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, std::uint32_t retryNumber, const std::function& fn); enum class NextStep { diff --git a/src/client/impl/internal/retry/retry_async.h b/src/client/impl/internal/retry/retry_async.h index fbaa06df329..9e06d033c54 100644 --- a/src/client/impl/internal/retry/retry_async.h +++ b/src/client/impl/internal/retry/retry_async.h @@ -1,9 +1,18 @@ #pragma once +#include + #include +#include #include +#include +#include +#include +#include +#include + namespace NYdb::inline V3::NRetry::Async { template @@ -18,9 +27,45 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { public: TAsyncStatusType Execute() { + ParentSpan_ = Client_.Impl_->CreateRetryRootSpan(); + + // parentScope publishes the root retry span as the active thread-local + // context only for the synchronous prefix below (until DoRetry returns). + // The attempt span is parented to ParentSpan_ explicitly via + // CreateRetryAttemptSpan(..., ParentSpan_) in StartAttemptSpan(), so the + // root <- attempt link does NOT depend on this scope and survives the + // asynchronous boundary in the .Apply() callback below. + auto parentScope = ParentSpan_ ? ParentSpan_->Activate() : nullptr; + this->RetryStartTime_ = TInstant::Now(); - this->Retry(); - return this->Promise_.GetFuture(); + TPtr self(this); + DoRetry(self); + + return this->Promise_.GetFuture().Apply( + [self](const auto& f) mutable { + try { + auto value = f.GetValue(); + if (self->ParentSpan_) { + self->ParentSpan_->SetRetryCount(self->RetryNumber_); + self->ParentSpan_->End(value.GetStatus()); + } + return value; + } catch (...) { + if (self->ParentSpan_) { + self->ParentSpan_->SetRetryCount(self->RetryNumber_); + try { + std::rethrow_exception(std::current_exception()); + } catch (const std::exception& e) { + self->ParentSpan_->RecordException(typeid(e).name(), e.what()); + } catch (...) { + self->ParentSpan_->RecordException("unknown", "unknown exception"); + } + self->ParentSpan_->End(EStatus::CLIENT_INTERNAL_ERROR); + } + throw; + } + } + ); } protected: @@ -35,21 +80,27 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { virtual TAsyncStatusType RunOperation() = 0; static void DoRetry(TPtr self) { + self->StartAttemptSpan(); + + auto scope = self->AttemptSpan_ ? self->AttemptSpan_->Activate() : nullptr; self->Retry(); } static void DoBackoff(TPtr self, bool fast) { auto backoffSettings = fast ? self->Settings_.FastBackoffSettings_ : self->Settings_.SlowBackoffSettings_; - AsyncBackoff(self->Client_.Impl_, backoffSettings, self->RetryNumber_, + const auto backoff = AsyncBackoff(self->Client_.Impl_, backoffSettings, self->RetryNumber_, [self]() {DoRetry(self);}); + self->LastBackoffMs_ = std::chrono::duration_cast(backoff).count(); } static void HandleExceptionAsync(TPtr self, std::exception_ptr e) { + self->EndAttemptSpan(EStatus::CLIENT_INTERNAL_ERROR); self->Promise_.SetException(e); } static void HandleStatusAsync(TPtr self, const TStatusType& status) { + self->EndAttemptSpan(status.GetStatus()); auto nextStep = self->GetNextStep(status); if (nextStep != NextStep::Finish) { self->RetryNumber_++; @@ -79,6 +130,23 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { } ); } + +private: + void StartAttemptSpan() { + AttemptSpan_ = Client_.Impl_->CreateRetryAttemptSpan( + this->RetryNumber_, LastBackoffMs_, ParentSpan_); + } + + void EndAttemptSpan(EStatus status) { + if (AttemptSpan_) { + AttemptSpan_->End(status); + AttemptSpan_.reset(); + } + } + + std::shared_ptr ParentSpan_; + std::shared_ptr AttemptSpan_; + std::int64_t LastBackoffMs_ = 0; }; template > diff --git a/src/client/impl/internal/retry/retry_sync.h b/src/client/impl/internal/retry/retry_sync.h index beefcb27714..4452368d548 100644 --- a/src/client/impl/internal/retry/retry_sync.h +++ b/src/client/impl/internal/retry/retry_sync.h @@ -1,9 +1,15 @@ #pragma once #include +#include #include #include +#include + +#include +#include +#include namespace NYdb::inline V3::NRetry::Sync { @@ -14,46 +20,96 @@ class TRetryContext : public TRetryContextBase { public: TStatusType Execute() { + ParentSpan_ = Client_.Impl_->CreateRetryRootSpan(); + + auto parentScope = ParentSpan_ ? ParentSpan_->Activate() : nullptr; + auto& parentSpan = ParentSpan_; + + try { + auto status = ExecuteImpl(); + if (parentSpan) { + parentSpan->SetRetryCount(this->RetryNumber_); + parentSpan->End(status.GetStatus()); + } + return status; + } catch (...) { + if (parentSpan) { + parentSpan->SetRetryCount(this->RetryNumber_); + try { + std::rethrow_exception(std::current_exception()); + } catch (const std::exception& e) { + parentSpan->RecordException(typeid(e).name(), e.what()); + } catch (...) { + parentSpan->RecordException("unknown", "unknown exception"); + } + parentSpan->End(EStatus::CLIENT_INTERNAL_ERROR); + } + throw; + } + } + +protected: + TRetryContext(TClient& client, const TRetryOperationSettings& settings) + : TRetryContextBase(settings) + , Client_(client) + {} + + virtual TStatusType Retry() = 0; + + virtual TStatusType RunOperation() = 0; + + std::chrono::microseconds DoBackoff(bool fast) { + const auto &settings = fast ? this->Settings_.FastBackoffSettings_ + : this->Settings_.SlowBackoffSettings_; + return Backoff(settings, this->RetryNumber_); + } + +private: + TStatusType ExecuteImpl() { this->RetryStartTime_ = TInstant::Now(); - TStatusType status = Retry(); // first attempt + std::int64_t lastBackoffMs = 0; + + TStatusType status = RunAttempt(lastBackoffMs); for (this->RetryNumber_ = 0; this->RetryNumber_ <= this->Settings_.MaxRetries_;) { auto nextStep = this->GetNextStep(status); + std::chrono::microseconds backoff{}; switch (nextStep) { case NextStep::RetryImmediately: break; case NextStep::RetryFastBackoff: - DoBackoff(true); + backoff = DoBackoff(true); break; case NextStep::RetrySlowBackoff: - DoBackoff(false); + backoff = DoBackoff(false); break; case NextStep::Finish: return status; } - // make next retry this->RetryNumber_++; this->LogRetry(status); this->Client_.Impl_->CollectRetryStatSync(status.GetStatus()); - status = Retry(); + lastBackoffMs = std::chrono::duration_cast(backoff).count(); + status = RunAttempt(lastBackoffMs); } return status; } -protected: - TRetryContext(TClient& client, const TRetryOperationSettings& settings) - : TRetryContextBase(settings) - , Client_(client) - {} + TStatusType RunAttempt(std::int64_t backoffMs) { + auto attemptSpan = Client_.Impl_->CreateRetryAttemptSpan(this->RetryNumber_, backoffMs, ParentSpan_); + std::unique_ptr scope; + if (attemptSpan) { + scope = attemptSpan->Activate(); + } - virtual TStatusType Retry() = 0; + TStatusType status = Retry(); - virtual TStatusType RunOperation() = 0; - - void DoBackoff(bool fast) { - const auto &settings = fast ? this->Settings_.FastBackoffSettings_ - : this->Settings_.SlowBackoffSettings_; - Backoff(settings, this->RetryNumber_); + if (attemptSpan) { + attemptSpan->End(status.GetStatus()); + } + return status; } + + std::shared_ptr ParentSpan_; }; template> diff --git a/src/client/impl/observability/CMakeLists.txt b/src/client/impl/observability/CMakeLists.txt index ea060beda47..02701526976 100644 --- a/src/client/impl/observability/CMakeLists.txt +++ b/src/client/impl/observability/CMakeLists.txt @@ -1,14 +1,20 @@ +add_subdirectory(error_category) + _ydb_sdk_add_library(impl-observability) target_link_libraries(impl-observability PUBLIC yutil client-metrics + client-trace client-impl-ydb_stats + impl-observability-error_category + impl-internal-db_driver_state ) target_sources(impl-observability PRIVATE metrics.cpp observation.cpp + operation_name.cpp span.cpp ) diff --git a/src/client/impl/observability/error_category/CMakeLists.txt b/src/client/impl/observability/error_category/CMakeLists.txt new file mode 100644 index 00000000000..1904d9f68a0 --- /dev/null +++ b/src/client/impl/observability/error_category/CMakeLists.txt @@ -0,0 +1,11 @@ +_ydb_sdk_add_library(impl-observability-error_category) + +target_link_libraries(impl-observability-error_category PUBLIC + yutil +) + +target_sources(impl-observability-error_category PRIVATE + error_category.cpp +) + +_ydb_sdk_install_targets(TARGETS impl-observability-error_category) diff --git a/src/client/impl/observability/error_category/error_category.cpp b/src/client/impl/observability/error_category/error_category.cpp new file mode 100644 index 00000000000..87ea95b0f9d --- /dev/null +++ b/src/client/impl/observability/error_category/error_category.cpp @@ -0,0 +1,13 @@ +#include "error_category.h" + +#include + +namespace NYdb::inline V3::NObservability { + +std::string_view CategorizeErrorType(EStatus status) noexcept { + return static_cast(status) >= TRANSPORT_STATUSES_FIRST + ? std::string_view("transport_error") + : std::string_view("ydb_error"); +} + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/error_category/error_category.h b/src/client/impl/observability/error_category/error_category.h new file mode 100644 index 00000000000..b98781da149 --- /dev/null +++ b/src/client/impl/observability/error_category/error_category.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +#include + +namespace NYdb::inline V3::NObservability { + +// Maps EStatus to OTel-style error.type category: +// "transport_error" for client/transport-layer statuses (>= TRANSPORT_STATUSES_FIRST), +// "ydb_error" for server-side YDB statuses. +std::string_view CategorizeErrorType(EStatus status) noexcept; + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/metrics.cpp b/src/client/impl/observability/metrics.cpp index 01a96f2cc58..8443c27f173 100644 --- a/src/client/impl/observability/metrics.cpp +++ b/src/client/impl/observability/metrics.cpp @@ -1,5 +1,7 @@ #include "metrics.h" +#include "operation_name.h" + #include #include @@ -32,14 +34,14 @@ TRequestMetrics::TRequestMetrics(NSdkStats::TStatCollector::TClientOperationStat , const std::string& requestName , const TLog& log ) : Collector_(operationCollector) - , RequestName_(requestName) + , RequestName_(NormalizeOperationName(requestName)) , Log_(log) { if (!Collector_) { return; } try { - Collector_->IncRequestCount(requestName); + Collector_->IncRequestCount(RequestName_); StartTime_ = std::chrono::steady_clock::now(); } catch (...) { SafeLogRequestMetricsError(Log_, "failed to initialize metrics", std::current_exception()); diff --git a/src/client/impl/observability/observation.cpp b/src/client/impl/observability/observation.cpp index 7483087bcfd..2e5d7fe0c62 100644 --- a/src/client/impl/observability/observation.cpp +++ b/src/client/impl/observability/observation.cpp @@ -1,26 +1,32 @@ #include "observation.h" +#define INCLUDE_YDB_INTERNAL_H +#include +#undef INCLUDE_YDB_INTERNAL_H + namespace NYdb::inline V3::NObservability { TRequestObservation::TRequestObservation(const std::string& ydbClientType , NSdkStats::TStatCollector::TClientOperationStatCollector* operationCollector , std::shared_ptr tracer , const std::string& operationName - , const std::string& discoveryEndpoint - , const std::string& database - , const TLog& log -) : Span_(std::make_shared(std::move(tracer), operationName, discoveryEndpoint, database, log, ydbClientType)) - , Metrics_(std::make_shared(operationCollector, operationName, log)) + , const std::shared_ptr& dbDriverState +) : Span_( + TRequestSpan::Create(ydbClientType + , std::move(tracer) + , operationName + , dbDriverState->DiscoveryEndpoint + , dbDriverState->Database + , dbDriverState->Log + ) + ), Metrics_( + std::make_shared(operationCollector, operationName, dbDriverState->Log) + ) {} -void TRequestObservation::SetPeerEndpoint(const std::string& endpoint) noexcept { +void TRequestObservation::End(EStatus status, const std::string& endpoint) noexcept { if (Span_) { Span_->SetPeerEndpoint(endpoint); - } -} - -void TRequestObservation::End(EStatus status) noexcept { - if (Span_) { Span_->End(status); } if (Metrics_) { diff --git a/src/client/impl/observability/observation.h b/src/client/impl/observability/observation.h index 544b0c4baff..542e2f6c95b 100644 --- a/src/client/impl/observability/observation.h +++ b/src/client/impl/observability/observation.h @@ -14,13 +14,10 @@ class TRequestObservation { , NSdkStats::TStatCollector::TClientOperationStatCollector* operationCollector , std::shared_ptr tracer , const std::string& operationName - , const std::string& discoveryEndpoint - , const std::string& database - , const TLog& log + , const std::shared_ptr& dbDriverState ); - void SetPeerEndpoint(const std::string& endpoint) noexcept; - void End(EStatus status) noexcept; + void End(EStatus status, const std::string& endpoint = "") noexcept; void EndWithClientInternalError() noexcept; private: diff --git a/src/client/impl/observability/operation_name.cpp b/src/client/impl/observability/operation_name.cpp new file mode 100644 index 00000000000..08dc75784ad --- /dev/null +++ b/src/client/impl/observability/operation_name.cpp @@ -0,0 +1,17 @@ +#include "operation_name.h" + +#include + +namespace NYdb::inline V3::NObservability { + +std::string NormalizeOperationName(const std::string& requestName) { + static constexpr std::string_view kPrefix = "ydb."; + if (requestName.size() >= kPrefix.size() + && std::string_view(requestName.data(), kPrefix.size()) == kPrefix) + { + return requestName; + } + return std::string(kPrefix) + requestName; +} + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/operation_name.h b/src/client/impl/observability/operation_name.h new file mode 100644 index 00000000000..33f18d86f29 --- /dev/null +++ b/src/client/impl/observability/operation_name.h @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace NYdb::inline V3::NObservability { + +std::string NormalizeOperationName(const std::string& requestName); + +} // namespace NYdb::NObservability diff --git a/src/client/impl/observability/span.cpp b/src/client/impl/observability/span.cpp index 46805563fee..8d86cfd25e8 100644 --- a/src/client/impl/observability/span.cpp +++ b/src/client/impl/observability/span.cpp @@ -1,7 +1,14 @@ #include "span.h" +#include "operation_name.h" + +#include #include +#define INCLUDE_YDB_INTERNAL_H +#include +#undef INCLUDE_YDB_INTERNAL_H + #include #include @@ -11,6 +18,8 @@ namespace NYdb::inline V3::NObservability { namespace { constexpr int DefaultGrpcPort = 2135; +constexpr const char* kRetryRootSpanName = "ydb.RunWithRetry"; +constexpr const char* kRetryAttemptSpanName = "ydb.Try"; std::string YdbClientApiAttributeValue(const std::string& clientType) noexcept { return clientType.empty() ? std::string("Unspecified") : clientType; @@ -48,6 +57,21 @@ void ParseEndpoint(const std::string& endpoint, std::string& host, int& port) { } } +void EmitExceptionEvent(NTrace::ISpan& span, + const std::string& type, + const std::string& message, + const std::string& stacktrace) +{ + std::map attrs{ + {"exception.type", type}, + {"exception.message", message}, + }; + if (!stacktrace.empty()) { + attrs.emplace("exception.stacktrace", stacktrace); + } + span.AddEvent("exception", attrs); +} + void SafeLogRequestSpanError(TLog& log, const char* message, std::exception_ptr exception) noexcept { try { if (!exception) { @@ -68,12 +92,74 @@ void SafeLogRequestSpanError(TLog& log, const char* message, std::exception_ptr } // namespace -TRequestSpan::TRequestSpan(std::shared_ptr tracer +std::shared_ptr TRequestSpan::Create(const std::string& ydbClientType + , std::shared_ptr tracer + , const std::string& requestName + , const std::string& discoveryEndpoint + , const std::string& database + , const TLog& log + , NTrace::ESpanKind kind + , const std::shared_ptr& parent +) { + NTrace::ISpan* parentRaw = parent ? parent->Span_.get() : nullptr; + return std::shared_ptr(new TRequestSpan( + ydbClientType, + std::move(tracer), + requestName, + discoveryEndpoint, + database, + log, + kind, + parentRaw + )); +} + +std::shared_ptr TRequestSpan::CreateForClientRetry(const std::string& ydbClientType + , std::shared_ptr tracer + , const std::shared_ptr& dbDriverState +) { + return Create( + ydbClientType, + std::move(tracer), + kRetryRootSpanName, + dbDriverState->DiscoveryEndpoint, + dbDriverState->Database, + dbDriverState->Log, + NTrace::ESpanKind::INTERNAL + ); +} + +std::shared_ptr TRequestSpan::CreateForRetryAttempt(const std::string& ydbClientType + , std::shared_ptr tracer + , const std::shared_ptr& dbDriverState + , std::uint32_t attempt + , std::int64_t backoffMs + , const std::shared_ptr& parent +) { + auto span = Create( + ydbClientType, + std::move(tracer), + kRetryAttemptSpanName, + dbDriverState->DiscoveryEndpoint, + dbDriverState->Database, + dbDriverState->Log, + NTrace::ESpanKind::INTERNAL, + parent + ); + if (span) { + span->SetRetryAttributes(attempt, backoffMs); + } + return span; +} + +TRequestSpan::TRequestSpan(const std::string& ydbClientType + , std::shared_ptr tracer , const std::string& requestName - , const std::string& endpoint + , const std::string& discoveryEndpoint , const std::string& database , const TLog& log - , const std::string& ydbClientType + , NTrace::ESpanKind kind + , NTrace::ISpan* parent ) : Log_(log) { if (!tracer) { return; @@ -81,16 +167,17 @@ TRequestSpan::TRequestSpan(std::shared_ptr tracer std::string host; int port; - ParseEndpoint(endpoint, host, port); + ParseEndpoint(discoveryEndpoint, host, port); try { - Span_ = tracer->StartSpan(requestName, NTrace::ESpanKind::CLIENT); + const auto operationName = NormalizeOperationName(requestName); + Span_ = tracer->StartSpan(operationName, kind, parent); if (!Span_) { return; } Span_->SetAttribute("db.system.name", "ydb"); Span_->SetAttribute("db.namespace", database); - Span_->SetAttribute("db.operation.name", requestName); + Span_->SetAttribute("db.operation.name", operationName); Span_->SetAttribute("ydb.client.api", YdbClientApiAttributeValue(ydbClientType)); Span_->SetAttribute("server.address", host); Span_->SetAttribute("server.port", static_cast(port)); @@ -136,12 +223,39 @@ void TRequestSpan::AddEvent(const std::string& name, const std::map TRequestSpan::Activate() noexcept { + if (!Span_) { + return nullptr; + } + try { + return Span_->Activate(); + } catch (...) { + SafeLogRequestSpanError(Log_, "failed to activate span", std::current_exception()); + return nullptr; + } +} + void TRequestSpan::End(EStatus status) noexcept { if (Span_) { try { - Span_->SetAttribute("db.response.status_code", ToString(status)); if (status != EStatus::SUCCESS) { - Span_->SetAttribute("error.type", ToString(status)); + const auto statusName = ToString(status); + const auto errorType = CategorizeErrorType(status); + Span_->SetAttribute("db.response.status_code", statusName); + Span_->SetAttribute("error.type", std::string(errorType)); + EmitExceptionEvent(*Span_, statusName, statusName, /*stacktrace=*/""); + Span_->SetStatus(NTrace::ESpanStatus::Error, statusName); } Span_->End(); } catch (...) { @@ -151,4 +265,27 @@ void TRequestSpan::End(EStatus status) noexcept { } } +void TRequestSpan::SetRetryCount(std::uint32_t count) noexcept { + if (!Span_ || count == 0) { + return; + } + try { + Span_->SetAttribute("ydb.retry.count", static_cast(count)); + } catch (...) { + SafeLogRequestSpanError(Log_, "failed to set retry count", std::current_exception()); + } +} + +void TRequestSpan::SetRetryAttributes(std::uint32_t attempt, std::int64_t backoffMs) noexcept { + if (!Span_ || attempt == 0) { + return; + } + try { + Span_->SetAttribute("ydb.retry.attempt", static_cast(attempt)); + Span_->SetAttribute("ydb.retry.backoff_ms", backoffMs); + } catch (...) { + SafeLogRequestSpanError(Log_, "failed to set retry attributes", std::current_exception()); + } +} + } // namespace NYdb::NObservability diff --git a/src/client/impl/observability/span.h b/src/client/impl/observability/span.h index dbbde456e63..5c1347617b4 100644 --- a/src/client/impl/observability/span.h +++ b/src/client/impl/observability/span.h @@ -2,32 +2,77 @@ #include #include +#include #include +#include #include #include #include +namespace NYdb::inline V3 { + +class TDbDriverState; + +} // namespace NYdb::inline V3 + namespace NYdb::inline V3::NObservability { class TRequestSpan { public: - TRequestSpan(std::shared_ptr tracer + static std::shared_ptr Create( + const std::string& ydbClientType + , std::shared_ptr tracer , const std::string& requestName - , const std::string& endpoint + , const std::string& discoveryEndpoint , const std::string& database , const TLog& log - , const std::string& ydbClientType = {} + , NTrace::ESpanKind kind = NTrace::ESpanKind::CLIENT + , const std::shared_ptr& parent = nullptr + ); + + static std::shared_ptr CreateForClientRetry( + const std::string& ydbClientType + , std::shared_ptr tracer + , const std::shared_ptr& dbDriverState + ); + + static std::shared_ptr CreateForRetryAttempt( + const std::string& ydbClientType + , std::shared_ptr tracer + , const std::shared_ptr& dbDriverState + , std::uint32_t attempt + , std::int64_t backoffMs + , const std::shared_ptr& parent = nullptr ); + ~TRequestSpan() noexcept; + TRequestSpan(const TRequestSpan&) = delete; + TRequestSpan& operator=(const TRequestSpan&) = delete; + void SetPeerEndpoint(const std::string& endpoint) noexcept; void AddEvent(const std::string& name, const std::map& attributes = {}) noexcept; + void RecordException(const std::string& type, const std::string& message, const std::string& stacktrace = {}) noexcept; + std::unique_ptr Activate() noexcept; + void SetRetryCount(std::uint32_t count) noexcept; void End(EStatus status) noexcept; private: + TRequestSpan(const std::string& ydbClientType + , std::shared_ptr tracer + , const std::string& requestName + , const std::string& discoveryEndpoint + , const std::string& database + , const TLog& log + , NTrace::ESpanKind kind + , NTrace::ISpan* parent + ); + + void SetRetryAttributes(std::uint32_t attempt, std::int64_t backoffMs) noexcept; + TLog Log_; std::shared_ptr Span_; }; diff --git a/src/client/impl/stats/CMakeLists.txt b/src/client/impl/stats/CMakeLists.txt index 15866af4bc6..0fac8860363 100644 --- a/src/client/impl/stats/CMakeLists.txt +++ b/src/client/impl/stats/CMakeLists.txt @@ -5,6 +5,7 @@ target_link_libraries(client-impl-ydb_stats PUBLIC grpc-client monlib-metrics client-metrics + impl-observability-error_category ) target_sources(client-impl-ydb_stats PRIVATE diff --git a/src/client/impl/stats/stats.h b/src/client/impl/stats/stats.h index b2e61c0ace8..570f9e4675a 100644 --- a/src/client/impl/stats/stats.h +++ b/src/client/impl/stats/stats.h @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -326,10 +327,10 @@ struct TStatCollector { {"db.namespace", Database_}, {"db.operation.name", operationName}, {"ydb.client.api", YdbClientApiAttributeValue(ClientType_)}, - {"db.response.status_code", TStringBuilder() << status}, }; if (status != EStatus::SUCCESS) { - labels["error.type"] = TStringBuilder() << status; + labels["db.response.status_code"] = TStringBuilder() << status; + labels["error.type"] = std::string(NObservability::CategorizeErrorType(status)); } ExternalRegistry_->Histogram( "db.client.operation.duration", diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index a6db0b273db..201ba27cc3e 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -19,7 +19,7 @@ #include #include -#include +#include #include @@ -104,14 +104,15 @@ class TQueryClient::TImpl: public TClientImplCommon, public CollectParamsSize(params ? ¶ms->GetProtoMap() : nullptr); auto obs = MakeObservation("ExecuteQuery"); + std::string sessionEndpoint = session.has_value() ? session->SessionImpl_->GetEndpoint() : std::string{}; return TExecQueryImpl::ExecuteQuery( Connections_, DbDriverState_, query, txControl, params, settings, session) - .Apply([obs](TAsyncExecuteQueryResult future) { + .Apply([obs, sessionEndpoint = std::move(sessionEndpoint)](TAsyncExecuteQueryResult future) { try { auto result = future.GetValue(); - obs->SetPeerEndpoint(result.GetEndpoint()); - obs->End(result.GetStatus()); + const auto& resultEndpoint = result.GetEndpoint(); + obs->End(result.GetStatus(), !resultEndpoint.empty() ? resultEndpoint : sessionEndpoint); return result; } catch (...) { obs->EndWithClientInternalError(); @@ -189,18 +190,17 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto responseCb = [promise, session, obs] (Ydb::Query::RollbackTransactionResponse* response, TPlainStatus status) mutable { try { - obs->SetPeerEndpoint(status.Endpoint); if (response) { NYdb::NIssue::TIssues opIssues; NYdb::NIssue::IssuesFromMessage(response->issues(), opIssues); TStatus rollbackTxStatus(TPlainStatus{static_cast(response->status()), std::move(opIssues), status.Endpoint, std::move(status.Metadata)}); - obs->End(rollbackTxStatus.GetStatus()); + obs->End(rollbackTxStatus.GetStatus(), rollbackTxStatus.GetEndpoint()); promise.SetValue(std::move(rollbackTxStatus)); } else { - obs->End(status.Status); + obs->End(status.Status, status.Endpoint); promise.SetValue(TStatus(std::move(status))); } } catch (...) { @@ -237,19 +237,18 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto responseCb = [promise, session, obs] (Ydb::Query::CommitTransactionResponse* response, TPlainStatus status) mutable { try { - obs->SetPeerEndpoint(status.Endpoint); if (response) { NYdb::NIssue::TIssues opIssues; NYdb::NIssue::IssuesFromMessage(response->issues(), opIssues); TStatus commitTxStatus(TPlainStatus{static_cast(response->status()), std::move(opIssues), status.Endpoint, std::move(status.Metadata)}); - obs->End(commitTxStatus.GetStatus()); + obs->End(commitTxStatus.GetStatus(), commitTxStatus.GetEndpoint()); TCommitTransactionResult commitTxResult(std::move(commitTxStatus)); promise.SetValue(std::move(commitTxResult)); } else { - obs->End(status.Status); + obs->End(status.Status, status.Endpoint); promise.SetValue(TCommitTransactionResult(TStatus(std::move(status)))); } } catch (...) { @@ -491,7 +490,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public ); if (Observation) { - Observation->End(EStatus::SUCCESS); + Observation->End(EStatus::SUCCESS, session->GetEndpoint()); } ScheduleReply(std::move(val)); } @@ -502,8 +501,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public { auto val = future.ExtractValue(); if (obs) { - obs->SetPeerEndpoint(val.GetEndpoint()); - obs->End(val.GetStatus()); + obs->End(val.GetStatus(), val.GetEndpoint()); } promise.SetValue(std::move(val)); }); @@ -533,7 +531,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public std::shared_ptr Observation; }; - auto obs = MakeObservation("GetSession"); + auto obs = MakeObservation("CreateSession"); auto ctx = std::make_unique(shared_from_this(), settings, obs); auto future = ctx->GetFuture(); SessionPool_.GetSession(std::move(ctx)); @@ -602,6 +600,29 @@ class TQueryClient::TImpl: public TClientImplCommon, public } } + std::shared_ptr CreateRetryRootSpan() { + return NObservability::TRequestSpan::CreateForClientRetry( + "Query", + Tracer_, + DbDriverState_ + ); + } + + std::shared_ptr CreateRetryAttemptSpan( + std::uint32_t attempt, + std::int64_t backoffMs, + const std::shared_ptr& parent = nullptr) + { + return NObservability::TRequestSpan::CreateForRetryAttempt( + "Query", + Tracer_, + DbDriverState_, + attempt, + backoffMs, + parent + ); + } + private: std::shared_ptr MakeObservation(const std::string& operationName) { return std::make_shared( @@ -609,9 +630,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public &OperationStatCollector_, Tracer_, operationName, - DbDriverState_->DiscoveryEndpoint, - DbDriverState_->Database, - DbDriverState_->Log + DbDriverState_ ); } @@ -692,28 +711,26 @@ int64_t TQueryClient::GetCurrentPoolSize() const { TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings) { - TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); - return ctx->Execute(); + return TRetryContextResultAsync::TPtr( + new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings))->Execute(); } TAsyncStatus TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings) { - TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); - return ctx->Execute(); + return TRetryContextAsync::TPtr( + new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings))->Execute(); } TAsyncStatus TQueryClient::RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRetryOperationSettings settings) { - TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithoutSession(*this, std::move(queryFunc), settings)); - return ctx->Execute(); + return TRetryContextAsync::TPtr( + new NRetry::Async::TRetryWithoutSession(*this, std::move(queryFunc), settings))->Execute(); } TStatus TQueryClient::RetryQuerySync(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings) { - NRetry::Sync::TRetryWithSession ctx(*this, queryFunc, settings); - return ctx.Execute(); + return NRetry::Sync::TRetryWithSession(*this, queryFunc, settings).Execute(); } TStatus TQueryClient::RetryQuerySync(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings) { - NRetry::Sync::TRetryWithoutSession ctx(*this, queryFunc, settings); - return ctx.Execute(); + return NRetry::Sync::TRetryWithoutSession(*this, queryFunc, settings).Execute(); } TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, const TTxControl& txControl, @@ -723,8 +740,8 @@ TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, cons auto queryFunc = [&query, &txControl](TSession session, TDuration duration) -> TAsyncExecuteQueryResult { return session.ExecuteQuery(query, txControl, TExecuteQuerySettings().ClientTimeout(duration)); }; - TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); - return ctx->Execute(); + return TRetryContextResultAsync::TPtr( + new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings))->Execute(); } //////////////////////////////////////////////////////////////////////////////// diff --git a/src/client/table/impl/table_client.cpp b/src/client/table/impl/table_client.cpp index 47626433eaa..4ddac0c0728 100644 --- a/src/client/table/impl/table_client.cpp +++ b/src/client/table/impl/table_client.cpp @@ -37,6 +37,29 @@ TTableClient::TImpl::TImpl(std::shared_ptr&& connections, SessionPool_.SetStatCollector(DbDriverState_->StatCollector.GetSessionPoolStatCollector("Table")); } +std::shared_ptr TTableClient::TImpl::CreateRetryRootSpan() { + return NObservability::TRequestSpan::CreateForClientRetry( + "Table", + Tracer_, + DbDriverState_ + ); +} + +std::shared_ptr TTableClient::TImpl::CreateRetryAttemptSpan( + std::uint32_t attempt + , std::int64_t backoffMs + , const std::shared_ptr& parent +) { + return NObservability::TRequestSpan::CreateForRetryAttempt( + "Table", + Tracer_, + DbDriverState_, + attempt, + backoffMs, + parent + ); +} + TTableClient::TImpl::~TImpl() { if (Connections_->GetDrainOnDtors()) { Drain().Wait(DRAIN_TIMEOUT); @@ -385,7 +408,7 @@ TAsyncCreateSessionResult TTableClient::TImpl::CreateSession(const TCreateSessio auto createSessionPromise = NewPromise(); auto self = shared_from_this(); - auto obs = MakeObservation("GetSession"); + auto obs = MakeObservation("CreateSession"); auto createSessionExtractor = [createSessionPromise, self, standalone, obs] (google::protobuf::Any* any, TPlainStatus status) mutable { @@ -403,8 +426,8 @@ TAsyncCreateSessionResult TTableClient::TImpl::CreateSession(const TCreateSessio // We do not use SessionStatusInterception for CreateSession request session.SessionImpl_->MarkBroken(); } + obs->End(status.Status, status.Endpoint); TCreateSessionResult val(TStatus(std::move(status)), std::move(session)); - obs->End(val.GetStatus()); createSessionPromise.SetValue(std::move(val)); }; @@ -778,7 +801,7 @@ TAsyncStatus TTableClient::TImpl::ExecuteSchemeQuery(const TSession& session, co return future.Apply([obs](NThreading::TFuture f) mutable { auto status = f.ExtractValue(); - obs->End(status.GetStatus()); + obs->End(status.GetStatus(), status.GetEndpoint()); return status; }); } @@ -806,9 +829,9 @@ TAsyncBeginTransactionResult TTableClient::TImpl::BeginTransaction(const TSessio txId = result.tx_meta().id(); } + obs->End(status.Status, status.Endpoint); TBeginTransactionResult beginTxResult(TStatus(std::move(status)), TTransaction(session, txId)); - obs->End(beginTxResult.GetStatus()); promise.SetValue(std::move(beginTxResult)); }; @@ -835,7 +858,7 @@ TAsyncCommitTransactionResult TTableClient::TImpl::CommitTransaction(const TSess request.set_tx_id(TStringType{txId}); request.set_collect_stats(GetStatsCollectionMode(settings.CollectQueryStats_)); - auto obs = MakeObservation("CommitTransaction"); + auto obs = MakeObservation("Commit"); auto promise = NewPromise(); @@ -851,8 +874,8 @@ TAsyncCommitTransactionResult TTableClient::TImpl::CommitTransaction(const TSess } } + obs->End(status.Status, status.Endpoint); TCommitTransactionResult commitTxResult(TStatus(std::move(status)), queryStats); - obs->End(commitTxResult.GetStatus()); promise.SetValue(std::move(commitTxResult)); }; @@ -878,7 +901,7 @@ TAsyncStatus TTableClient::TImpl::RollbackTransaction(const TSession& session, c request.set_session_id(TStringType{session.GetId()}); request.set_tx_id(TStringType{txId}); - auto obs = MakeObservation("RollbackTransaction"); + auto obs = MakeObservation("Rollback"); auto future = RunSimple( std::move(request), @@ -888,7 +911,7 @@ TAsyncStatus TTableClient::TImpl::RollbackTransaction(const TSession& session, c return future.Apply([obs](TAsyncStatus fut) { auto status = fut.GetValue(); - obs->End(status.GetStatus()); + obs->End(status.GetStatus(), status.GetEndpoint()); return status; }); } @@ -1165,8 +1188,8 @@ TAsyncBulkUpsertResult TTableClient::TImpl::BulkUpsert(const std::string& table, auto promise = NewPromise(); auto extractor = [promise, obs](google::protobuf::Any* any, TPlainStatus status) mutable { Y_UNUSED(any); + obs->End(status.Status, status.Endpoint); TBulkUpsertResult val(TStatus(std::move(status))); - obs->End(val.GetStatus()); promise.SetValue(std::move(val)); }; @@ -1216,8 +1239,8 @@ TAsyncBulkUpsertResult TTableClient::TImpl::BulkUpsert(const std::string& table, auto extractor = [promise, obs] (google::protobuf::Any* any, TPlainStatus status) mutable { Y_UNUSED(any); + obs->End(status.Status, status.Endpoint); TBulkUpsertResult val(TStatus(std::move(status))); - obs->End(val.GetStatus()); promise.SetValue(std::move(val)); }; diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 27c35c843ad..54fe07bf00c 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -12,7 +12,7 @@ #include #include #include -#include +#include #include "client_session.h" #include "data_query.h" @@ -157,6 +157,12 @@ class TTableClient::TImpl: public TClientImplCommon, public void CollectRetryStatAsync(EStatus status); void CollectRetryStatSync(EStatus status); + std::shared_ptr CreateRetryRootSpan(); + std::shared_ptr CreateRetryAttemptSpan(std::uint32_t attempt + , std::int64_t backoffMs + , const std::shared_ptr& parent = nullptr + ); + public: TClientSettings Settings_; @@ -288,11 +294,10 @@ class TTableClient::TImpl: public TClientImplCommon, public sessionPtr->SessionImpl_->AddQueryToCache(*dataQuery); } + obs->End(status.Status, status.Endpoint); TDataQueryResult dataQueryResult(TStatus(std::move(status)), std::move(res), tx, dataQuery, fromCache, queryStats); - obs->End(dataQueryResult.GetStatus()); - delete sessionPtr; tx.reset(); dataQuery.reset(); @@ -347,9 +352,7 @@ class TTableClient::TImpl: public TClientImplCommon, public &OperationStatCollector_, Tracer_, operationName, - DbDriverState_->DiscoveryEndpoint, - DbDriverState_->Database, - DbDriverState_->Log + DbDriverState_ ); } }; diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index bc0a885f237..298c875b99c 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -9,8 +9,8 @@ #include #undef INCLUDE_YDB_INTERNAL_H -#include -#include +#include +#include #include #include #include @@ -1590,23 +1590,21 @@ TTypeBuilder TTableClient::GetTypeBuilder() { //////////////////////////////////////////////////////////////////////////////// TAsyncStatus TTableClient::RetryOperation(TOperationFunc&& operation, const TRetryOperationSettings& settings) { - TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(operation), settings)); - return ctx->Execute(); + return TRetryContextAsync::TPtr( + new NRetry::Async::TRetryWithSession(*this, std::move(operation), settings))->Execute(); } TAsyncStatus TTableClient::RetryOperation(TOperationWithoutSessionFunc&& operation, const TRetryOperationSettings& settings) { - TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithoutSession(*this, std::move(operation), settings)); - return ctx->Execute(); + return TRetryContextAsync::TPtr( + new NRetry::Async::TRetryWithoutSession(*this, std::move(operation), settings))->Execute(); } TStatus TTableClient::RetryOperationSync(const TOperationWithoutSessionSyncFunc& operation, const TRetryOperationSettings& settings) { - NRetry::Sync::TRetryWithoutSession ctx(*this, operation, settings); - return ctx.Execute(); + return NRetry::Sync::TRetryWithoutSession(*this, operation, settings).Execute(); } TStatus TTableClient::RetryOperationSync(const TOperationSyncFunc& operation, const TRetryOperationSettings& settings) { - NRetry::Sync::TRetryWithSession ctx(*this, operation, settings); - return ctx.Execute(); + return NRetry::Sync::TRetryWithSession(*this, operation, settings).Execute(); } NThreading::TFuture TTableClient::Stop() { @@ -2511,10 +2509,6 @@ uint64_t TIndexDescription::GetSizeBytes() const { return SizeBytes_; } -void TIndexDescription::SetParallel(uint32_t parallel) { - Parallel_ = parallel; -} - TIndexDescription TIndexDescription::CreateGlobalIndex( const std::string& name, const std::vector& indexColumns, @@ -3059,9 +3053,6 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { if constexpr (std::is_same_v) { result.SizeBytes_ = proto.size_bytes(); } - if constexpr (std::is_same_v) { - result.Parallel_ = proto.parallel(); - } return result; } @@ -3074,8 +3065,6 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { *proto.mutable_data_columns() = {DataColumns_.begin(), DataColumns_.end()}; - proto.set_parallel(Parallel_); - switch (IndexType_) { case EIndexType::GlobalSync: { auto& settings = *proto.mutable_global_index()->mutable_settings(); diff --git a/tests/common/fake_trace_provider.h b/tests/common/fake_trace_provider.h new file mode 100644 index 00000000000..55ddb01bf12 --- /dev/null +++ b/tests/common/fake_trace_provider.h @@ -0,0 +1,163 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace NYdb::NTests { + +struct TFakeEvent { + std::string Name; + std::map Attributes; +}; + +class TFakeScope : public NTrace::IScope { +}; + +class TFakeSpan : public NTrace::ISpan { +public: + void End() override { + std::lock_guard lock(Mutex_); + Ended_ = true; + } + + void SetAttribute(const std::string& key, const std::string& value) override { + std::lock_guard lock(Mutex_); + StringAttributes_[key] = value; + } + + void SetAttribute(const std::string& key, int64_t value) override { + std::lock_guard lock(Mutex_); + IntAttributes_[key] = value; + } + + void AddEvent(const std::string& name, const std::map& attributes) override { + std::lock_guard lock(Mutex_); + Events_.push_back({name, attributes}); + } + + std::unique_ptr Activate() override { + std::lock_guard lock(Mutex_); + Activated_ = true; + return std::make_unique(); + } + + bool IsEnded() const { + std::lock_guard lock(Mutex_); + return Ended_; + } + + bool IsActivated() const { + std::lock_guard lock(Mutex_); + return Activated_; + } + + std::string GetStringAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + auto it = StringAttributes_.find(key); + return it != StringAttributes_.end() ? it->second : ""; + } + + bool HasStringAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + return StringAttributes_.contains(key); + } + + int64_t GetIntAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + auto it = IntAttributes_.find(key); + return it != IntAttributes_.end() ? it->second : 0; + } + + bool HasIntAttribute(const std::string& key) const { + std::lock_guard lock(Mutex_); + return IntAttributes_.contains(key); + } + + std::vector GetEvents() const { + std::lock_guard lock(Mutex_); + return Events_; + } + +private: + mutable std::mutex Mutex_; + bool Ended_ = false; + bool Activated_ = false; + std::map StringAttributes_; + std::map IntAttributes_; + std::vector Events_; +}; + +class TFakeTracer : public NTrace::ITracer { +public: + std::shared_ptr StartSpan( + const std::string& name, + NTrace::ESpanKind kind, + NTrace::ISpan* parent + ) override { + auto span = std::make_shared(); + std::lock_guard lock(Mutex_); + Spans_.push_back({name, kind, span, parent}); + return span; + } + + struct TSpanRecord { + std::string Name; + NTrace::ESpanKind Kind; + std::shared_ptr Span; + NTrace::ISpan* Parent = nullptr; + }; + + std::vector GetSpans() const { + std::lock_guard lock(Mutex_); + return Spans_; + } + + std::shared_ptr GetLastSpan() const { + std::lock_guard lock(Mutex_); + return Spans_.empty() ? nullptr : Spans_.back().Span; + } + + TSpanRecord GetLastSpanRecord() const { + std::lock_guard lock(Mutex_); + return Spans_.back(); + } + + size_t SpanCount() const { + std::lock_guard lock(Mutex_); + return Spans_.size(); + } + +private: + mutable std::mutex Mutex_; + std::vector Spans_; +}; + +class TFakeTraceProvider : public NTrace::ITraceProvider { +public: + std::shared_ptr GetTracer(const std::string& name) override { + std::lock_guard lock(Mutex_); + auto it = Tracers_.find(name); + if (it != Tracers_.end()) { + return it->second; + } + auto tracer = std::make_shared(); + Tracers_[name] = tracer; + return tracer; + } + + std::shared_ptr GetFakeTracer(const std::string& name) const { + std::lock_guard lock(Mutex_); + auto it = Tracers_.find(name); + return it != Tracers_.end() ? it->second : nullptr; + } + +private: + mutable std::mutex Mutex_; + std::map> Tracers_; +}; + +} // namespace NYdb::NTests diff --git a/tests/integration/metrics/CMakeLists.txt b/tests/integration/metrics/CMakeLists.txt index 6c9bb8b3abd..f63f5be0461 100644 --- a/tests/integration/metrics/CMakeLists.txt +++ b/tests/integration/metrics/CMakeLists.txt @@ -7,6 +7,7 @@ add_ydb_test(NAME metrics_it GTEST yutil YDB-CPP-SDK::Query client-metrics + impl-observability-error_category LABELS integration ) diff --git a/tests/integration/metrics/main.cpp b/tests/integration/metrics/main.cpp index 4668d5349ab..1db9fe28e1b 100644 --- a/tests/integration/metrics/main.cpp +++ b/tests/integration/metrics/main.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -65,10 +66,10 @@ std::shared_ptr GetDuration( {"db.namespace", dbNamespace}, {"db.operation.name", operation}, {"ydb.client.api", "Query"}, - {"db.response.status_code", ToString(status)}, }; if (status != EStatus::SUCCESS) { - labels["error.type"] = ToString(status); + labels["db.response.status_code"] = ToString(status); + labels["error.type"] = std::string(NObservability::CategorizeErrorType(status)); } return registry->GetHistogram("db.client.operation.duration", labels); } @@ -95,15 +96,15 @@ TEST(QueryMetricsIntegration, ExecuteQuerySuccessRecordsMetrics) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); - auto requests = GetCounter(registry, database, "db.client.operation.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, database, "db.client.operation.requests", "ydb.ExecuteQuery"); ASSERT_NE(requests, nullptr) << "ExecuteQuery request counter not created"; EXPECT_GE(requests->Get(), 1); - auto errors = GetCounter(registry, database, "db.client.operation.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, database, "db.client.operation.errors", "ydb.ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto duration = GetDuration(registry, database, "ExecuteQuery", EStatus::SUCCESS); + auto duration = GetDuration(registry, database, "ydb.ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "ExecuteQuery duration histogram not created"; EXPECT_GE(duration->Count(), 1u); for (double v : duration->GetValues()) { @@ -127,15 +128,15 @@ TEST(QueryMetricsIntegration, ExecuteQueryErrorRecordsErrorMetric) { ).ExtractValueSync(); EXPECT_NE(result.GetStatus(), EStatus::SUCCESS); - auto requests = GetCounter(registry, database, "db.client.operation.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, database, "db.client.operation.requests", "ydb.ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_GE(requests->Get(), 1); - auto errors = GetCounter(registry, database, "db.client.operation.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, database, "db.client.operation.errors", "ydb.ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_GE(errors->Get(), 1); - auto duration = GetDuration(registry, database, "ExecuteQuery", result.GetStatus()); + auto duration = GetDuration(registry, database, "ydb.ExecuteQuery", result.GetStatus()); ASSERT_NE(duration, nullptr); EXPECT_GE(duration->Count(), 1u); @@ -150,11 +151,11 @@ TEST(QueryMetricsIntegration, CreateSessionRecordsMetrics) { auto session = client.GetSession().ExtractValueSync(); ASSERT_TRUE(session.IsSuccess()) << session.GetIssues().ToString(); - auto requests = GetCounter(registry, database, "db.client.operation.requests", "GetSession"); + auto requests = GetCounter(registry, database, "db.client.operation.requests", "ydb.CreateSession"); ASSERT_NE(requests, nullptr) << "CreateSession request counter not created"; EXPECT_GE(requests->Get(), 1); - auto duration = GetDuration(registry, database, "GetSession", EStatus::SUCCESS); + auto duration = GetDuration(registry, database, "ydb.CreateSession", EStatus::SUCCESS); ASSERT_NE(duration, nullptr) << "CreateSession duration histogram not created"; EXPECT_GE(duration->Count(), 1u); @@ -184,11 +185,11 @@ TEST(QueryMetricsIntegration, CommitTransactionRecordsMetrics) { auto commitResult = execResult.GetTransaction()->Commit().ExtractValueSync(); ASSERT_TRUE(commitResult.IsSuccess()) << commitResult.GetIssues().ToString(); - auto commitRequests = GetCounter(registry, database, "db.client.operation.requests", "Commit"); + auto commitRequests = GetCounter(registry, database, "db.client.operation.requests", "ydb.Commit"); ASSERT_NE(commitRequests, nullptr) << "Commit request counter not created"; EXPECT_GE(commitRequests->Get(), 1); - auto commitDuration = GetDuration(registry, database, "Commit", EStatus::SUCCESS); + auto commitDuration = GetDuration(registry, database, "ydb.Commit", EStatus::SUCCESS); ASSERT_NE(commitDuration, nullptr); EXPECT_GE(commitDuration->Count(), 1u); } @@ -212,15 +213,15 @@ TEST(QueryMetricsIntegration, RollbackTransactionRecordsMetrics) { auto rollbackResult = tx.Rollback().ExtractValueSync(); ASSERT_TRUE(rollbackResult.IsSuccess()) << rollbackResult.GetIssues().ToString(); - auto rollbackRequests = GetCounter(registry, database, "db.client.operation.requests", "Rollback"); + auto rollbackRequests = GetCounter(registry, database, "db.client.operation.requests", "ydb.Rollback"); ASSERT_NE(rollbackRequests, nullptr) << "Rollback request counter not created"; EXPECT_GE(rollbackRequests->Get(), 1); - auto rollbackErrors = GetCounter(registry, database, "db.client.operation.errors", "Rollback"); + auto rollbackErrors = GetCounter(registry, database, "db.client.operation.errors", "ydb.Rollback"); ASSERT_NE(rollbackErrors, nullptr); EXPECT_EQ(rollbackErrors->Get(), 0); - auto rollbackDuration = GetDuration(registry, database, "Rollback", EStatus::SUCCESS); + auto rollbackDuration = GetDuration(registry, database, "ydb.Rollback", EStatus::SUCCESS); ASSERT_NE(rollbackDuration, nullptr); EXPECT_GE(rollbackDuration->Count(), 1u); @@ -245,15 +246,15 @@ TEST(QueryMetricsIntegration, MultipleQueriesAccumulateMetrics) { ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); } - auto requests = GetCounter(registry, database, "db.client.operation.requests", "ExecuteQuery"); + auto requests = GetCounter(registry, database, "db.client.operation.requests", "ydb.ExecuteQuery"); ASSERT_NE(requests, nullptr); EXPECT_EQ(requests->Get(), numQueries); - auto errors = GetCounter(registry, database, "db.client.operation.errors", "ExecuteQuery"); + auto errors = GetCounter(registry, database, "db.client.operation.errors", "ydb.ExecuteQuery"); ASSERT_NE(errors, nullptr); EXPECT_EQ(errors->Get(), 0); - auto duration = GetDuration(registry, database, "ExecuteQuery", EStatus::SUCCESS); + auto duration = GetDuration(registry, database, "ydb.ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); EXPECT_EQ(duration->Count(), static_cast(numQueries)); @@ -300,7 +301,7 @@ TEST(QueryMetricsIntegration, DurationValuesAreRealistic) { ).ExtractValueSync(); ASSERT_EQ(result.GetStatus(), EStatus::SUCCESS) << result.GetIssues().ToString(); - auto duration = GetDuration(registry, database, "ExecuteQuery", EStatus::SUCCESS); + auto duration = GetDuration(registry, database, "ydb.ExecuteQuery", EStatus::SUCCESS); ASSERT_NE(duration, nullptr); ASSERT_GE(duration->Count(), 1u); diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index be4e49bfc9f..d901992fefe 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -132,4 +132,20 @@ add_ydb_test(NAME client-ydb_metrics_ut GTEST client-metrics LABELS unit +) + +add_ydb_test(NAME client-ydb_spans_ut GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + observability/spans_ut.cpp + LINK_LIBRARIES + yutil + impl-observability + client-ydb_query-impl + client-ydb_table-impl + client-metrics + client-trace + LABELS + unit ) \ No newline at end of file diff --git a/tests/unit/client/observability/metrics_ut.cpp b/tests/unit/client/observability/metrics_ut.cpp index be073c3dcad..d342b0b4fa6 100644 --- a/tests/unit/client/observability/metrics_ut.cpp +++ b/tests/unit/client/observability/metrics_ut.cpp @@ -13,6 +13,10 @@ using namespace NYdb::NSdkStats; namespace { constexpr const char kTestDbNamespace[] = "/Root/testdb"; + + std::string YdbOp(const std::string& op) { + return op.rfind("ydb.", 0) == 0 ? op : "ydb." + op; + } } // namespace class RequestMetricsTest : public ::testing::Test { @@ -27,7 +31,7 @@ class RequestMetricsTest : public ::testing::Test { return Registry->GetCounter("db.client.operation.requests", { {"db.system.name", "ydb"}, {"db.namespace", kTestDbNamespace}, - {"db.operation.name", op}, + {"db.operation.name", YdbOp(op)}, {"ydb.client.api", "Unspecified"}, }); } @@ -36,7 +40,7 @@ class RequestMetricsTest : public ::testing::Test { return Registry->GetCounter("db.client.operation.errors", { {"db.system.name", "ydb"}, {"db.namespace", kTestDbNamespace}, - {"db.operation.name", op}, + {"db.operation.name", YdbOp(op)}, {"ydb.client.api", "Unspecified"}, }); } @@ -45,12 +49,12 @@ class RequestMetricsTest : public ::testing::Test { TLabels labels = { {"db.system.name", "ydb"}, {"db.namespace", kTestDbNamespace}, - {"db.operation.name", op}, + {"db.operation.name", YdbOp(op)}, {"ydb.client.api", "Unspecified"}, - {"db.response.status_code", ToString(status)}, }; if (status != EStatus::SUCCESS) { - labels["error.type"] = ToString(status); + labels["db.response.status_code"] = ToString(status); + labels["error.type"] = std::string(NObservability::CategorizeErrorType(status)); } return Registry->GetHistogram("db.client.operation.duration", labels); } @@ -222,7 +226,7 @@ TEST(RequestMetricsDbNamespaceTest, DifferentNamespacesAreSeparateMetricSeries) return NMetrics::TLabels{ {"db.system.name", "ydb"}, {"db.namespace", "/db/alpha"}, - {"db.operation.name", op}, + {"db.operation.name", YdbOp(op)}, {"ydb.client.api", "Unspecified"}, }; }; @@ -230,13 +234,13 @@ TEST(RequestMetricsDbNamespaceTest, DifferentNamespacesAreSeparateMetricSeries) return NMetrics::TLabels{ {"db.system.name", "ydb"}, {"db.namespace", "/db/beta"}, - {"db.operation.name", op}, + {"db.operation.name", YdbOp(op)}, {"ydb.client.api", "Unspecified"}, }; }; - auto reqAlpha = registry->GetCounter("db.client.operation.requests", labelsAlpha("GetSession")); - auto reqBeta = registry->GetCounter("db.client.operation.requests", labelsBeta("GetSession")); + auto reqAlpha = registry->GetCounter("db.client.operation.requests", labelsAlpha("ydb.GetSession")); + auto reqBeta = registry->GetCounter("db.client.operation.requests", labelsBeta("ydb.GetSession")); ASSERT_NE(reqAlpha, nullptr); ASSERT_NE(reqBeta, nullptr); EXPECT_EQ(reqAlpha->Get(), 1); @@ -244,18 +248,10 @@ TEST(RequestMetricsDbNamespaceTest, DifferentNamespacesAreSeparateMetricSeries) auto durAlpha = registry->GetHistogram( "db.client.operation.duration", - [&] { - auto l = labelsAlpha("GetSession"); - l["db.response.status_code"] = ToString(EStatus::SUCCESS); - return l; - }()); + labelsAlpha("ydb.GetSession")); auto durBeta = registry->GetHistogram( "db.client.operation.duration", - [&] { - auto l = labelsBeta("GetSession"); - l["db.response.status_code"] = ToString(EStatus::SUCCESS); - return l; - }()); + labelsBeta("ydb.GetSession")); ASSERT_NE(durAlpha, nullptr); ASSERT_NE(durBeta, nullptr); EXPECT_EQ(durAlpha->Count(), 1u); @@ -275,7 +271,7 @@ TEST(RequestMetricsClientAliasesTest, QueryOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteQuery"}, + {"db.operation.name", "ydb.ExecuteQuery"}, {"ydb.client.api", "Query"}, } ), @@ -287,7 +283,7 @@ TEST(RequestMetricsClientAliasesTest, QueryOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteQuery"}, + {"db.operation.name", "ydb.ExecuteQuery"}, {"ydb.client.api", "Query"}, } ), @@ -299,9 +295,8 @@ TEST(RequestMetricsClientAliasesTest, QueryOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteQuery"}, + {"db.operation.name", "ydb.ExecuteQuery"}, {"ydb.client.api", "Query"}, - {"db.response.status_code", ToString(EStatus::SUCCESS)}, } ), nullptr @@ -321,7 +316,7 @@ TEST(RequestMetricsClientAliasesTest, TableOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteDataQuery"}, + {"db.operation.name", "ydb.ExecuteDataQuery"}, {"ydb.client.api", "Table"}, } ), @@ -333,7 +328,7 @@ TEST(RequestMetricsClientAliasesTest, TableOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteDataQuery"}, + {"db.operation.name", "ydb.ExecuteDataQuery"}, {"ydb.client.api", "Table"}, } ), @@ -345,9 +340,8 @@ TEST(RequestMetricsClientAliasesTest, TableOperationsUseOtelStandardMetrics) { { {"db.system.name", "ydb"}, {"db.namespace", ""}, - {"db.operation.name", "ExecuteDataQuery"}, + {"db.operation.name", "ydb.ExecuteDataQuery"}, {"ydb.client.api", "Table"}, - {"db.response.status_code", ToString(EStatus::SUCCESS)}, } ), nullptr diff --git a/tests/unit/client/observability/spans_ut.cpp b/tests/unit/client/observability/spans_ut.cpp new file mode 100644 index 00000000000..bf0ec03a007 --- /dev/null +++ b/tests/unit/client/observability/spans_ut.cpp @@ -0,0 +1,498 @@ +#include +#include + +#include + +#include + +using namespace NYdb; +using namespace NYdb::NTests; + +namespace { + +constexpr const char kTestDbNamespace[] = "/Root/testdb"; + +struct TSpanTestParams { + std::string Name; + std::string ClientType; + + std::string ExecuteOp; + std::string ExecuteOpName; + + std::string CreateSessionOp; + std::string CreateSessionOpName; + + std::string CommitOp; + std::string CommitOpName; + + std::string RollbackOp; + std::string RollbackOpName; + + std::string RetryOp; + std::string RetryOpName; +}; + +} // namespace + +class SpanTest : public ::testing::TestWithParam { +protected: + void SetUp() override { + Tracer = std::make_shared(); + } + + std::shared_ptr MakeRequestSpan( + const std::string& operationName, + const std::string& endpoint, + NTrace::ESpanKind kind = NTrace::ESpanKind::CLIENT + ) { + return NYdb::NObservability::TRequestSpan::Create( + GetParam().ClientType, + Tracer, + operationName, + endpoint, + kTestDbNamespace, + TLog{}, + kind + ); + } + + std::shared_ptr Tracer; +}; + +TEST_P(SpanTest, SpanNameFormat) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + ASSERT_EQ(Tracer->SpanCount(), 1u); + EXPECT_EQ(Tracer->GetLastSpanRecord().Name, p.ExecuteOpName); +} + +TEST_P(SpanTest, SpanKindIsClient) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.CreateSessionOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + ASSERT_EQ(Tracer->SpanCount(), 1u); + EXPECT_EQ(Tracer->GetLastSpanRecord().Kind, NTrace::ESpanKind::CLIENT); +} + +TEST_P(SpanTest, DbSystemAttribute) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("db.system.name"), "ydb"); +} + +TEST_P(SpanTest, DbNamespaceAndClientApi) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("db.namespace"), kTestDbNamespace); + EXPECT_EQ(fakeSpan->GetStringAttribute("ydb.client.api"), p.ClientType); + EXPECT_EQ(fakeSpan->GetStringAttribute("db.operation.name"), p.ExecuteOpName); +} + +TEST_P(SpanTest, ServerAddressAndPort) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.CommitOp, "ydb.server:2135"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "ydb.server"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_P(SpanTest, ServerAddressCustomPort) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.RollbackOp, "myhost:9090"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "myhost"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 9090); +} + +TEST_P(SpanTest, ServerAddressNoPortDefaultsTo2135) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "myhost"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "myhost"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_P(SpanTest, IPv6EndpointParsing) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "[::1]:2136"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "::1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2136); +} + +TEST_P(SpanTest, IPv6EndpointNoPort) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "[fe80::1]"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("server.address"), "fe80::1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("server.port"), 2135); +} + +TEST_P(SpanTest, PeerEndpointAttributes) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "discovery.ydb:2135"); + span->SetPeerEndpoint("10.0.0.1:2136"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("network.peer.address"), "10.0.0.1"); + EXPECT_EQ(fakeSpan->GetIntAttribute("network.peer.port"), 2136); +} + +TEST_P(SpanTest, SuccessStatusDoesNotSetErrorAttrs) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.CommitOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_FALSE(fakeSpan->HasStringAttribute("db.response.status_code")); + EXPECT_FALSE(fakeSpan->HasStringAttribute("error.type")); +} + +TEST_P(SpanTest, ErrorStatusSetsErrorType) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.RollbackOp, "localhost:2135"); + span->End(EStatus::UNAVAILABLE); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ(fakeSpan->GetStringAttribute("db.response.status_code"), "UNAVAILABLE"); + EXPECT_TRUE(fakeSpan->HasStringAttribute("error.type")); + EXPECT_FALSE(fakeSpan->GetStringAttribute("error.type").empty()); +} + +TEST_P(SpanTest, SpanIsEndedAfterEnd) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + + EXPECT_FALSE(fakeSpan->IsEnded()); + span->End(EStatus::SUCCESS); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_P(SpanTest, NullTracerDoesNotCrash) { + const auto& p = GetParam(); + EXPECT_NO_THROW({ + auto span = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, + nullptr, + p.ExecuteOp, + "localhost:2135", + kTestDbNamespace, + TLog{} + ); + span->SetPeerEndpoint("10.0.0.1:2136"); + span->AddEvent("retry", {{"attempt", "1"}}); + span->End(EStatus::SUCCESS); + }); +} + +TEST_P(SpanTest, DestructorEndsSpan) { + const auto& p = GetParam(); + auto fakeSpan = [&]() -> std::shared_ptr { + auto span = MakeRequestSpan(p.CreateSessionOp, "localhost:2135"); + return Tracer->GetLastSpan(); + }(); + + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_P(SpanTest, ExplicitEndThenDestructorDoesNotDoubleEnd) { + const auto& p = GetParam(); + auto fakeSpan = [&]() -> std::shared_ptr { + auto span = MakeRequestSpan(p.CommitOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + return Tracer->GetLastSpan(); + }(); + + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsEnded()); +} + +TEST_P(SpanTest, AddEventForwarded) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->AddEvent("retry", {{"ydb.attempt", "2"}, {"error.type", "UNAVAILABLE"}}); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + auto events = fakeSpan->GetEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events[0].Name, "retry"); + EXPECT_EQ(events[0].Attributes.at("ydb.attempt"), "2"); + EXPECT_EQ(events[0].Attributes.at("error.type"), "UNAVAILABLE"); +} + +TEST_P(SpanTest, EmptyPeerEndpointIgnored) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.CreateSessionOp, "localhost:2135"); + span->SetPeerEndpoint(""); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_FALSE(fakeSpan->HasStringAttribute("network.peer.address")); + EXPECT_FALSE(fakeSpan->HasIntAttribute("network.peer.port")); +} + +TEST_P(SpanTest, OperationNamesAreNormalized) { + const auto& p = GetParam(); + const std::vector> ops = { + {p.CreateSessionOp, p.CreateSessionOpName}, + {p.ExecuteOp, p.ExecuteOpName}, + {p.CommitOp, p.CommitOpName}, + {p.RollbackOp, p.RollbackOpName}, + }; + + for (const auto& [op, _] : ops) { + auto span = MakeRequestSpan(op, "localhost:2135"); + span->End(EStatus::SUCCESS); + } + + auto spans = Tracer->GetSpans(); + ASSERT_EQ(spans.size(), ops.size()); + for (size_t i = 0; i < ops.size(); ++i) { + EXPECT_EQ(spans[i].Name, ops[i].second); + EXPECT_EQ(spans[i].Kind, NTrace::ESpanKind::CLIENT); + } +} + +TEST_P(SpanTest, MultipleErrorStatuses) { + const auto& p = GetParam(); + const std::vector> errorStatuses = { + {EStatus::BAD_REQUEST, "ydb_error"}, + {EStatus::UNAUTHORIZED, "ydb_error"}, + {EStatus::INTERNAL_ERROR, "ydb_error"}, + {EStatus::UNAVAILABLE, "ydb_error"}, + {EStatus::OVERLOADED, "ydb_error"}, + {EStatus::TIMEOUT, "ydb_error"}, + {EStatus::NOT_FOUND, "ydb_error"}, + {EStatus::CLIENT_INTERNAL_ERROR, "transport_error"}, + }; + + for (const auto& [status, expectedErrorType] : errorStatuses) { + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->End(status); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_EQ( + fakeSpan->GetStringAttribute("db.response.status_code"), + ToString(status) + ); + EXPECT_EQ( + fakeSpan->GetStringAttribute("error.type"), + expectedErrorType + ) << "wrong error.type category for status " << static_cast(status); + } +} + +TEST_P(SpanTest, EmptyEndpointDoesNotCrash) { + const auto& p = GetParam(); + EXPECT_NO_THROW({ + auto span = MakeRequestSpan(p.ExecuteOp, ""); + span->End(EStatus::SUCCESS); + }); +} + +TEST_P(SpanTest, ActivateReturnsScope) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.RetryOp, "localhost:2135"); + auto scope = span->Activate(); + EXPECT_NE(scope, nullptr); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + EXPECT_TRUE(fakeSpan->IsActivated()); + + span->End(EStatus::SUCCESS); +} + +TEST_P(SpanTest, ActivateNullTracerReturnsNull) { + const auto& p = GetParam(); + auto span = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, + nullptr, + p.RetryOp, + "localhost:2135", + kTestDbNamespace, + TLog{} + ); + auto scope = span->Activate(); + EXPECT_EQ(scope, nullptr); +} + +TEST_P(SpanTest, InternalSpanKindIsPropagated) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.RetryOp, "localhost:2135", NTrace::ESpanKind::INTERNAL); + span->End(EStatus::SUCCESS); + + ASSERT_EQ(Tracer->SpanCount(), 1u); + EXPECT_EQ(Tracer->GetLastSpanRecord().Name, p.RetryOpName); + EXPECT_EQ(Tracer->GetLastSpanRecord().Kind, NTrace::ESpanKind::INTERNAL); +} + +TEST_P(SpanTest, ErrorStatusAddsExceptionEvent) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->End(EStatus::UNAVAILABLE); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + auto events = fakeSpan->GetEvents(); + ASSERT_FALSE(events.empty()); + + bool found = false; + for (const auto& event : events) { + if (event.Name == "exception") { + found = true; + EXPECT_EQ(event.Attributes.at("exception.type"), "UNAVAILABLE"); + EXPECT_EQ(event.Attributes.at("exception.message"), "UNAVAILABLE"); + } + } + EXPECT_TRUE(found) << "expected an 'exception' event on a failed span"; +} + +TEST_P(SpanTest, SuccessStatusNoExceptionEvent) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.CommitOp, "localhost:2135"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + for (const auto& event : fakeSpan->GetEvents()) { + EXPECT_NE(event.Name, "exception"); + } +} + +TEST_P(SpanTest, RecordExceptionEmitsEvent) { + const auto& p = GetParam(); + auto span = MakeRequestSpan(p.ExecuteOp, "localhost:2135"); + span->RecordException("TimeoutException", "operation timed out"); + span->End(EStatus::SUCCESS); + + auto fakeSpan = Tracer->GetLastSpan(); + ASSERT_NE(fakeSpan, nullptr); + auto events = fakeSpan->GetEvents(); + ASSERT_FALSE(events.empty()); + bool found = false; + for (const auto& event : events) { + if (event.Name == "exception" + && event.Attributes.count("exception.type") + && event.Attributes.at("exception.type") == "TimeoutException") + { + found = true; + EXPECT_EQ(event.Attributes.at("exception.message"), "operation timed out"); + } + } + EXPECT_TRUE(found); +} + +TEST_P(SpanTest, ExplicitParentIsPropagatedToTracer) { + const auto& p = GetParam(); + + auto parent = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, Tracer, p.RetryOp, "localhost:2135", kTestDbNamespace, TLog{}, NTrace::ESpanKind::INTERNAL); + ASSERT_NE(parent, nullptr); + auto parentRecord = Tracer->GetLastSpanRecord(); + auto* parentRaw = parentRecord.Span.get(); + + auto child = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, Tracer, p.ExecuteOp, "localhost:2135", kTestDbNamespace, + TLog{}, NTrace::ESpanKind::CLIENT, parent); + ASSERT_NE(child, nullptr); + auto childRecord = Tracer->GetLastSpanRecord(); + + EXPECT_EQ(childRecord.Parent, parentRaw) + << "child span must receive parent pointer through ITracer::StartSpan"; +} + +TEST_P(SpanTest, RetryAttemptParentedToRoot) { + const auto& p = GetParam(); + + auto parent = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, Tracer, p.RetryOp, "localhost:2135", kTestDbNamespace, TLog{}, NTrace::ESpanKind::INTERNAL); + ASSERT_NE(parent, nullptr); + auto parentRecord = Tracer->GetLastSpanRecord(); + auto* parentRaw = parentRecord.Span.get(); + + // Imitating what retry contexts do: create attempt span with explicit parent. + auto attempt = NYdb::NObservability::TRequestSpan::Create( + p.ClientType, Tracer, "ydb.Try", "localhost:2135", kTestDbNamespace, + TLog{}, NTrace::ESpanKind::INTERNAL, parent); + ASSERT_NE(attempt, nullptr); + auto attemptRecord = Tracer->GetLastSpanRecord(); + + EXPECT_EQ(attemptRecord.Parent, parentRaw) + << "retry attempt span must be parented to retry root span explicitly"; +} + +INSTANTIATE_TEST_SUITE_P( + Clients, + SpanTest, + ::testing::Values( + TSpanTestParams{ + /*Name=*/ "Query", + /*ClientType=*/ "Query", + /*ExecuteOp=*/ "ExecuteQuery", + /*ExecuteOpName=*/ "ydb.ExecuteQuery", + /*CreateSessionOp=*/ "CreateSession", + /*CreateSessionOpName=*/"ydb.CreateSession", + /*CommitOp=*/ "Commit", + /*CommitOpName=*/ "ydb.Commit", + /*RollbackOp=*/ "Rollback", + /*RollbackOpName=*/ "ydb.Rollback", + /*RetryOp=*/ "ydb.RunWithRetry", + /*RetryOpName=*/ "ydb.RunWithRetry" + }, + TSpanTestParams{ + /*Name=*/ "Table", + /*ClientType=*/ "Table", + /*ExecuteOp=*/ "ydb.ExecuteDataQuery", + /*ExecuteOpName=*/ "ydb.ExecuteDataQuery", + /*CreateSessionOp=*/ "ydb.CreateSession", + /*CreateSessionOpName=*/"ydb.CreateSession", + /*CommitOp=*/ "ydb.Commit", + /*CommitOpName=*/ "ydb.Commit", + /*RollbackOp=*/ "ydb.Rollback", + /*RollbackOpName=*/ "ydb.Rollback", + /*RetryOp=*/ "ydb.RunWithRetry", + /*RetryOpName=*/ "ydb.RunWithRetry" + } + ), + [](const ::testing::TestParamInfo& info) { + return info.param.Name; + } +); From 6fb7b3d75c093f7f8ad4aabb5c1c26c0ef9e6130 Mon Sep 17 00:00:00 2001 From: maladetska Date: Wed, 29 Apr 2026 01:51:17 +0300 Subject: [PATCH 93/93] update --- src/client/query/client.cpp | 2 +- src/client/table/impl/table_client.h | 2 +- src/client/table/table.cpp | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 201ba27cc3e..448e50b147e 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -19,7 +19,7 @@ #include #include -#include +#include #include diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 54fe07bf00c..8bddced9b79 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -12,7 +12,7 @@ #include #include #include -#include +#include #include "client_session.h" #include "data_query.h" diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 298c875b99c..4f6e411bf73 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -9,8 +9,8 @@ #include #undef INCLUDE_YDB_INTERNAL_H -#include -#include +#include +#include #include #include #include