diff --git a/.github/import_generation.txt b/.github/import_generation.txt index 9902f17848a..f04c001f3f7 100644 --- a/.github/import_generation.txt +++ b/.github/import_generation.txt @@ -1 +1 @@ -28 +29 diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 74c170ce520..b766f8b3117 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -42606baa568dd1f0408b0651dce31577f28bcfe1 +eaa66fe4cc75727a60d89a9ec9d99970e03ce0e0 diff --git a/include/ydb-cpp-sdk/client/monitoring/monitoring.h b/include/ydb-cpp-sdk/client/monitoring/monitoring.h index 54ce3211844..8de502cf363 100644 --- a/include/ydb-cpp-sdk/client/monitoring/monitoring.h +++ b/include/ydb-cpp-sdk/client/monitoring/monitoring.h @@ -5,6 +5,7 @@ namespace Ydb { namespace Monitoring { class SelfCheckResult; + class ClusterStateResult; } } @@ -34,6 +35,12 @@ struct TSelfCheckSettings : public TOperationRequestSettings FLUENT_SETTING_OPTIONAL(uint32_t, MaximumLevel); }; +struct TClusterStateSettings : public TOperationRequestSettings { + FLUENT_SETTING_OPTIONAL(uint32_t, DurationSeconds); + FLUENT_SETTING_OPTIONAL(uint32_t, PeriodSeconds); +}; + + class TSelfCheckResult : public TStatus { friend class NYdb::TProtoAccessor; public: @@ -45,6 +52,17 @@ class TSelfCheckResult : public TStatus { using TAsyncSelfCheckResult = NThreading::TFuture; +class TClusterStateResult : public TStatus { + friend class NYdb::TProtoAccessor; +public: + TClusterStateResult(TStatus&& status, Ydb::Monitoring::ClusterStateResult&& result); +private: + class TImpl; + std::shared_ptr Impl_; +}; + +using TAsyncClusterStateResult = NThreading::TFuture; + class TMonitoringClient { class TImpl; @@ -52,6 +70,8 @@ class TMonitoringClient { TMonitoringClient(const TDriver& driver, const TCommonClientSettings& settings = TCommonClientSettings()); TAsyncSelfCheckResult SelfCheck(const TSelfCheckSettings& settings = TSelfCheckSettings()); + + TAsyncClusterStateResult ClusterState(const TClusterStateSettings& settings = TClusterStateSettings()); private: std::shared_ptr Impl_; }; diff --git a/include/ydb-cpp-sdk/client/proto/accessor.h b/include/ydb-cpp-sdk/client/proto/accessor.h index b13dd6a1dbd..47cbea7c65b 100644 --- a/include/ydb-cpp-sdk/client/proto/accessor.h +++ b/include/ydb-cpp-sdk/client/proto/accessor.h @@ -49,6 +49,7 @@ class TProtoAccessor { static const Ydb::Topic::DescribeTopicResult& GetProto(const NYdb::NTopic::TTopicDescription& topicDescription); static const Ydb::Topic::DescribeConsumerResult& GetProto(const NYdb::NTopic::TConsumerDescription& consumerDescription); static const Ydb::Monitoring::SelfCheckResult& GetProto(const NYdb::NMonitoring::TSelfCheckResult& selfCheckResult); + static const Ydb::Monitoring::ClusterStateResult& GetProto(const NYdb::NMonitoring::TClusterStateResult& clusterStateResult); static const Ydb::Coordination::DescribeNodeResult& GetProto(const NYdb::NCoordination::TNodeDescription& describeNodeResult); static const Ydb::Import::ListObjectsInS3ExportResult& GetProto(const NYdb::NImport::TListObjectsInS3ExportResult& result); diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 85d6a3ff648..16c9b703f33 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -31,6 +31,7 @@ class ExplicitPartitions; class GlobalIndexSettings; class VectorIndexSettings; class KMeansTreeSettings; +class FulltextIndexSettings; class PartitioningSettings; class ReadReplicasSettings; class DateTypeColumnModeSettings; @@ -295,6 +296,49 @@ struct TKMeansTreeSettings { void Out(IOutputStream &o) const; }; +struct TFulltextIndexSettings { +public: + enum class ELayout { + Unspecified = 0, + Flat, + }; + + enum class ETokenizer { + Unspecified = 0, + Whitespace, + Standard, + Keyword, + }; + + struct TAnalyzers { + std::optional Tokenizer; + std::optional Language; + std::optional UseFilterLowercase; + std::optional UseFilterStopwords; + std::optional UseFilterNgram; + std::optional UseFilterEdgeNgram; + std::optional FilterNgramMinLength; + std::optional FilterNgramMaxLength; + std::optional UseFilterLength; + std::optional FilterLengthMin; + std::optional FilterLengthMax; + }; + + struct TColumnAnalyzers { + std::optional Column; + std::optional Analyzers; + }; + + std::optional Layout; + std::vector Columns; + + static TFulltextIndexSettings FromProto(const Ydb::Table::FulltextIndexSettings& proto); + + void SerializeTo(Ydb::Table::FulltextIndexSettings& settings) const; + + void Out(IOutputStream& o) const; +}; + //! Represents index description class TIndexDescription { friend class NYdb::TProtoAccessor; @@ -306,7 +350,7 @@ class TIndexDescription { const std::vector& indexColumns, const std::vector& dataColumns = {}, const std::vector& globalIndexSettings = {}, - const std::variant& specializedIndexSettings = {} + const std::variant& specializedIndexSettings = {} ); TIndexDescription( @@ -320,7 +364,7 @@ class TIndexDescription { EIndexType GetIndexType() const; const std::vector& GetIndexColumns() const; const std::vector& GetDataColumns() const; - const std::variant& GetIndexSettings() const; + const std::variant& GetIndexSettings() const; uint64_t GetSizeBytes() const; void SerializeTo(Ydb::Table::TableIndex& proto) const; @@ -340,7 +384,7 @@ class TIndexDescription { std::vector IndexColumns_; std::vector DataColumns_; std::vector GlobalIndexSettings_; - std::variant SpecializedIndexSettings_; + std::variant SpecializedIndexSettings_; uint64_t SizeBytes_ = 0; }; @@ -755,6 +799,9 @@ class TTableDescription { // vector KMeansTree 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); // default void AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); @@ -996,6 +1043,10 @@ class TTableBuilder { TTableBuilder& AddVectorKMeansTreeIndex(const std::string& indexName, const std::vector& indexColumns, const TKMeansTreeSettings& indexSettings); TTableBuilder& AddVectorKMeansTreeIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TKMeansTreeSettings& indexSettings); + // 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); + // default TTableBuilder& AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); TTableBuilder& 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 5d6aa6111f9..ff6b141d68b 100644 --- a/include/ydb-cpp-sdk/client/table/table_enum.h +++ b/include/ydb-cpp-sdk/client/table/table_enum.h @@ -35,6 +35,7 @@ enum class EIndexType { GlobalAsync, GlobalUnique, GlobalVectorKMeansTree, + GlobalFulltext, Unknown = std::numeric_limits::max() }; diff --git a/include/ydb-cpp-sdk/client/types/request_settings.h b/include/ydb-cpp-sdk/client/types/request_settings.h index 8d481726e3c..51f3b40307c 100644 --- a/include/ydb-cpp-sdk/client/types/request_settings.h +++ b/include/ydb-cpp-sdk/client/types/request_settings.h @@ -20,7 +20,7 @@ struct TRequestSettings { FLUENT_SETTING(std::string, TraceId); FLUENT_SETTING(std::string, RequestType); FLUENT_SETTING(THeader, Header); - FLUENT_SETTING(TDuration, ClientTimeout); + FLUENT_SETTING_DEFAULT(TDuration, ClientTimeout, TDuration::Max()); FLUENT_SETTING(std::string, TraceParent); TRequestSettings() = default; diff --git a/src/api/grpc/ydb_monitoring_v1.proto b/src/api/grpc/ydb_monitoring_v1.proto index d14b503754b..f67a3718e1b 100644 --- a/src/api/grpc/ydb_monitoring_v1.proto +++ b/src/api/grpc/ydb_monitoring_v1.proto @@ -10,4 +10,6 @@ service MonitoringService { rpc SelfCheck(Monitoring.SelfCheckRequest) returns (Monitoring.SelfCheckResponse); // Checks current node health rpc NodeCheck(Monitoring.NodeCheckRequest) returns (Monitoring.NodeCheckResponse); + // Get cluster state + rpc ClusterState(Monitoring.ClusterStateRequest) returns (Monitoring.ClusterStateResponse); } diff --git a/src/api/protos/ydb_monitoring.proto b/src/api/protos/ydb_monitoring.proto index 6d652244412..58e6e2c2164 100644 --- a/src/api/protos/ydb_monitoring.proto +++ b/src/api/protos/ydb_monitoring.proto @@ -182,12 +182,19 @@ message LocationBridgePile { string name = 1; } +message LocationStateStorage { + uint32 ring = 1; + LocationNode node = 2; + LocationBridgePile pile = 3; +} + message LocationCompute { LocationNode node = 1; LocationComputePool pool = 2; LocationComputeTablet tablet = 3; LocationComputeSchema schema = 4; LocationBridgePile pile = 5; + LocationStateStorage state_storage = 6; } message LocationDatabase { @@ -227,3 +234,18 @@ message SelfCheckResult { repeated DatabaseStatus database_status = 3; LocationNode location = 4; } + +message ClusterStateRequest { + Ydb.Operations.OperationParams operation_params = 1; // basic operation params, including timeout + uint32 duration_seconds = 2; + uint32 period_seconds = 3; +} + +message ClusterStateResponse { + // After successfull completion must contain ClusterStateResult. + Ydb.Operations.Operation operation = 1; +} + +message ClusterStateResult { + string result = 1; +} diff --git a/src/api/protos/ydb_scheme.proto b/src/api/protos/ydb_scheme.proto index 2273746bfc8..d4c98313654 100644 --- a/src/api/protos/ydb_scheme.proto +++ b/src/api/protos/ydb_scheme.proto @@ -68,7 +68,7 @@ message Entry { RESOURCE_POOL = 21; TRANSFER = 23; SYS_VIEW = 24; - STREAMING_QUERY = 25; + STREAMING_QUERY = 26; } // Name of scheme entry (dir2 of /dir1/dir2) diff --git a/src/client/common_client/impl/CMakeLists.txt b/src/client/common_client/impl/CMakeLists.txt index 2be569fc25b..26b4096b7ac 100644 --- a/src/client/common_client/impl/CMakeLists.txt +++ b/src/client/common_client/impl/CMakeLists.txt @@ -3,6 +3,7 @@ _ydb_sdk_add_library(client-ydb_common_client-impl) target_link_libraries(client-ydb_common_client-impl PUBLIC yutil impl-internal-grpc_connections + impl-internal-rpc_request_settings ) target_sources(client-ydb_common_client-impl PRIVATE diff --git a/src/client/common_client/impl/client.h b/src/client/common_client/impl/client.h index da8c8306dbe..1ce6be40afa 100644 --- a/src/client/common_client/impl/client.h +++ b/src/client/common_client/impl/client.h @@ -53,7 +53,7 @@ class TClientImplCommon return DbDriverState_->DiscoveryCompleted(); } - void ScheduleTask(const std::function& fn, TDuration timeout) override { + void ScheduleTask(const std::function& fn, TDeadline::Duration timeout) override { std::weak_ptr weak = this->shared_from_this(); auto cbGuard = [weak, fn]() { auto strongClient = weak.lock(); diff --git a/src/client/common_client/impl/iface.h b/src/client/common_client/impl/iface.h index 87b6ad6472a..5fa14e50f12 100644 --- a/src/client/common_client/impl/iface.h +++ b/src/client/common_client/impl/iface.h @@ -1,14 +1,15 @@ #pragma once +#include + #include -#include namespace NYdb::inline V3 { class IClientImplCommon { public: virtual ~IClientImplCommon() = default; - virtual void ScheduleTask(const std::function& fn, TDuration timeout) = 0; + virtual void ScheduleTask(const std::function& fn, TDeadline::Duration timeout) = 0; }; } diff --git a/src/client/driver/driver.cpp b/src/client/driver/driver.cpp index e1e88dea226..3c8e8f20b96 100644 --- a/src/client/driver/driver.cpp +++ b/src/client/driver/driver.cpp @@ -265,9 +265,9 @@ TDriverConfig TDriver::GetConfig() const { ); config.SetDrainOnDtors(Impl_->DrainOnDtors_); config.SetBalancingPolicy(std::make_unique(Impl_->BalancingSettings_)); - config.SetGRpcKeepAliveTimeout(Impl_->GRpcKeepAliveTimeout_); + config.SetGRpcKeepAliveTimeout(std::chrono::duration_cast(Impl_->GRpcKeepAliveTimeout_)); config.SetGRpcKeepAlivePermitWithoutCalls(Impl_->GRpcKeepAlivePermitWithoutCalls_); - config.SetSocketIdleTimeout(Impl_->SocketIdleTimeout_); + config.SetSocketIdleTimeout(std::chrono::duration_cast(Impl_->SocketIdleTimeout_)); config.SetMaxInboundMessageSize(Impl_->MaxInboundMessageSize_); config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_); config.SetMaxMessageSize(Impl_->MaxMessageSize_); diff --git a/src/client/federated_topic/impl/federation_observer.cpp b/src/client/federated_topic/impl/federation_observer.cpp index 2085d064ac3..437e88c4ee7 100644 --- a/src/client/federated_topic/impl/federation_observer.cpp +++ b/src/client/federated_topic/impl/federation_observer.cpp @@ -13,9 +13,9 @@ TFederatedDbObserverImpl::TFederatedDbObserverImpl(std::shared_ptrRunDeferred( @@ -90,7 +93,7 @@ void TFederatedDbObserverImpl::RunFederationDiscoveryImpl() { &Ydb::FederationDiscovery::V1::FederationDiscoveryService::Stub::AsyncListFederationDatabases, DbDriverState_, {}, // no polling unready operations, so no need in delay parameter - RpcSettings, + settings, FederationDiscoveryDelayContext); } diff --git a/src/client/federated_topic/impl/federation_observer.h b/src/client/federated_topic/impl/federation_observer.h index 4499782c6ca..3eddb068768 100644 --- a/src/client/federated_topic/impl/federation_observer.h +++ b/src/client/federated_topic/impl/federation_observer.h @@ -84,6 +84,7 @@ class TFederatedDbObserverImpl : public TClientImplCommon FederatedDbState; NThreading::TPromise PromiseToInitState; TRpcRequestSettings RpcSettings; + TDeadline::Duration ConnectionTimeout; TSpinLock Lock; NTopic::IRetryPolicy::TPtr FederationDiscoveryRetryPolicy; diff --git a/src/client/impl/internal/CMakeLists.txt b/src/client/impl/internal/CMakeLists.txt index 5370c34f570..6fd29a5c4a2 100644 --- a/src/client/impl/internal/CMakeLists.txt +++ b/src/client/impl/internal/CMakeLists.txt @@ -4,6 +4,7 @@ add_subdirectory(grpc_connections) add_subdirectory(logger) add_subdirectory(make_request) add_subdirectory(plain_status) +add_subdirectory(rpc_request_settings) add_subdirectory(retry) add_subdirectory(thread_pool) add_subdirectory(value_helpers) diff --git a/src/client/impl/internal/db_driver_state/state.cpp b/src/client/impl/internal/db_driver_state/state.cpp index ffaaf2f0e11..9a836fc527a 100644 --- a/src/client/impl/internal/db_driver_state/state.cpp +++ b/src/client/impl/internal/db_driver_state/state.cpp @@ -20,9 +20,11 @@ namespace { namespace NYdb::inline V3 { +using namespace std::chrono_literals; + constexpr int PESSIMIZATION_DISCOVERY_THRESHOLD = 50; // percent of endpoints pessimized by transport error to start recheck constexpr TDuration ENDPOINT_UPDATE_PERIOD = TDuration::Minutes(1); // period to perform endpoints update in "normal" case -constexpr TDuration DISCOVERY_RECHECK_PERIOD = TDuration::Seconds(5); // period to run periodic discovery task +constexpr TDeadline::Duration DISCOVERY_RECHECK_PERIOD = 5s; // period to run periodic discovery task TDbDriverState::TDbDriverState( const std::string& database, @@ -237,7 +239,7 @@ TDbDriverStatePtr TDbDriverStateTracker::GetDriverState( return strongState; } -void TDbDriverState::AddPeriodicTask(TPeriodicCb&& cb, TDuration period) { +void TDbDriverState::AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) { Client->AddPeriodicTask(std::move(cb), period); } diff --git a/src/client/impl/internal/db_driver_state/state.h b/src/client/impl/internal/db_driver_state/state.h index ce72760bf65..db7ceef77a0 100644 --- a/src/client/impl/internal/db_driver_state/state.h +++ b/src/client/impl/internal/db_driver_state/state.h @@ -39,7 +39,7 @@ class TDbDriverState void SignalDiscoveryCompleted(); - void AddPeriodicTask(TPeriodicCb&& cb, TDuration period) override; + void AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) override; void AddCb(TCb&& cb, ENotifyType type); void ForEachEndpoint(const TEndpointElectorSafe::THandleCb& cb, const void* tag) const; diff --git a/src/client/impl/internal/grpc_connections/CMakeLists.txt b/src/client/impl/internal/grpc_connections/CMakeLists.txt index 052bd17505d..a2de8250f3c 100644 --- a/src/client/impl/internal/grpc_connections/CMakeLists.txt +++ b/src/client/impl/internal/grpc_connections/CMakeLists.txt @@ -6,6 +6,7 @@ target_link_libraries(impl-internal-grpc_connections PUBLIC api-protos impl-internal-db_driver_state impl-internal-plain_status + impl-internal-rpc_request_settings client-impl-ydb_stats client-resources client-types-exceptions diff --git a/src/client/impl/internal/grpc_connections/actions.cpp b/src/client/impl/internal/grpc_connections/actions.cpp index 196f40d38f7..e7f8d44ad0e 100644 --- a/src/client/impl/internal/grpc_connections/actions.cpp +++ b/src/client/impl/internal/grpc_connections/actions.cpp @@ -4,9 +4,11 @@ #include +using namespace std::chrono_literals; + namespace NYdb::inline V3 { -constexpr TDuration MAX_DEFERRED_CALL_DELAY = TDuration::Seconds(10); // The max delay between GetOperation calls for one operation +constexpr TDeadline::Duration MAX_DEFERRED_CALL_DELAY = 10s; // The max delay between GetOperation calls for one operation TSimpleCbResult::TSimpleCbResult( TSimpleCb&& cb, @@ -24,18 +26,18 @@ TDeferredAction::TDeferredAction(const std::string& operationId, TDeferredOperationCb&& userCb, TGRpcConnectionsImpl* connection, std::shared_ptr context, - TDuration delay, + TDeadline::Duration delay, + TDeadline globalDeadline, TDbDriverStatePtr dbState, const std::string& endpoint) : TAlarmActionBase(std::move(userCb), connection, std::move(context)) - , NextDelay_(Min(delay * 2, MAX_DEFERRED_CALL_DELAY)) + , NextDelay_(std::min(delay * 2, MAX_DEFERRED_CALL_DELAY)) + , GlobalDeadline_(globalDeadline) , DbDriverState_(dbState) , OperationId_(operationId) , Endpoint_(endpoint) { - Deadline_ = gpr_time_add( - gpr_now(GPR_CLOCK_MONOTONIC), - gpr_time_from_micros(delay.MicroSeconds(), GPR_TIMESPAN)); + Deadline_ = std::min(GlobalDeadline_, TDeadline::AfterDuration(delay)); } void TDeferredAction::OnAlarm() { @@ -46,7 +48,8 @@ void TDeferredAction::OnAlarm() { TRpcRequestSettings settings; settings.PreferredEndpoint = TEndpointKey(Endpoint_, 0); - + settings.Deadline = GlobalDeadline_; + Connection_->RunDeferred( std::move(getOperationRequest), std::move(UserResponseCb_), @@ -56,7 +59,7 @@ void TDeferredAction::OnAlarm() { settings, true, std::move(Context_)); - } +} void TDeferredAction::OnError() { Y_ABORT_UNLESS(Connection_); @@ -76,13 +79,11 @@ TPeriodicAction::TPeriodicAction( TPeriodicCb&& userCb, TGRpcConnectionsImpl* connection, std::shared_ptr context, - TDuration period) + TDeadline::Duration period) : TAlarmActionBase(std::move(userCb), connection, std::move(context)) , Period_(period) { - Deadline_ = gpr_time_add( - gpr_now(GPR_CLOCK_MONOTONIC), - gpr_time_from_micros(Period_.MicroSeconds(), GPR_TIMESPAN)); + Deadline_ = TDeadline::AfterDuration(period); } void TPeriodicAction::OnAlarm() { diff --git a/src/client/impl/internal/grpc_connections/actions.h b/src/client/impl/internal/grpc_connections/actions.h index e29f96fa962..1c5b46be8df 100644 --- a/src/client/impl/internal/grpc_connections/actions.h +++ b/src/client/impl/internal/grpc_connections/actions.h @@ -100,7 +100,7 @@ class TAlarmActionBase } protected: - gpr_timespec Deadline_ = {}; + TDeadline Deadline_; private: std::mutex Mutex_; @@ -200,7 +200,8 @@ class TDeferredAction TDeferredOperationCb&& userCb, TGRpcConnectionsImpl* connection, std::shared_ptr context, - TDuration timeout, + TDeadline::Duration delay, + TDeadline globalDeadline, TDbDriverStatePtr dbState, const std::string& endpoint); @@ -208,7 +209,9 @@ class TDeferredAction void OnError() override; private: - TDuration NextDelay_; + TDeadline::Duration NextDelay_; + TDeadline GlobalDeadline_; + TDbDriverStatePtr DbDriverState_; const std::string OperationId_; const std::string Endpoint_; @@ -222,12 +225,12 @@ class TPeriodicAction TPeriodicCb&& userCb, TGRpcConnectionsImpl* connection, std::shared_ptr context, - TDuration period); + TDeadline::Duration period); void OnAlarm() override; void OnError() override; private: - TDuration Period_; + TDeadline::Duration Period_; }; } // namespace NYdb diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index ab12fb06e63..8c6cde808e8 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -148,7 +148,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MaxQueuedResponses_(params->GetMaxQueuedResponses()) , DrainOnDtors_(params->GetDrinOnDtors()) , BalancingSettings_(params->GetBalancingSettings()) - , GRpcKeepAliveTimeout_(params->GetGRpcKeepAliveTimeout()) + , GRpcKeepAliveTimeout_(TDeadline::SafeDurationCast(params->GetGRpcKeepAliveTimeout())) , GRpcKeepAlivePermitWithoutCalls_(params->GetGRpcKeepAlivePermitWithoutCalls()) , MemoryQuota_(params->GetMemoryQuota()) , MaxInboundMessageSize_(params->GetMaxInboundMessageSize()) @@ -156,9 +156,9 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , MaxMessageSize_(params->GetMaxMessageSize()) , QueuedRequests_(0) , TcpKeepAliveSettings_(params->GetTcpKeepAliveSettings()) - , SocketIdleTimeout_(params->GetSocketIdleTimeout()) + , SocketIdleTimeout_(TDeadline::SafeDurationCast(params->GetSocketIdleTimeout())) #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL - , ChannelPool_(TcpKeepAliveSettings_, SocketIdleTimeout_) + , ChannelPool_(TcpKeepAliveSettings_, params->GetSocketIdleTimeout()) #endif , NetworkThreadsNum_(params->GetNetworkThreadsNum()) , UsePerChannelTcpConnection_(params->GetUsePerChannelTcpConnection()) @@ -166,7 +166,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p , Log(params->GetLog()) { #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL - if (SocketIdleTimeout_ != TDuration::Max()) { + if (SocketIdleTimeout_ != TDeadline::Duration::max()) { auto channelPoolUpdateWrapper = [this] (NYdb::NIssue::TIssues&&, EStatus status) mutable { @@ -177,7 +177,7 @@ TGRpcConnectionsImpl::TGRpcConnectionsImpl(std::shared_ptr p ChannelPool_.DeleteExpiredStubsHolders(); return true; }; - AddPeriodicTask(channelPoolUpdateWrapper, SocketIdleTimeout_ * 0.1); + AddPeriodicTask(channelPoolUpdateWrapper, SocketIdleTimeout_ / 10); } #endif if (params->GetExecutor()) { @@ -204,7 +204,7 @@ TGRpcConnectionsImpl::~TGRpcConnectionsImpl() { ResponseQueue_->Stop(); } -void TGRpcConnectionsImpl::AddPeriodicTask(TPeriodicCb&& cb, TDuration period) { +void TGRpcConnectionsImpl::AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) { std::shared_ptr context; if (!TryCreateContext(context)) { NYdb::NIssue::TIssues issues; @@ -219,7 +219,7 @@ void TGRpcConnectionsImpl::AddPeriodicTask(TPeriodicCb&& cb, TDuration period) { } } -void TGRpcConnectionsImpl::ScheduleOneTimeTask(TSimpleCb&& fn, TDuration timeout) { +void TGRpcConnectionsImpl::ScheduleOneTimeTask(TSimpleCb&& fn, TDeadline::Duration timeout) { auto cbLow = [this, fn = std::move(fn)](NYdb::NIssue::TIssues&&, EStatus status) mutable { if (status != EStatus::SUCCESS) { return false; @@ -242,7 +242,7 @@ void TGRpcConnectionsImpl::ScheduleOneTimeTask(TSimpleCb&& fn, TDuration timeout return false; }; - if (timeout) { + if (timeout > TDeadline::Duration::zero()) { AddPeriodicTask(std::move(cbLow), timeout); } else { cbLow(NYdb::NIssue::TIssues(), EStatus::SUCCESS); @@ -315,8 +315,8 @@ void TGRpcConnectionsImpl::Stop(bool wait) { GRpcClientLow_.Stop(wait); } -void TGRpcConnectionsImpl::SetGrpcKeepAlive(NYdbGrpc::TGRpcClientConfig& config, const TDuration& timeout, bool permitWithoutCalls) { - ui64 timeoutMs = timeout.MilliSeconds(); +void TGRpcConnectionsImpl::SetGrpcKeepAlive(NYdbGrpc::TGRpcClientConfig& config, const TDeadline::Duration& timeout, bool permitWithoutCalls) { + std::uint64_t timeoutMs = std::chrono::duration_cast(timeout).count(); config.IntChannelParams[GRPC_ARG_KEEPALIVE_TIME_MS] = timeoutMs >> 3; config.IntChannelParams[GRPC_ARG_KEEPALIVE_TIMEOUT_MS] = timeoutMs; config.IntChannelParams[GRPC_ARG_HTTP2_MAX_PINGS_WITHOUT_DATA] = 0; @@ -340,7 +340,7 @@ TAsyncListEndpointsResult TGRpcConnectionsImpl::GetEndpoints(TDbDriverStatePtr d }; TRpcRequestSettings rpcSettings; - rpcSettings.ClientTimeout = GET_ENDPOINTS_TIMEOUT; + rpcSettings.Deadline = TDeadline::AfterDuration(GET_ENDPOINTS_TIMEOUT); RunDeferred( std::move(request), diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index e28c7af322e..e5b44a6f862 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -18,9 +18,9 @@ namespace NYdb::inline V3 { -constexpr TDuration GRPC_KEEP_ALIVE_TIMEOUT_FOR_DISCOVERY = TDuration::Seconds(10); -constexpr TDuration INITIAL_DEFERRED_CALL_DELAY = TDuration::MilliSeconds(10); // The delay before first deferred service call -constexpr TDuration GET_ENDPOINTS_TIMEOUT = TDuration::Seconds(10); // Time wait for ListEndpoints request, after this time we pass error to client +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 using NYdbGrpc::TCallMeta; using NYdbGrpc::IQueueClientContextPtr; @@ -45,8 +45,8 @@ class TGRpcConnectionsImpl TGRpcConnectionsImpl(std::shared_ptr params); ~TGRpcConnectionsImpl(); - void AddPeriodicTask(TPeriodicCb&& cb, TDuration period) override; - void ScheduleOneTimeTask(TSimpleCb&& fn, TDuration timeout); + void AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) override; + void ScheduleOneTimeTask(TSimpleCb&& fn, TDeadline::Duration timeout); NThreading::TFuture ScheduleFuture( TDuration timeout, IQueueClientContextPtr token = nullptr @@ -74,7 +74,7 @@ class TGRpcConnectionsImpl template using TServiceConnection = NYdbGrpc::TServiceConnection; - static void SetGrpcKeepAlive(NYdbGrpc::TGRpcClientConfig& config, const TDuration& timeout, bool permitWithoutCalls); + static void SetGrpcKeepAlive(NYdbGrpc::TGRpcClientConfig& config, const TDeadline::Duration& timeout, bool permitWithoutCalls); template std::pair>, TEndpointKey> GetServiceConnection( @@ -113,7 +113,7 @@ class TGRpcConnectionsImpl } clientConfig.Locator = endpoint.Endpoint; clientConfig.SslTargetNameOverride = endpoint.SslTargetNameOverride; - if (GRpcKeepAliveTimeout_) { + if (GRpcKeepAliveTimeout_ > TDeadline::Duration::zero()) { SetGrpcKeepAlive(clientConfig, GRpcKeepAliveTimeout_, GRpcKeepAlivePermitWithoutCalls_); } } @@ -233,7 +233,7 @@ class TGRpcConnectionsImpl } TCallMeta meta; - meta.Timeout = requestSettings.ClientTimeout; + meta.Timeout = requestSettings.Deadline; #ifndef YDB_GRPC_UNSECURE_AUTH meta.CallCredentials = dbState->CallCredentials; #else @@ -330,7 +330,7 @@ class TGRpcConnectionsImpl TDeferredOperationCb&& userResponseCb, TSimpleRpc rpc, TDbDriverStatePtr dbState, - TDuration deferredTimeout, + TDeadline::Duration delay, const TRpcRequestSettings& requestSettings, bool poll = false, std::shared_ptr context = nullptr) @@ -341,7 +341,7 @@ class TGRpcConnectionsImpl return; } - auto responseCb = [this, userResponseCb = std::move(userResponseCb), dbState, deferredTimeout, poll, context] + auto responseCb = [this, userResponseCb = std::move(userResponseCb), dbState, delay, deadline = requestSettings.Deadline, poll, context] (TResponse* response, TPlainStatus status) mutable { if (response) { @@ -352,7 +352,8 @@ class TGRpcConnectionsImpl std::move(userResponseCb), this, std::move(context), - deferredTimeout, + delay, + deadline, dbState, status.Endpoint); @@ -409,7 +410,7 @@ class TGRpcConnectionsImpl TDeferredResultCb&& userResponseCb, TSimpleRpc rpc, TDbDriverStatePtr dbState, - TDuration deferredTimeout, + TDeadline::Duration delay, const TRpcRequestSettings& requestSettings, std::shared_ptr context = nullptr) { @@ -427,7 +428,7 @@ class TGRpcConnectionsImpl operationCb, rpc, dbState, - deferredTimeout, + delay, requestSettings, true, // poll context); @@ -466,7 +467,7 @@ class TGRpcConnectionsImpl } TCallMeta meta; - meta.Timeout = requestSettings.ClientTimeout; + meta.Timeout = requestSettings.Deadline; #ifndef YDB_GRPC_UNSECURE_AUTH meta.CallCredentials = dbState->CallCredentials; #else @@ -763,7 +764,7 @@ class TGRpcConnectionsImpl const std::int64_t MaxQueuedResponses_; const bool DrainOnDtors_; const TBalancingPolicy::TImpl BalancingSettings_; - const TDuration GRpcKeepAliveTimeout_; + const TDeadline::Duration GRpcKeepAliveTimeout_; const bool GRpcKeepAlivePermitWithoutCalls_; const std::uint64_t MemoryQuota_; const std::uint64_t MaxInboundMessageSize_; @@ -772,7 +773,7 @@ class TGRpcConnectionsImpl std::atomic_int64_t QueuedRequests_; const NYdbGrpc::TTcpKeepAliveSettings TcpKeepAliveSettings_; - const TDuration SocketIdleTimeout_; + const TDeadline::Duration SocketIdleTimeout_; #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL NYdbGrpc::TChannelPool ChannelPool_; #endif diff --git a/src/client/impl/internal/internal_client/client.h b/src/client/impl/internal/internal_client/client.h index f53015355aa..3e52f984480 100644 --- a/src/client/impl/internal/internal_client/client.h +++ b/src/client/impl/internal/internal_client/client.h @@ -22,7 +22,7 @@ struct TListEndpointsResult; class IInternalClient { public: virtual NThreading::TFuture GetEndpoints(std::shared_ptr dbState) = 0; - virtual void AddPeriodicTask(TPeriodicCb&& cb, TDuration period) = 0; + virtual void AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) = 0; #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL virtual void DeleteChannels(const std::vector& endpoints) = 0; #endif diff --git a/src/client/impl/internal/make_request/make.h b/src/client/impl/internal/make_request/make.h index 01c8ff3666d..f10055d94a5 100644 --- a/src/client/impl/internal/make_request/make.h +++ b/src/client/impl/internal/make_request/make.h @@ -25,7 +25,7 @@ void FillOperationParams(const TRequestSettings& settings, TProtoRequest& params if (settings.OperationTimeout_) { SetDuration(settings.OperationTimeout_, *operationParams.mutable_operation_timeout()); - } else if (settings.ClientTimeout_ && settings.UseClientTimeoutForOperation_) { + } else if (settings.ClientTimeout_ && settings.ClientTimeout_ != TDuration::Max() && settings.UseClientTimeoutForOperation_) { SetDuration(settings.ClientTimeout_, *operationParams.mutable_operation_timeout()); } diff --git a/src/client/impl/internal/retry/retry.cpp b/src/client/impl/internal/retry/retry.cpp index dbeb516677d..73880d0e5c6 100644 --- a/src/client/impl/internal/retry/retry.cpp +++ b/src/client/impl/internal/retry/retry.cpp @@ -2,40 +2,40 @@ #include -#include -#include - #include -#include +#include + +#include namespace NYdb::inline V3::NRetry { -constexpr ui32 MAX_BACKOFF_DURATION_MS = TDuration::Hours(1).MilliSeconds(); +using TBackoffDuration = std::chrono::duration; -ui32 CalcBackoffTime(const TBackoffSettings& settings, ui32 retryNumber) { - ui32 backoffSlots = 1 << std::min(retryNumber, settings.Ceiling_); - TDuration maxDuration = settings.SlotDuration_ * backoffSlots; +constexpr TBackoffDuration MAX_BACKOFF_DURATION = std::chrono::hours(1); + +namespace { + +TBackoffDuration CalcBackoffTime(const TBackoffSettings& settings, std::uint32_t retryNumber) { + std::uint32_t backoffSlots = 1 << std::min(retryNumber, settings.Ceiling_); + TBackoffDuration maxDuration(settings.SlotDuration_.MicroSeconds() * backoffSlots); double uncertaintyRatio = std::max(std::min(settings.UncertainRatio_, 1.0), 0.0); double uncertaintyMultiplier = RandomNumber() * uncertaintyRatio - uncertaintyRatio + 1.0; - double durationMs = round(maxDuration.MilliSeconds() * uncertaintyMultiplier); + return std::max(std::min(maxDuration * uncertaintyMultiplier, MAX_BACKOFF_DURATION), TBackoffDuration::zero()); +} - return std::max(std::min(durationMs, (double)MAX_BACKOFF_DURATION_MS), 0.0); } -void Backoff(const NRetry::TBackoffSettings& settings, ui32 retryNumber) { - auto durationMs = CalcBackoffTime(settings, retryNumber); - Sleep(TDuration::MilliSeconds(durationMs)); +void Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber) { + std::this_thread::sleep_for(CalcBackoffTime(settings, retryNumber)); } void AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, - ui32 retryNumber, const std::function& fn) + std::uint32_t retryNumber, const std::function& fn) { - auto durationMs = CalcBackoffTime(settings, retryNumber); - client->ScheduleTask(fn, TDuration::MilliSeconds(durationMs)); + client->ScheduleTask(fn, std::chrono::duration_cast(CalcBackoffTime(settings, retryNumber))); } } - diff --git a/src/client/impl/internal/retry/retry.h b/src/client/impl/internal/retry/retry.h index d8f5f8bb792..1228796793c 100644 --- a/src/client/impl/internal/retry/retry.h +++ b/src/client/impl/internal/retry/retry.h @@ -19,10 +19,9 @@ class IClientImplCommon; namespace NYdb::inline V3::NRetry { -ui32 CalcBackoffTime(const TBackoffSettings& settings, ui32 retryNumber); -void Backoff(const NRetry::TBackoffSettings& settings, ui32 retryNumber); +void Backoff(const NRetry::TBackoffSettings& settings, std::uint32_t retryNumber); void AsyncBackoff(std::shared_ptr client, const TBackoffSettings& settings, - ui32 retryNumber, const std::function& fn); + std::uint32_t retryNumber, const std::function& fn); enum class NextStep { RetryImmediately, @@ -34,7 +33,7 @@ enum class NextStep { class TRetryContextBase : TNonCopyable { protected: TRetryOperationSettings Settings_; - ui32 RetryNumber_; + std::uint32_t RetryNumber_; TInstant RetryStartTime_; protected: diff --git a/src/client/impl/internal/rpc_request_settings/CMakeLists.txt b/src/client/impl/internal/rpc_request_settings/CMakeLists.txt new file mode 100644 index 00000000000..21ca5d556e9 --- /dev/null +++ b/src/client/impl/internal/rpc_request_settings/CMakeLists.txt @@ -0,0 +1,9 @@ +_ydb_sdk_add_library(impl-internal-rpc_request_settings INTERFACE) + +target_link_libraries(impl-internal-rpc_request_settings INTERFACE + yutil + ydb-time + client-impl-ydb_endpoints +) + +_ydb_sdk_install_targets(TARGETS impl-internal-rpc_request_settings) diff --git a/src/client/impl/internal/rpc_request_settings/settings.h b/src/client/impl/internal/rpc_request_settings/settings.h index 9006ccf3bfe..8fbf1f8c86d 100644 --- a/src/client/impl/internal/rpc_request_settings/settings.h +++ b/src/client/impl/internal/rpc_request_settings/settings.h @@ -3,6 +3,8 @@ #include #include +#include + namespace NYdb::inline V3 { struct TRpcRequestSettings { @@ -16,7 +18,7 @@ struct TRpcRequestSettings { UseDiscoveryEndpoint // Use single discovery endpoint } EndpointPolicy = TEndpointPolicy::UsePreferredEndpointOptionally; bool UseAuth = true; - TDuration ClientTimeout; + NYdb::TDeadline Deadline = NYdb::TDeadline::Max(); template static TRpcRequestSettings Make(const TRequestSettings& settings, const TEndpointKey& preferredEndpoint = {}, TEndpointPolicy endpointPolicy = TEndpointPolicy::UsePreferredEndpointOptionally) { @@ -28,11 +30,11 @@ struct TRpcRequestSettings { if (!settings.TraceParent_.empty()) { rpcSettings.Header.emplace_back("traceparent", settings.TraceParent_); } - + rpcSettings.PreferredEndpoint = preferredEndpoint; rpcSettings.EndpointPolicy = endpointPolicy; rpcSettings.UseAuth = true; - rpcSettings.ClientTimeout = settings.ClientTimeout_; + rpcSettings.Deadline = NYdb::TDeadline::AfterDuration(settings.ClientTimeout_); return rpcSettings; } }; diff --git a/src/client/impl/session/session_pool.h b/src/client/impl/session/session_pool.h index 9533460efd3..93f5fc409b3 100644 --- a/src/client/impl/session/session_pool.h +++ b/src/client/impl/session/session_pool.h @@ -20,10 +20,9 @@ class IGetSessionCtx : private TNonCopyable { }; //How often run session pool keep alive check -constexpr TDuration PERIODIC_ACTION_INTERVAL = TDuration::Seconds(5); -constexpr TDuration MAX_WAIT_SESSION_TIMEOUT = TDuration::Seconds(5); //Max time to wait session -constexpr ui64 PERIODIC_ACTION_BATCH_SIZE = 10; //Max number of tasks to perform during one interval -constexpr TDuration CREATE_SESSION_INTERNAL_TIMEOUT = TDuration::Seconds(2); //Timeout for createSession call inside session pool +constexpr TDeadline::Duration PERIODIC_ACTION_INTERVAL = std::chrono::seconds(5); +constexpr TDuration MAX_WAIT_SESSION_TIMEOUT = TDuration::Seconds(5); // Max time to wait session +constexpr std::uint64_t PERIODIC_ACTION_BATCH_SIZE = 10; // Max number of tasks to perform during one interval TStatus GetStatus(const TOperation& operation); TStatus GetStatus(const TStatus& status); diff --git a/src/client/monitoring/monitoring.cpp b/src/client/monitoring/monitoring.cpp index 16a4945ebd1..cdf1a5af725 100644 --- a/src/client/monitoring/monitoring.cpp +++ b/src/client/monitoring/monitoring.cpp @@ -22,11 +22,24 @@ class TSelfCheckResult::TImpl { Ydb::Monitoring::SelfCheckResult Result; }; +class TClusterStateResult::TImpl { +public: + TImpl(Ydb::Monitoring::ClusterStateResult&& result) + : Result(std::move(result)) + {} + Ydb::Monitoring::ClusterStateResult Result; +}; + TSelfCheckResult::TSelfCheckResult(TStatus&& status, Ydb::Monitoring::SelfCheckResult&& result) : TStatus(std::move(status)) , Impl_(std::make_shared(std::move(result))) {} +TClusterStateResult::TClusterStateResult(TStatus&& status, Ydb::Monitoring::ClusterStateResult&& result) + : TStatus(std::move(status)) + , Impl_(std::make_shared(std::move(result))) +{} + class TMonitoringClient::TImpl : public TClientImplCommon { public: TImpl(std::shared_ptr&& connections, const TCommonClientSettings& settings) @@ -87,6 +100,45 @@ class TMonitoringClient::TImpl : public TClientImplCommon(settings); + + if (settings.DurationSeconds_) { + request.set_duration_seconds(settings.DurationSeconds_.value()); + } + + if (settings.PeriodSeconds_) { + request.set_period_seconds(settings.PeriodSeconds_.value()); + } + auto promise = NThreading::NewPromise(); + + auto extractor = [promise] + (google::protobuf::Any* any, TPlainStatus status) mutable { + Ydb::Monitoring::ClusterStateResult result; + if (any) { + any->UnpackTo(&result); + } + TClusterStateResult val( + TStatus(std::move(status)), + std::move(result)); + + promise.SetValue(std::move(val)); + }; + + using Ydb::Monitoring::ClusterStateRequest; + using Ydb::Monitoring::ClusterStateResponse; + + Connections_->RunDeferred( + std::move(request), + extractor, + &Ydb::Monitoring::V1::MonitoringService::Stub::AsyncClusterState, + DbDriverState_, + INITIAL_DEFERRED_CALL_DELAY, + TRpcRequestSettings::Make(settings)); + + return promise.GetFuture(); + } }; TMonitoringClient::TMonitoringClient(const TDriver& driver, const TCommonClientSettings& settings) @@ -97,10 +149,17 @@ TAsyncSelfCheckResult TMonitoringClient::SelfCheck(const TSelfCheckSettings& set return Impl_->SelfCheck(settings); } +TAsyncClusterStateResult TMonitoringClient::ClusterState(const TClusterStateSettings& settings) { + return Impl_->ClusterState(settings); +} + } const Ydb::Monitoring::SelfCheckResult& TProtoAccessor::GetProto(const NYdb::NMonitoring::TSelfCheckResult& selfCheckResult) { return selfCheckResult.Impl_->Result; } +const Ydb::Monitoring::ClusterStateResult& TProtoAccessor::GetProto(const NYdb::NMonitoring::TClusterStateResult& clusterStateResult) { + return clusterStateResult.Impl_->Result; +} } diff --git a/src/client/operation/impl.h b/src/client/operation/impl.h index 22c42e2f353..6790d64363b 100644 --- a/src/client/operation/impl.h +++ b/src/client/operation/impl.h @@ -13,7 +13,7 @@ namespace NYdb::inline V3::NOperation { -constexpr TDuration OPERATION_CLIENT_TIMEOUT = TDuration::Seconds(5); +constexpr TDeadline::Duration OPERATION_CLIENT_TIMEOUT = std::chrono::seconds(5); class TOperationClient::TImpl : public TClientImplCommon { template @@ -37,7 +37,7 @@ class TOperationClient::TImpl : public TClientImplCommonRun( std::move(request), @@ -103,7 +103,7 @@ class TOperationClient::TImpl : public TClientImplCommonRunRunDeferred( diff --git a/src/client/persqueue_public/impl/write_session.cpp b/src/client/persqueue_public/impl/write_session.cpp index 07e90296885..3ad3306449b 100644 --- a/src/client/persqueue_public/impl/write_session.cpp +++ b/src/client/persqueue_public/impl/write_session.cpp @@ -86,7 +86,7 @@ TSimpleBlockingWriteSession::TSimpleBlockingWriteSession( subSettings.EventHandlers_.CommonHandler({}); } Writer = std::make_shared(subSettings, client, connections, dbDriverState); - Writer->Start(TDuration::Max()); + Writer->Start(TDuration::Zero()); } ui64 TSimpleBlockingWriteSession::GetInitSeqNo() { diff --git a/src/client/persqueue_public/impl/write_session_impl.cpp b/src/client/persqueue_public/impl/write_session_impl.cpp index 022bea0714e..9e0ff2c7c66 100644 --- a/src/client/persqueue_public/impl/write_session_impl.cpp +++ b/src/client/persqueue_public/impl/write_session_impl.cpp @@ -167,7 +167,7 @@ void TWriteSessionImpl::DoCdsRequest(TDuration delay) { INITIAL_DEFERRED_CALL_DELAY, TRpcRequestSettings::Make(settings)); // TODO: make client timeout setting }; - Connections->ScheduleOneTimeTask(std::move(cdsRequestCall), delay); + Connections->ScheduleOneTimeTask(std::move(cdsRequestCall), TDeadline::SafeDurationCast(delay)); return; } } diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 118c96f1624..66c4d998e7e 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -267,7 +267,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public } TAsyncFetchScriptResultsResult FetchScriptResultsImpl(Ydb::Query::FetchScriptResultsRequest&& request, const TFetchScriptResultsSettings& settings) { - using namespace Ydb::Query; + auto rpcSettings = TRpcRequestSettings::Make(settings); + if (!settings.FetchToken_.empty()) { request.set_fetch_token(TStringType{settings.FetchToken_}); } @@ -276,7 +277,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public auto promise = NThreading::NewPromise(); auto extractor = [promise] - (FetchScriptResultsResponse* response, TPlainStatus status) mutable { + (Ydb::Query::FetchScriptResultsResponse* response, TPlainStatus status) mutable { if (response) { NYdb::NIssue::TIssues opIssues; NYdb::NIssue::IssuesFromMessage(response->issues(), opIssues); @@ -300,13 +301,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public } }; - TRpcRequestSettings rpcSettings; - rpcSettings.ClientTimeout = TDuration::Seconds(60); - - Connections_->Run( + Connections_->Run( std::move(request), extractor, - &V1::QueryService::Stub::AsyncFetchScriptResults, + &Ydb::Query::V1::QueryService::Stub::AsyncFetchScriptResults, DbDriverState_, rpcSettings); @@ -383,7 +381,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public rpcSettings); } - TAsyncCreateSessionResult CreateAttachedSession(TDuration timeout) { + TAsyncCreateSessionResult CreateAttachedSession(const TRpcRequestSettings& rpcSettings) { using namespace Ydb::Query; Ydb::Query::CreateSessionRequest request; @@ -408,9 +406,6 @@ class TQueryClient::TImpl: public TClientImplCommon, public } }; - TRpcRequestSettings rpcSettings; - rpcSettings.ClientTimeout = timeout; - Connections_->Run( std::move(request), extractor, @@ -422,14 +417,14 @@ class TQueryClient::TImpl: public TClientImplCommon, public } TAsyncCreateSessionResult GetSession(const TCreateSessionSettings& settings) { - using namespace NSessionPool; + auto rpcSettings = TRpcRequestSettings::Make(settings); - class TQueryClientGetSessionCtx : public IGetSessionCtx { + class TQueryClientGetSessionCtx : public NSessionPool::IGetSessionCtx { public: - TQueryClientGetSessionCtx(std::shared_ptr client, TDuration timeout) + TQueryClientGetSessionCtx(std::shared_ptr client, const TRpcRequestSettings& settings) : Promise(NThreading::NewPromise()) , Client(client) - , ClientTimeout(timeout) + , RpcSettings(settings) {} TAsyncCreateSessionResult GetFuture() { @@ -454,7 +449,7 @@ class TQueryClient::TImpl: public TClientImplCommon, public } void ReplyNewSession() override { - Client->CreateAttachedSession(ClientTimeout).Subscribe( + Client->CreateAttachedSession(RpcSettings).Subscribe( [promise{std::move(Promise)}](TAsyncCreateSessionResult future) mutable { promise.SetValue(future.ExtractValue()); @@ -467,10 +462,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public } NThreading::TPromise Promise; std::shared_ptr Client; - TDuration ClientTimeout; + TRpcRequestSettings RpcSettings; }; - auto ctx = std::make_unique(shared_from_this(), settings.ClientTimeout_); + auto ctx = std::make_unique(shared_from_this(), rpcSettings); auto future = ctx->GetFuture(); SessionPool_.GetSession(std::move(ctx)); return future; diff --git a/src/client/table/impl/table_client.cpp b/src/client/table/impl/table_client.cpp index 26594dfb0b6..771db4bd070 100644 --- a/src/client/table/impl/table_client.cpp +++ b/src/client/table/impl/table_client.cpp @@ -78,7 +78,7 @@ NThreading::TFuture TTableClient::TImpl::Stop() { return Drain(); } -void TTableClient::TImpl::ScheduleTaskUnsafe(std::function&& fn, TDuration timeout) { +void TTableClient::TImpl::ScheduleTaskUnsafe(std::function&& fn, TDeadline::Duration timeout) { Connections_->ScheduleOneTimeTask(std::move(fn), timeout); } @@ -332,7 +332,7 @@ TAsyncCreateSessionResult TTableClient::TImpl::GetSession(const TCreateSessionSe //TODO: Do we realy need it? Client->ScheduleTaskUnsafe([promise{std::move(Promise)}, val{std::move(val)}]() mutable { promise.SetValue(std::move(val)); - }, TDuration()); + }, TDeadline::Duration::zero()); } NThreading::TPromise Promise; std::shared_ptr Client; diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 18dd8720e97..e014fabc0a8 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -25,7 +25,7 @@ namespace NYdb::inline V3 { namespace NTable { //How ofter run host scan to perform session balancing -constexpr TDuration HOSTSCAN_PERIODIC_ACTION_INTERVAL = TDuration::Seconds(2); +constexpr TDeadline::Duration HOSTSCAN_PERIODIC_ACTION_INTERVAL = std::chrono::seconds(2); constexpr TDuration KEEP_ALIVE_CLIENT_TIMEOUT = TDuration::Seconds(5); TDuration GetMinTimeToTouch(const TSessionPoolSettings& settings); @@ -43,7 +43,7 @@ class TTableClient::TImpl: public TClientImplCommon, public void InitStopper(); NThreading::TFuture Drain(); NThreading::TFuture Stop(); - void ScheduleTaskUnsafe(std::function&& fn, TDuration timeout); + void ScheduleTaskUnsafe(std::function&& fn, TDeadline::Duration timeout); void StartPeriodicSessionPoolTask(); static ui64 ScanForeignLocations(std::shared_ptr client); static std::pair ScanLocation(std::shared_ptr client, diff --git a/src/client/table/out.cpp b/src/client/table/out.cpp index 9d7a957219e..b827e8b2511 100644 --- a/src/client/table/out.cpp +++ b/src/client/table/out.cpp @@ -83,3 +83,90 @@ Y_DECLARE_OUT_SPEC(, NYdb::NTable::TKMeansTreeSettings, stream, value) { ", levels: " << value.Levels << " }"; } + +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::Unspecified: + stream << "unspecified"; + break; + } +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::ETokenizer, stream, value) { + switch (value) { + case NYdb::NTable::TFulltextIndexSettings::ETokenizer::Whitespace: + stream << "whitespace"; + break; + case NYdb::NTable::TFulltextIndexSettings::ETokenizer::Standard: + stream << "standard"; + break; + case NYdb::NTable::TFulltextIndexSettings::ETokenizer::Keyword: + stream << "keyword"; + break; + case NYdb::NTable::TFulltextIndexSettings::ETokenizer::Unspecified: + stream << "unspecified"; + break; + } +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::TAnalyzers, stream, value) { + stream << "{ tokenizer: " << value.Tokenizer.value_or(NYdb::NTable::TFulltextIndexSettings::ETokenizer::Unspecified); + if (value.Language.has_value()) { + stream << ", language: " << *value.Language; + } + if (value.UseFilterLowercase.has_value()) { + stream << ", use_filter_lowercase: " << (*value.UseFilterLowercase ? "true" : "false"); + } + if (value.UseFilterStopwords.has_value()) { + stream << ", use_filter_stopwords: " << (*value.UseFilterStopwords ? "true" : "false"); + } + if (value.UseFilterNgram.has_value()) { + stream << ", use_filter_ngram: " << (*value.UseFilterNgram ? "true" : "false"); + } + if (value.UseFilterEdgeNgram.has_value()) { + stream << ", use_filter_edge_ngram: " << (*value.UseFilterEdgeNgram ? "true" : "false"); + } + if (value.FilterNgramMinLength.has_value()) { + stream << ", filter_ngram_min_length: " << *value.FilterNgramMinLength; + } + if (value.FilterNgramMaxLength.has_value()) { + stream << ", filter_ngram_max_length: " << *value.FilterNgramMaxLength; + } + if (value.UseFilterLength.has_value()) { + stream << ", use_filter_length: " << (*value.UseFilterLength ? "true" : "false"); + } + if (value.FilterLengthMin.has_value()) { + stream << ", filter_length_min: " << *value.FilterLengthMin; + } + if (value.FilterLengthMax.has_value()) { + stream << ", filter_length_max: " << *value.FilterLengthMax; + } + stream << " }"; +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings::TColumnAnalyzers, stream, value) { + stream << "{ "; + if (value.Column.has_value()) { + stream << "column: " << *value.Column << ", "; + } + if (value.Analyzers.has_value()) { + stream << "analyzers: " << *value.Analyzers; + } + stream << " }"; +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TFulltextIndexSettings, stream, value) { + stream << "{ layout: " << value.Layout.value_or(NYdb::NTable::TFulltextIndexSettings::ELayout::Unspecified); + if (!value.Columns.empty()) { + stream << ", columns: ["; + for (size_t i = 0; i < value.Columns.size(); ++i) { + if (i > 0) stream << ", "; + stream << value.Columns[i]; + } + stream << "]"; + } + stream << " }"; +} diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 6e3332b2d28..3b00f587551 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -480,6 +480,14 @@ class TTableDescription::TImpl { Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, dataColumns, {}, indexSettings)); } + void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { + Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, {}, {}, indexSettings)); + } + + void AddFulltextIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { + Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, dataColumns, {}, indexSettings)); + } + void AddChangefeed(const std::string& name, EChangefeedMode mode, EChangefeedFormat format) { Changefeeds_.emplace_back(name, mode, format); } @@ -784,6 +792,14 @@ 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) { + Impl_->AddFulltextIndex(indexName, EIndexType::GlobalFulltext, indexColumns, indexSettings); +} + +void TTableDescription::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { + Impl_->AddFulltextIndex(indexName, EIndexType::GlobalFulltext, indexColumns, dataColumns, indexSettings); +} + void TTableDescription::AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns) { AddSyncSecondaryIndex(indexName, indexColumns); } @@ -1273,6 +1289,16 @@ TTableBuilder& TTableBuilder::AddVectorKMeansTreeIndex(const std::string& indexN return *this; } +TTableBuilder& TTableBuilder::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TFulltextIndexSettings& indexSettings) { + TableDescription_.AddFulltextIndex(indexName, indexColumns, dataColumns, indexSettings); + return *this; +} + +TTableBuilder& TTableBuilder::AddFulltextIndex(const std::string& indexName, const std::vector& indexColumns, const TFulltextIndexSettings& indexSettings) { + TableDescription_.AddFulltextIndex(indexName, indexColumns, indexSettings); + return *this; +} + TTableBuilder& TTableBuilder::AddSecondaryIndex(const std::string& indexName, const std::string& indexColumn) { return AddSyncSecondaryIndex(indexName, indexColumn); } @@ -2330,7 +2356,7 @@ TIndexDescription::TIndexDescription( const std::vector& indexColumns, const std::vector& dataColumns, const std::vector& globalIndexSettings, - const std::variant& specializedIndexSettings + const std::variant& specializedIndexSettings ) : IndexName_(name) , IndexType_(type) , IndexColumns_(indexColumns) @@ -2371,7 +2397,7 @@ const std::vector& TIndexDescription::GetDataColumns() const { return DataColumns_; } -const std::variant& TIndexDescription::GetIndexSettings() const { +const std::variant& TIndexDescription::GetIndexSettings() const { return SpecializedIndexSettings_; } @@ -2544,13 +2570,187 @@ void TKMeansTreeSettings::Out(IOutputStream& o) const { o << *this; } +TFulltextIndexSettings::TAnalyzers FromProto(const Ydb::Table::FulltextIndexSettings::Analyzers& proto) { + using ETokenizer = TFulltextIndexSettings::ETokenizer; + using TAnalyzers = TFulltextIndexSettings::TAnalyzers; + + auto convertTokenizer = [&] { + switch (proto.tokenizer()) { + case Ydb::Table::FulltextIndexSettings::WHITESPACE: + return ETokenizer::Whitespace; + case Ydb::Table::FulltextIndexSettings::STANDARD: + return ETokenizer::Standard; + case Ydb::Table::FulltextIndexSettings::KEYWORD: + return ETokenizer::Keyword; + default: + return ETokenizer::Unspecified; + } + }; + + TAnalyzers result; + result.Tokenizer = convertTokenizer(); + + if (proto.has_language()) { + result.Language = proto.language(); + } + if (proto.has_use_filter_lowercase()) { + result.UseFilterLowercase = proto.use_filter_lowercase(); + } + if (proto.has_use_filter_stopwords()) { + result.UseFilterStopwords = proto.use_filter_stopwords(); + } + if (proto.has_use_filter_ngram()) { + result.UseFilterNgram = proto.use_filter_ngram(); + } + if (proto.has_use_filter_edge_ngram()) { + result.UseFilterEdgeNgram = proto.use_filter_edge_ngram(); + } + if (proto.has_filter_ngram_min_length()) { + result.FilterNgramMinLength = proto.filter_ngram_min_length(); + } + if (proto.has_filter_ngram_max_length()) { + result.FilterNgramMaxLength = proto.filter_ngram_max_length(); + } + if (proto.has_use_filter_length()) { + result.UseFilterLength = proto.use_filter_length(); + } + if (proto.has_filter_length_min()) { + result.FilterLengthMin = proto.filter_length_min(); + } + if (proto.has_filter_length_max()) { + result.FilterLengthMax = proto.filter_length_max(); + } + + return result; +} + +Ydb::Table::FulltextIndexSettings::Analyzers ToProto(const TFulltextIndexSettings::TAnalyzers& analyzers) { + using ETokenizer = TFulltextIndexSettings::ETokenizer; + + auto convertTokenizer = [&] { + switch (*analyzers.Tokenizer) { + case ETokenizer::Whitespace: + return Ydb::Table::FulltextIndexSettings::WHITESPACE; + case ETokenizer::Standard: + return Ydb::Table::FulltextIndexSettings::STANDARD; + case ETokenizer::Keyword: + return Ydb::Table::FulltextIndexSettings::KEYWORD; + case ETokenizer::Unspecified: + return Ydb::Table::FulltextIndexSettings::TOKENIZER_UNSPECIFIED; + } + return Ydb::Table::FulltextIndexSettings::TOKENIZER_UNSPECIFIED; + }; + + Ydb::Table::FulltextIndexSettings::Analyzers proto; + if (analyzers.Tokenizer) { + proto.set_tokenizer(convertTokenizer()); + } + if (analyzers.Language.has_value()) { + proto.set_language(*analyzers.Language); + } + if (analyzers.UseFilterLowercase.has_value()) { + proto.set_use_filter_lowercase(*analyzers.UseFilterLowercase); + } + if (analyzers.UseFilterStopwords.has_value()) { + proto.set_use_filter_stopwords(*analyzers.UseFilterStopwords); + } + if (analyzers.UseFilterNgram.has_value()) { + proto.set_use_filter_ngram(*analyzers.UseFilterNgram); + } + if (analyzers.UseFilterEdgeNgram.has_value()) { + proto.set_use_filter_edge_ngram(*analyzers.UseFilterEdgeNgram); + } + if (analyzers.FilterNgramMinLength.has_value()) { + proto.set_filter_ngram_min_length(*analyzers.FilterNgramMinLength); + } + if (analyzers.FilterNgramMaxLength.has_value()) { + proto.set_filter_ngram_max_length(*analyzers.FilterNgramMaxLength); + } + if (analyzers.UseFilterLength.has_value()) { + proto.set_use_filter_length(*analyzers.UseFilterLength); + } + if (analyzers.FilterLengthMin.has_value()) { + proto.set_filter_length_min(*analyzers.FilterLengthMin); + } + if (analyzers.FilterLengthMax.has_value()) { + proto.set_filter_length_max(*analyzers.FilterLengthMax); + } + + return proto; +} + +TFulltextIndexSettings::TColumnAnalyzers FromProto(const Ydb::Table::FulltextIndexSettings::ColumnAnalyzers& proto) { + TFulltextIndexSettings::TColumnAnalyzers result; + if (proto.has_column()) { + result.Column = proto.column(); + } + if (proto.has_analyzers()) { + result.Analyzers = FromProto(proto.analyzers()); + } + return result; +} + +Ydb::Table::FulltextIndexSettings::ColumnAnalyzers ToProto(const TFulltextIndexSettings::TColumnAnalyzers& columnAnalyzers) { + Ydb::Table::FulltextIndexSettings::ColumnAnalyzers proto; + if (columnAnalyzers.Column.has_value()) { + proto.set_column(*columnAnalyzers.Column); + } + if (columnAnalyzers.Analyzers.has_value()) { + *proto.mutable_analyzers() = ToProto(*columnAnalyzers.Analyzers); + } + return proto; +} + +TFulltextIndexSettings TFulltextIndexSettings::FromProto(const Ydb::Table::FulltextIndexSettings& proto) { + auto convertLayout = [&] { + switch (proto.layout()) { + case Ydb::Table::FulltextIndexSettings::FLAT: + return ELayout::Flat; + 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::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); + } +} + +void TFulltextIndexSettings::Out(IOutputStream& o) const { + o << *this; +} + template TIndexDescription TIndexDescription::FromProto(const TProto& proto) { EIndexType type; std::vector indexColumns; std::vector dataColumns; std::vector globalIndexSettings; - std::variant specializedIndexSettings = std::monostate{}; + std::variant specializedIndexSettings = std::monostate{}; indexColumns.assign(proto.index_columns().begin(), proto.index_columns().end()); dataColumns.assign(proto.data_columns().begin(), proto.data_columns().end()); @@ -2580,6 +2780,13 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { specializedIndexSettings = TKMeansTreeSettings::FromProto(vectorProto.vector_settings()); break; } + case TProto::kGlobalFulltextIndex: { + type = EIndexType::GlobalFulltext; + const auto& fulltextProto = proto.global_fulltext_index(); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(fulltextProto.settings())); + specializedIndexSettings = TFulltextIndexSettings::FromProto(fulltextProto.fulltext_settings()); + break; + } default: // fallback to global sync type = EIndexType::GlobalSync; globalIndexSettings.resize(1); @@ -2635,6 +2842,18 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { } break; } + case EIndexType::GlobalFulltext: { + auto* global_fulltext_index = proto.mutable_global_fulltext_index(); + auto& settings = *global_fulltext_index->mutable_settings(); + auto& fulltext_settings = *global_fulltext_index->mutable_fulltext_settings(); + if (GlobalIndexSettings_.size() == 1) { + GlobalIndexSettings_[0].SerializeTo(settings); + } + if (const auto* ftSettings = std::get_if(&SpecializedIndexSettings_)) { + ftSettings->SerializeTo(fulltext_settings); + } + break; + } case EIndexType::Unknown: break; } @@ -2656,11 +2875,23 @@ void TIndexDescription::Out(IOutputStream& o) const { o << ", data_columns: [" << JoinSeq(", ", DataColumns_) << "]"; } - std::visit([&](const T& settings) { - if constexpr (!std::is_same_v) { - o << ", vector_settings: " << settings; + switch (IndexType_) { + case EIndexType::GlobalSync: + case EIndexType::GlobalAsync: + case EIndexType::GlobalUnique: + case EIndexType::Unknown: + break; + case EIndexType::GlobalVectorKMeansTree: + if (auto settings = std::get_if(&SpecializedIndexSettings_)) { + o << ", vector_settings: " << *settings; + } + break; + case EIndexType::GlobalFulltext: + if (auto settings = std::get_if(&SpecializedIndexSettings_)) { + o << ", fulltext_settings: " << *settings; } - }, SpecializedIndexSettings_); + break; + } o << " }"; } diff --git a/src/client/topic/impl/write_session.cpp b/src/client/topic/impl/write_session.cpp index 968ac4a6a0b..3b51c48e723 100644 --- a/src/client/topic/impl/write_session.cpp +++ b/src/client/topic/impl/write_session.cpp @@ -107,7 +107,7 @@ TSimpleBlockingWriteSession::TSimpleBlockingWriteSession( subSettings.EventHandlers_.CommonHandler({}); } Writer = std::make_shared(subSettings, client, connections, dbDriverState); - Writer->Start(TDuration::Max()); + Writer->Start(TDuration::Zero()); } uint64_t TSimpleBlockingWriteSession::GetInitSeqNo() { diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 774bf0c1302..0f31c09fb5d 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -296,7 +296,7 @@ void TWriteSessionImpl::ConnectToPreferredPartitionLocation(const TDuration& del context); }; - Connections->ScheduleOneTimeTask(std::move(callback), delay); + Connections->ScheduleOneTimeTask(std::move(callback), TDeadline::SafeDurationCast(delay)); } void TWriteSessionImpl::OnDescribePartition(const TStatus& status, const Ydb::Topic::DescribePartitionResult& proto, const NYdbGrpc::IQueueClientContextPtr& describePartitionContext) diff --git a/src/client/topic/ut/slow/txusage_slow_ut.cpp b/src/client/topic/ut/slow/txusage_slow_ut.cpp new file mode 100644 index 00000000000..c4fae9af75d --- /dev/null +++ b/src/client/topic/ut/slow/txusage_slow_ut.cpp @@ -0,0 +1,138 @@ +#include +#include + +namespace NYdb::inline V3::NTopic::NTests::NTxUsage { + +Y_UNIT_TEST_SUITE(TxUsage) { + +Y_UNIT_TEST_F(Transactions_Conflict_On_SeqNo_Table, TFixtureTable) +{ + TestTransactionsConflictOnSeqNo(); +} + +Y_UNIT_TEST_F(Transactions_Conflict_On_SeqNo_Query, TFixtureQuery) +{ + TestTransactionsConflictOnSeqNo(); +} + +Y_UNIT_TEST_F(WriteToTopic_Demo_44_Table, TFixtureTable) +{ + TestWriteToTopic44(); +} + +Y_UNIT_TEST_F(WriteToTopic_Demo_44_Query, TFixtureQuery) +{ + TestWriteToTopic44(); +} + +Y_UNIT_TEST_F(Write_Random_Sized_Messages_In_Wide_Transactions_Table, TFixtureTable) +{ + TestWriteRandomSizedMessagesInWideTransactions(); +} + +Y_UNIT_TEST_F(Write_Random_Sized_Messages_In_Wide_Transactions_Query, TFixtureQuery) +{ + TestWriteRandomSizedMessagesInWideTransactions(); +} + +Y_UNIT_TEST_F(Write_Only_Big_Messages_In_Wide_Transactions_Table, TFixtureTable) +{ + TestWriteOnlyBigMessagesInWideTransactions(); +} + +Y_UNIT_TEST_F(Write_Only_Big_Messages_In_Wide_Transactions_Query, TFixtureQuery) +{ + TestWriteOnlyBigMessagesInWideTransactions(); +} + +Y_UNIT_TEST_F(Write_And_Read_Big_Messages_1, TFixtureNoClient) +{ + TestWriteAndReadMessages(27, 64'000 * 12, false); +} + +Y_UNIT_TEST_F(Write_And_Read_Big_Messages_2, TFixtureNoClient) +{ + TestWriteAndReadMessages(27, 64'000 * 12, true); +} + +Y_UNIT_TEST_F(Write_And_Read_Huge_Messages_1, TFixtureNoClient) +{ + TestWriteAndReadMessages(4, 9'000'000, false); +} + +Y_UNIT_TEST_F(Write_And_Read_Huge_Messages_2, TFixtureNoClient) +{ + TestWriteAndReadMessages(4, 9'000'000, true); +} + +Y_UNIT_TEST_F(Write_And_Read_Gigant_Messages_1, TFixtureNoClient) +{ + TestWriteAndReadMessages(4, 61'000'000, false); +} + +Y_UNIT_TEST_F(Write_And_Read_Gigant_Messages_2, TFixtureNoClient) +{ + TestWriteAndReadMessages(4, 61'000'000, true); +} + +Y_UNIT_TEST_F(Write_50k_100times_50tx, TFixtureTable) +{ + // 100 transactions. Write 100 50KB messages in each folder. Call the commit at the same time. + // As a result, there will be a lot of small blobs in the FastWrite zone of the main batch, + // which will be picked up by a compact. The scenario is similar to the work of Ya.Metrika. + + const std::size_t PARTITIONS_COUNT = 2; + const std::size_t TXS_COUNT = 50; + + auto makeSourceId = [](unsigned txId, unsigned partitionId) { + std::string sourceId = TEST_MESSAGE_GROUP_ID; + sourceId += "_"; + sourceId += ToString(txId); + sourceId += "_"; + sourceId += ToString(partitionId); + return sourceId; + }; + + CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); + + SetPartitionWriteSpeed("topic_A", 50'000'000); + + std::vector> sessions; + std::vector> transactions; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + sessions.push_back(CreateSession()); + auto& session = sessions.back(); + + transactions.push_back(session->BeginTx()); + auto& tx = transactions.back(); + + auto sourceId = makeSourceId(i, 0); + for (size_t j = 0; j < 100; ++j) { + WriteToTopic("topic_A", sourceId, std::string(50'000, 'x'), tx.get(), 0); + } + WaitForAcks("topic_A", sourceId); + + sourceId = makeSourceId(i, 1); + WriteToTopic("topic_A", sourceId, std::string(50'000, 'x'), tx.get(), 1); + WaitForAcks("topic_A", sourceId); + } + + // We are doing an asynchronous commit of transactions. They will be executed simultaneously. + std::vector futures; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); + } + + // All transactions must be completed successfully. + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures[i].Wait(); + const auto& result = futures[i].GetValueSync(); + UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); + } +} + +} + +} diff --git a/src/client/topic/ut/topic_to_table_ut.cpp b/src/client/topic/ut/topic_to_table_ut.cpp index 4a9ccb7b9de..38a5c89ba38 100644 --- a/src/client/topic/ut/topic_to_table_ut.cpp +++ b/src/client/topic/ut/topic_to_table_ut.cpp @@ -1,2774 +1,219 @@ -#include "ut_utils/topic_sdk_test_setup.h" - -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include - -#include +#include #include -#include -#include -#include -#include - -#include - - -using namespace std::chrono_literals; -namespace NYdb::inline V3::NTopic::NTests { - -const auto TEST_MESSAGE_GROUP_ID_1 = TEST_MESSAGE_GROUP_ID + "_1"; -const auto TEST_MESSAGE_GROUP_ID_2 = TEST_MESSAGE_GROUP_ID + "_2"; -const auto TEST_MESSAGE_GROUP_ID_3 = TEST_MESSAGE_GROUP_ID + "_3"; -const auto TEST_MESSAGE_GROUP_ID_4 = TEST_MESSAGE_GROUP_ID + "_4"; +namespace NYdb::inline V3::NTopic::NTests::NTxUsage { Y_UNIT_TEST_SUITE(TxUsage) { -class TFixture : public NUnitTest::TBaseFixture { -protected: - using TTopicReadSession = NTopic::IReadSession; - using TTopicReadSessionPtr = std::shared_ptr; - using TTopicWriteSession = NTopic::IWriteSession; - using TTopicWriteSessionPtr = std::shared_ptr; - - struct TTopicWriteSessionContext { - TTopicWriteSessionPtr Session; - std::optional ContinuationToken; - size_t WriteCount = 0; - size_t WrittenAckCount = 0; - size_t WrittenInTxAckCount = 0; - - void WaitForContinuationToken(); - void Write(const std::string& message, TTransactionBase* tx = nullptr); - - size_t AckCount() const { return WrittenAckCount + WrittenInTxAckCount; } - - void WaitForEvent(); - }; - - struct TFeatureFlags { - bool EnablePQConfigTransactionsAtSchemeShard = true; - }; - - class ISession { - public: - using TExecuteInTxResult = std::pair, std::unique_ptr>; - - virtual std::vector Execute(const std::string& query, - TTransactionBase* tx, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) = 0; - - virtual TExecuteInTxResult ExecuteInTx(const std::string& query, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) = 0; - - virtual std::unique_ptr BeginTx() = 0; - virtual void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) = 0; - virtual void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) = 0; - - virtual void Close() = 0; - - virtual TAsyncStatus AsyncCommitTx(TTransactionBase& tx) = 0; - - virtual ~ISession() = default; - }; - - void SetUp(NUnitTest::TTestContext&) override; - - void NotifySchemeShard(const TFeatureFlags& flags); - - std::unique_ptr CreateSession(); - - struct TReadMessageSettings { - TTransactionBase& Tx; - bool CommitOffsets = false; - std::optional Offset; - }; - - void CreateTopic(const std::string& path = TEST_TOPIC, - const std::string& consumer = TEST_CONSUMER, - std::size_t partitionCount = 1, - std::optional maxPartitionCount = std::nullopt, - const TDuration retention = TDuration::Hours(1), - bool important = false); - - void AddConsumer(const std::string& topicPath, const std::vector& consumers); - - void SetPartitionWriteSpeed(const std::string& topicName, std::size_t bytesPerSeconds); - - TTopicWriteSessionPtr CreateTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - std::optional partitionId); - TTopicWriteSessionContext& GetTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - std::optional partitionId); - - TTopicReadSessionPtr CreateTopicReadSession(const std::string& topicPath, - const std::string& consumerName, - std::optional partitionId); - TTopicReadSessionPtr GetTopicReadSession(const std::string& topicPath, - const std::string& consumerName, - std::optional partitionId); - - void WriteToTopic(const std::string& topicPath, - const std::string& messageGroupId, - const std::string& message, - TTransactionBase* tx = nullptr, - std::optional partitionId = std::nullopt); - std::vector ReadFromTopic(const std::string& topicPath, - const std::string& consumerName, - const TDuration& duration, - TTransactionBase* tx = nullptr, - std::optional partitionId = std::nullopt); - void WaitForAcks(const std::string& topicPath, - const std::string& messageGroupId, - std::size_t writtenInTxCount = std::numeric_limits::max()); - void WaitForSessionClose(const std::string& topicPath, - const std::string& messageGroupId, - NYdb::EStatus status); - void CloseTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - bool force = false); - void CloseTopicReadSession(const std::string& topicPath, - const std::string& consumerName); - - enum EEndOfTransaction { - Commit, - Rollback, - CloseTableSession - }; - - struct TTransactionCompletionTestDescription { - std::vector Topics; - EEndOfTransaction EndOfTransaction = Commit; - }; - - void TestTheCompletionOfATransaction(const TTransactionCompletionTestDescription& d); - void RestartPQTablet(const std::string& topicPath, std::uint32_t partition); - void DumpPQTabletKeys(const std::string& topicName, std::uint32_t partition); - void PQTabletPrepareFromResource(const std::string& topicPath, - std::uint32_t partitionId, - const std::string& resourceName); - - void DeleteSupportivePartition(const std::string& topicName, - std::uint32_t partition); - - struct TTableRecord { - TTableRecord() = default; - TTableRecord(const std::string& key, const std::string& value); - - std::string Key; - std::string Value; - }; - - std::vector MakeTableRecords(); - std::string MakeJsonDoc(const std::vector& records); - - void CreateTable(const std::string& path); - void UpsertToTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx); - void InsertToTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx); - void DeleteFromTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx); - size_t GetTableRecordsCount(const std::string& tablePath); - - enum ERestartPQTabletMode { - ERestartNo, - ERestartBeforeCommit, - ERestartAfterCommit, - }; - - struct TTestTxWithBigBlobsParams { - size_t OldHeadCount = 0; - size_t BigBlobsCount = 2; - size_t NewHeadCount = 0; - ERestartPQTabletMode RestartMode = ERestartNo; - }; - - void TestTxWithBigBlobs(const TTestTxWithBigBlobsParams& params); - - void WriteMessagesInTx(std::size_t big, size_t small); - - const TDriver& GetDriver() const; - NTable::TTableClient& GetTableClient(); - - void CheckTabletKeys(const std::string& topicName); - void DumpPQTabletKeys(const std::string& topicName); - - std::vector Read_Exactly_N_Messages_From_Topic(const std::string& topicPath, - const std::string& consumerName, - size_t count); - - void TestWriteRandomSizedMessagesInWideTransactions(); - - void TestWriteOnlyBigMessagesInWideTransactions(); - - void TestTransactionsConflictOnSeqNo(); - - void TestWriteToTopic1(); - - void TestWriteToTopic4(); - - void TestWriteToTopic7(); - - void TestWriteToTopic9(); - - void TestWriteToTopic10(); - - void TestWriteToTopic11(); - - void TestWriteToTopic12(); - - void TestWriteToTopic13(); - - void TestWriteToTopic14(); - - void TestWriteToTopic16(); - - void TestWriteToTopic24(); - - void TestWriteToTopic26(); - - void TestWriteToTopic27(); - - void TestWriteToTopic38(); - - void TestWriteToTopic40(); - - void TestWriteToTopic41(); - - void TestWriteToTopic42(); - - void TestWriteToTopic43(); - - void TestWriteToTopic44(); - - void TestWriteToTopic45(); - - void TestWriteToTopic46(); - - void TestWriteToTopic47(); - - void TestWriteToTopic50(); - - struct TAvgWriteBytes { - std::uint64_t PerSec = 0; - std::uint64_t PerMin = 0; - std::uint64_t PerHour = 0; - std::uint64_t PerDay = 0; - }; - - TAvgWriteBytes GetAvgWriteBytes(const std::string& topicPath, - std::uint32_t partitionId); - - void CheckAvgWriteBytes(const std::string& topicPath, - std::uint32_t partitionId, - std::size_t minSize, std::size_t maxSize); - - void SplitPartition(const std::string& topicPath, - std::uint32_t partitionId, - const std::string& boundary); - - virtual bool GetEnableOltpSink() const; - virtual bool GetEnableOlapSink() const; - virtual bool GetEnableHtapTx() const; - virtual bool GetAllowOlapDataQuery() const; - - size_t GetPQCacheRenameKeysCount(); - - enum class EClientType { - Table, - Query, - None - }; - - virtual EClientType GetClientType() const = 0; - virtual ~TFixture() = default; - - void TestWriteAndReadMessages(size_t count, size_t size, bool restart); - -private: - class TTableSession : public ISession { - public: - TTableSession(NTable::TTableClient& client); - - std::vector Execute(const std::string& query, - TTransactionBase* tx, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) override; - - TExecuteInTxResult ExecuteInTx(const std::string& query, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) override; - - std::unique_ptr BeginTx() override; - void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; - void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; - - TAsyncStatus AsyncCommitTx(TTransactionBase& tx) override; - - void Close() override; - - private: - NTable::TSession Init(NTable::TTableClient& client); - - NTable::TSession Session_; - }; - - class TQuerySession : public ISession { - public: - TQuerySession(NQuery::TQueryClient& client, - const std::string& endpoint, - const std::string& database); - - std::vector Execute(const std::string& query, - TTransactionBase* tx, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) override; - - TExecuteInTxResult ExecuteInTx(const std::string& query, - bool commit = true, - const TParams& params = TParamsBuilder().Build()) override; - - std::unique_ptr BeginTx() override; - void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; - void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; - - TAsyncStatus AsyncCommitTx(TTransactionBase& tx) override; - - void Close() override; - - private: - NQuery::TSession Init(NQuery::TQueryClient& client); - - NQuery::TSession Session_; - std::string Endpoint_; - std::string Database_; - }; - - template - E ReadEvent(TTopicReadSessionPtr reader, TTransactionBase& tx); - template - E ReadEvent(TTopicReadSessionPtr reader); - - std::uint64_t GetTopicTabletId(const TActorId& actorId, - const std::string& topicPath, - std::uint32_t partition); - std::vector GetTabletKeys(const TActorId& actorId, - std::uint64_t tabletId); - NPQ::TWriteId GetTransactionWriteId(const TActorId& actorId, - std::uint64_t tabletId); - void SendLongTxLockStatus(const TActorId& actorId, - std::uint64_t tabletId, - const NPQ::TWriteId& writeId, - NKikimrLongTxService::TEvLockStatus::EStatus status); - void WaitForTheTabletToDeleteTheWriteInfo(const TActorId& actorId, - std::uint64_t tabletId, - const NPQ::TWriteId& writeId); - - std::uint64_t GetSchemeShardTabletId(const TActorId& actorId); - - std::unique_ptr Setup; - std::unique_ptr Driver; - std::unique_ptr TableClient; - std::unique_ptr QueryClient; - - std::unordered_map, TTopicWriteSessionContext> TopicWriteSessions; - std::unordered_map TopicReadSessions; - - ui64 SchemaTxId = 1000; -}; - -class TFixtureTable : public TFixture { -protected: - EClientType GetClientType() const override { - return EClientType::Table; - } -}; - -class TFixtureQuery : public TFixture { -protected: - EClientType GetClientType() const override { - return EClientType::Query; - } -}; - -class TFixtureNoClient : public TFixture { -protected: - EClientType GetClientType() const override { - return EClientType::None; - } -}; - -TFixture::TTableRecord::TTableRecord(const std::string& key, const std::string& value) : - Key(key), - Value(value) -{ -} - -void TFixture::SetUp(NUnitTest::TTestContext&) -{ - NKikimr::Tests::TServerSettings settings = TTopicSdkTestSetup::MakeServerSettings(); - settings.SetEnableTopicServiceTx(true); - settings.SetEnableTopicSplitMerge(true); - settings.SetEnablePQConfigTransactionsAtSchemeShard(true); - settings.SetEnableOltpSink(GetEnableOltpSink()); - settings.SetEnableOlapSink(GetEnableOlapSink()); - settings.SetEnableHtapTx(GetEnableHtapTx()); - settings.SetAllowOlapDataQuery(GetAllowOlapDataQuery()); - - Setup = std::make_unique(TEST_CASE_NAME, settings); - - Driver = std::make_unique(Setup->MakeDriver()); - auto tableSettings = NTable::TClientSettings().SessionPoolSettings(NTable::TSessionPoolSettings() - .MaxActiveSessions(3000) - ); - - auto querySettings = NQuery::TClientSettings().SessionPoolSettings(NQuery::TSessionPoolSettings() - .MaxActiveSessions(3000) - ); - - TableClient = std::make_unique(*Driver, tableSettings); - QueryClient = std::make_unique(*Driver, querySettings); -} - -void TFixture::NotifySchemeShard(const TFeatureFlags& flags) -{ - auto request = std::make_unique(); - *request->Record.MutableConfig() = *Setup->GetServer().ServerSettings.AppConfig; - request->Record.MutableConfig()->MutableFeatureFlags()->SetEnablePQConfigTransactionsAtSchemeShard(flags.EnablePQConfigTransactionsAtSchemeShard); - - auto& runtime = Setup->GetRuntime(); - auto actorId = runtime.AllocateEdgeActor(); - - std::uint64_t ssId = GetSchemeShardTabletId(actorId); - - runtime.SendToPipe(ssId, actorId, request.release()); - runtime.GrabEdgeEvent(); -} - -TFixture::TTableSession::TTableSession(NTable::TTableClient& client) - : Session_(Init(client)) -{ -} - -NTable::TSession TFixture::TTableSession::Init(NTable::TTableClient& client) -{ - auto result = client.GetSession().ExtractValueSync(); - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return result.GetSession(); -} - -std::vector TFixture::TTableSession::Execute(const std::string& query, - TTransactionBase* tx, - bool commit, - const TParams& params) -{ - while (true) { - auto txTable = dynamic_cast(tx); - auto txControl = NTable::TTxControl::Tx(*txTable).CommitTx(commit); - - auto result = Session_.ExecuteDataQuery(query, txControl, params).GetValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return std::move(result).ExtractResultSets(); - } - std::this_thread::sleep_for(100ms); - } -} - -TFixture::ISession::TExecuteInTxResult TFixture::TTableSession::ExecuteInTx(const std::string& query, - bool commit, - const TParams& params) -{ - while (true) { - auto txControl = NTable::TTxControl::BeginTx().CommitTx(commit); - - auto result = Session_.ExecuteDataQuery(query, txControl, params).GetValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return {std::move(result).ExtractResultSets(), std::make_unique(*result.GetTransaction())}; - } - std::this_thread::sleep_for(100ms); - } -} - -std::unique_ptr TFixture::TTableSession::BeginTx() -{ - while (true) { - auto result = Session_.BeginTransaction().ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return std::make_unique(result.GetTransaction()); - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TTableSession::CommitTx(TTransactionBase& tx, EStatus status) -{ - auto txTable = dynamic_cast(tx); - while (true) { - auto result = txTable.Commit().ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); - return; - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TTableSession::RollbackTx(TTransactionBase& tx, EStatus status) -{ - auto txTable = dynamic_cast(tx); - while (true) { - auto result = txTable.Rollback().ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); - return; - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TTableSession::Close() -{ - Session_.Close(); -} - -TAsyncStatus TFixture::TTableSession::AsyncCommitTx(TTransactionBase& tx) -{ - auto txTable = dynamic_cast(tx); - return txTable.Commit().Apply([](auto result) { - return TStatus(result.GetValue()); - }); -} - -TFixture::TQuerySession::TQuerySession(NQuery::TQueryClient& client, - const std::string& endpoint, - const std::string& database) - : Session_(Init(client)) - , Endpoint_(endpoint) - , Database_(database) -{ -} - -NQuery::TSession TFixture::TQuerySession::Init(NQuery::TQueryClient& client) -{ - auto result = client.GetSession().ExtractValueSync(); - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return result.GetSession(); -} - -std::vector TFixture::TQuerySession::Execute(const std::string& query, - TTransactionBase* tx, - bool commit, - const TParams& params) -{ - while (true) { - auto txQuery = dynamic_cast(tx); - auto txControl = NQuery::TTxControl::Tx(*txQuery).CommitTx(commit); - - auto result = Session_.ExecuteQuery(query, txControl, params).ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return result.GetResultSets(); - } - std::this_thread::sleep_for(100ms); - } -} - -TFixture::ISession::TExecuteInTxResult TFixture::TQuerySession::ExecuteInTx(const std::string& query, - bool commit, - const TParams& params) -{ - while (true) { - auto txControl = NQuery::TTxControl::BeginTx().CommitTx(commit); - - auto result = Session_.ExecuteQuery(query, txControl, params).ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return {result.GetResultSets(), std::make_unique(*result.GetTransaction())}; - } - std::this_thread::sleep_for(100ms); - } -} - -std::unique_ptr TFixture::TQuerySession::BeginTx() -{ - while (true) { - auto result = Session_.BeginTransaction(NQuery::TTxSettings()).ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); - return std::make_unique(result.GetTransaction()); - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TQuerySession::CommitTx(TTransactionBase& tx, EStatus status) -{ - auto txQuery = dynamic_cast(tx); - while (true) { - auto result = txQuery.Commit().ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); - return; - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TQuerySession::RollbackTx(TTransactionBase& tx, EStatus status) -{ - auto txQuery = dynamic_cast(tx); - while (true) { - auto result = txQuery.Rollback().ExtractValueSync(); - if (result.GetStatus() != EStatus::SESSION_BUSY) { - UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); - return; - } - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::TQuerySession::Close() -{ - // SDK doesn't provide a method to close the session for Query Client, so we use grpc API directly - auto credentials = grpc::InsecureChannelCredentials(); - auto channel = grpc::CreateChannel(TStringType(Endpoint_), credentials); - auto stub = Ydb::Query::V1::QueryService::NewStub(channel); - - grpc::ClientContext context; - context.AddMetadata("x-ydb-database", TStringType(Database_)); - - Ydb::Query::DeleteSessionRequest request; - request.set_session_id(Session_.GetId()); - - Ydb::Query::DeleteSessionResponse response; - auto status = stub->DeleteSession(&context, request, &response); - - NIssue::TIssues issues; - NYdb::NIssue::IssuesFromMessage(response.issues(), issues); - UNIT_ASSERT_C(status.ok(), status.error_message()); - UNIT_ASSERT_VALUES_EQUAL_C(response.status(), Ydb::StatusIds::SUCCESS, issues.ToString()); -} - -TAsyncStatus TFixture::TQuerySession::AsyncCommitTx(TTransactionBase& tx) -{ - auto txQuery = dynamic_cast(tx); - return txQuery.Commit().Apply([](auto result) { - return TStatus(result.GetValue()); - }); -} - -std::unique_ptr TFixture::CreateSession() -{ - switch (GetClientType()) { - case EClientType::Table: { - UNIT_ASSERT_C(TableClient, "TableClient is not initialized"); - return std::make_unique(*TableClient); - } - case EClientType::Query: { - UNIT_ASSERT_C(QueryClient, "QueryClient is not initialized"); - return std::make_unique(*QueryClient, - Setup->GetEndpoint(), - Setup->GetDatabase()); - } - case EClientType::None: { - UNIT_FAIL("CreateSession is forbidden for None client type"); - } - } - - return nullptr; -} - -template -E TFixture::ReadEvent(TTopicReadSessionPtr reader, TTransactionBase& tx) -{ - NTopic::TReadSessionGetEventSettings options; - options.Block(true); - options.MaxEventsCount(1); - options.Tx(tx); - - auto event = reader->GetEvent(options); - UNIT_ASSERT(event); - - auto ev = std::get_if(&*event); - UNIT_ASSERT(ev); - - return *ev; -} - -template -E TFixture::ReadEvent(TTopicReadSessionPtr reader) -{ - auto event = reader->GetEvent(true, 1); - UNIT_ASSERT(event); - - auto ev = std::get_if(&*event); - UNIT_ASSERT(ev); - - return *ev; -} - -void TFixture::CreateTopic(const std::string& path, - const std::string& consumer, - std::size_t partitionCount, - std::optional maxPartitionCount, - const TDuration retention, - bool important) -{ - Setup->CreateTopic(path, consumer, partitionCount, maxPartitionCount, retention, important); -} - -void TFixture::AddConsumer(const std::string& topicPath, - const std::vector& consumers) -{ - NTopic::TTopicClient client(GetDriver()); - NTopic::TAlterTopicSettings settings; - - for (const auto& consumer : consumers) { - settings.BeginAddConsumer(consumer); - } - - auto result = client.AlterTopic(topicPath, settings).GetValueSync(); - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); -} - -void TFixture::SetPartitionWriteSpeed(const std::string& topicName, std::size_t bytesPerSeconds) -{ - NTopic::TTopicClient client(GetDriver()); - NTopic::TAlterTopicSettings settings; - - settings.SetPartitionWriteSpeedBytesPerSecond(bytesPerSeconds); - - auto result = client.AlterTopic(Setup->GetTopicPath(topicName), settings).GetValueSync(); - Y_ENSURE_BT(result.IsSuccess(), ToString(static_cast(result))); -} - -const TDriver& TFixture::GetDriver() const +Y_UNIT_TEST_F(WriteToTopic_Demo_11_Table, TFixtureTable) { - return *Driver; + TestWriteToTopic11(); } -NTable::TTableClient& TFixture::GetTableClient() +Y_UNIT_TEST_F(WriteToTopic_Demo_11_Query, TFixtureQuery) { - return *TableClient; + TestWriteToTopic11(); } -auto TFixture::CreateTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - std::optional partitionId) -> TTopicWriteSessionPtr +Y_UNIT_TEST_F(WriteToTopic_Demo_12_Table, TFixtureTable) { - NTopic::TTopicClient client(GetDriver()); - NTopic::TWriteSessionSettings options; - options.Path(topicPath); - options.ProducerId(messageGroupId); - options.MessageGroupId(messageGroupId); - options.PartitionId(partitionId); - options.Codec(ECodec::RAW); - return client.CreateWriteSession(options); + TestWriteToTopic12(); } -auto TFixture::GetTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - std::optional partitionId) -> TTopicWriteSessionContext& +Y_UNIT_TEST_F(WriteToTopic_Demo_12_Query, TFixtureQuery) { - std::pair key(topicPath, messageGroupId); - auto i = TopicWriteSessions.find(key); - - if (i == TopicWriteSessions.end()) { - TTopicWriteSessionContext context; - context.Session = CreateTopicWriteSession(topicPath, messageGroupId, partitionId); - - TopicWriteSessions.emplace(key, std::move(context)); - - i = TopicWriteSessions.find(key); - } - - return i->second; + TestWriteToTopic12(); } -NTopic::TTopicReadSettings MakeTopicReadSettings(const std::string& topicPath, - std::optional partitionId) +Y_UNIT_TEST_F(WriteToTopic_Demo_13_Table, TFixtureTable) { - TTopicReadSettings options; - options.Path(topicPath); - if (partitionId) { - options.AppendPartitionIds(*partitionId); - } - return options; + TestWriteToTopic13(); } -NTopic::TReadSessionSettings MakeTopicReadSessionSettings(const std::string& topicPath, - const std::string& consumerName, - std::optional partitionId) +Y_UNIT_TEST_F(WriteToTopic_Demo_13_Query, TFixtureQuery) { - NTopic::TReadSessionSettings options; - options.AppendTopics(MakeTopicReadSettings(topicPath, partitionId)); - options.ConsumerName(consumerName); - return options; + TestWriteToTopic13(); } -auto TFixture::CreateTopicReadSession(const std::string& topicPath, - const std::string& consumerName, - std::optional partitionId) -> TTopicReadSessionPtr +Y_UNIT_TEST_F(WriteToTopic_Demo_14_Table, TFixtureTable) { - NTopic::TTopicClient client(GetDriver()); - return client.CreateReadSession(MakeTopicReadSessionSettings(topicPath, - consumerName, - partitionId)); + TestWriteToTopic14(); } -auto TFixture::GetTopicReadSession(const std::string& topicPath, - const std::string& consumerName, - std::optional partitionId) -> TTopicReadSessionPtr +Y_UNIT_TEST_F(WriteToTopic_Demo_14_Query, TFixtureQuery) { - TTopicReadSessionPtr session; - - if (auto i = TopicReadSessions.find(topicPath); i == TopicReadSessions.end()) { - session = CreateTopicReadSession(topicPath, consumerName, partitionId); - auto event = ReadEvent(session); - event.Confirm(); - TopicReadSessions.emplace(topicPath, session); - } else { - session = i->second; - } - - return session; + TestWriteToTopic14(); } -void TFixture::TTopicWriteSessionContext::WaitForContinuationToken() +Y_UNIT_TEST_F(WriteToTopic_Demo_16_Table, TFixtureTable) { - while (!ContinuationToken.has_value()) { - WaitForEvent(); - } + TestWriteToTopic16(); } -void TFixture::TTopicWriteSessionContext::WaitForEvent() +Y_UNIT_TEST_F(WriteToTopic_Demo_16_Query, TFixtureQuery) { - Session->WaitEvent().Wait(); - for (auto& event : Session->GetEvents()) { - if (auto* e = std::get_if(&event)) { - ContinuationToken = std::move(e->ContinuationToken); - } else if (auto* e = std::get_if(&event)) { - for (auto& ack : e->Acks) { - switch (ack.State) { - case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN: - ++WrittenAckCount; - break; - case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN_IN_TX: - ++WrittenInTxAckCount; - break; - default: - break; - } - } - } else if ([[maybe_unused]] auto* e = std::get_if(&event)) { - UNIT_FAIL(""); - } - } + TestWriteToTopic16(); } -void TFixture::TTopicWriteSessionContext::Write(const std::string& message, TTransactionBase* tx) -{ - NTopic::TWriteMessage params(message); - - if (tx) { - params.Tx(*tx); - } - - Session->Write(std::move(*ContinuationToken), - std::move(params)); - - ++WriteCount; - ContinuationToken = std::nullopt; +#define Y_UNIT_TEST_WITH_REBOOTS(name, oldHeadCount, bigBlobsCount, newHeadCount) \ +Y_UNIT_TEST_F(name##_RestartNo_Table, TFixtureTable) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartNo}); \ +} \ +Y_UNIT_TEST_F(name##_RestartNo_Query, TFixtureQuery) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartNo}); \ +} \ +Y_UNIT_TEST_F(name##_RestartBeforeCommit_Table, TFixtureTable) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartBeforeCommit}); \ +} \ +Y_UNIT_TEST_F(name##_RestartBeforeCommit_Query, TFixtureQuery) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartBeforeCommit}); \ +} \ +Y_UNIT_TEST_F(name##_RestartAfterCommit_Table, TFixtureTable) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartAfterCommit}); \ +} \ +Y_UNIT_TEST_F(name##_RestartAfterCommit_Query, TFixtureQuery) { \ + TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartAfterCommit}); \ } -void TFixture::CloseTopicWriteSession(const std::string& topicPath, - const std::string& messageGroupId, - bool force) -{ - std::pair key(topicPath, messageGroupId); - auto i = TopicWriteSessions.find(key); - - UNIT_ASSERT(i != TopicWriteSessions.end()); - - TTopicWriteSessionContext& context = i->second; - - context.Session->Close(force ? TDuration::MilliSeconds(0) : TDuration::Max()); - TopicWriteSessions.erase(key); -} +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_18, 10, 2, 10); +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_19, 10, 0, 10); +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_20, 10, 2, 0); -void TFixture::CloseTopicReadSession(const std::string& topicPath, - const std::string& consumerName) -{ - Y_UNUSED(consumerName); - TopicReadSessions.erase(topicPath); -} +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_21, 0, 2, 10); +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_22, 0, 0, 10); +Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_23, 0, 2, 0); -void TFixture::WriteToTopic(const std::string& topicPath, - const std::string& messageGroupId, - const std::string& message, - TTransactionBase* tx, - std::optional partitionId) +Y_UNIT_TEST_F(WriteToTopic_Demo_24_Table, TFixtureTable) { - TTopicWriteSessionContext& context = GetTopicWriteSession(topicPath, messageGroupId, partitionId); - context.WaitForContinuationToken(); - UNIT_ASSERT(context.ContinuationToken.has_value()); - context.Write(message, tx); + TestWriteToTopic24(); } -std::vector TFixture::ReadFromTopic(const std::string& topicPath, - const std::string& consumerName, - const TDuration& duration, - TTransactionBase* tx, - std::optional partitionId) -{ - std::vector messages; - - TInstant end = TInstant::Now() + duration; - TDuration remain = duration; - - auto session = GetTopicReadSession(topicPath, consumerName, partitionId); - - while (TInstant::Now() < end) { - if (!session->WaitEvent().Wait(remain)) { - return messages; - } - - NTopic::TReadSessionGetEventSettings settings; - if (tx) { - settings.Tx(*tx); - } - - for (auto& event : session->GetEvents(settings)) { - if (auto* e = std::get_if(&event)) { - std::cerr << e->HasCompressedMessages() << " " << e->GetMessagesCount() << std::endl; - for (auto& m : e->GetMessages()) { - messages.emplace_back(m.GetData()); - } - - if (!tx) { - e->Commit(); - } - } - } - - remain = end - TInstant::Now(); - } - - return messages; -} - -void TFixture::WaitForAcks(const std::string& topicPath, const std::string& messageGroupId, std::size_t writtenInTxCount) -{ - std::pair key(topicPath, messageGroupId); - auto i = TopicWriteSessions.find(key); - UNIT_ASSERT(i != TopicWriteSessions.end()); - - auto& context = i->second; - - UNIT_ASSERT(context.AckCount() <= context.WriteCount); - - while (context.AckCount() < context.WriteCount) { - context.WaitForEvent(); - } - - UNIT_ASSERT((context.WrittenAckCount + context.WrittenInTxAckCount) == context.WriteCount); - - if (writtenInTxCount != std::numeric_limits::max()) { - UNIT_ASSERT_VALUES_EQUAL(context.WrittenInTxAckCount, writtenInTxCount); - } -} - -void TFixture::WaitForSessionClose(const std::string& topicPath, - const std::string& messageGroupId, - NYdb::EStatus status) -{ - std::pair key(topicPath, messageGroupId); - auto i = TopicWriteSessions.find(key); - UNIT_ASSERT(i != TopicWriteSessions.end()); - - auto& context = i->second; - - UNIT_ASSERT(context.AckCount() <= context.WriteCount); - - for(bool stop = false; !stop; ) { - context.Session->WaitEvent().Wait(); - for (auto& event : context.Session->GetEvents()) { - if (auto* e = std::get_if(&event)) { - context.ContinuationToken = std::move(e->ContinuationToken); - } else if (auto* e = std::get_if(&event)) { - for (auto& ack : e->Acks) { - switch (ack.State) { - case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN: - ++context.WrittenAckCount; - break; - case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN_IN_TX: - ++context.WrittenInTxAckCount; - break; - default: - break; - } - } - } else if (auto* e = std::get_if(&event)) { - UNIT_ASSERT_VALUES_EQUAL(e->GetStatus(), status); - UNIT_ASSERT_GT(e->GetIssues().Size(), 0); - stop = true; - } - } - } - - UNIT_ASSERT(context.AckCount() <= context.WriteCount); -} - -std::uint64_t TFixture::GetSchemeShardTabletId(const TActorId& actorId) -{ - auto navigate = std::make_unique(); - navigate->DatabaseName = "/Root"; - - NSchemeCache::TSchemeCacheNavigate::TEntry entry; - entry.Path = SplitPath("/Root"); - entry.SyncVersion = true; - entry.ShowPrivatePath = true; - entry.Operation = NSchemeCache::TSchemeCacheNavigate::OpList; - - navigate->ResultSet.push_back(std::move(entry)); - //navigate->UserToken = "root@builtin"; - navigate->Cookie = 12345; - - auto& runtime = Setup->GetRuntime(); - - runtime.Send(MakeSchemeCacheID(), actorId, - new TEvTxProxySchemeCache::TEvNavigateKeySet(navigate.release()), - 0, - true); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT_VALUES_EQUAL(response->Request->Cookie, 12345); - UNIT_ASSERT_VALUES_EQUAL(response->Request->ErrorCount, 0); - - auto& front = response->Request->ResultSet.front(); - - return front.Self->Info.GetSchemeshardId(); -} - -std::uint64_t TFixture::GetTopicTabletId(const TActorId& actorId, const std::string& topicPath, std::uint32_t partition) -{ - auto navigate = std::make_unique(); - navigate->DatabaseName = "/Root"; - - NSchemeCache::TSchemeCacheNavigate::TEntry entry; - entry.Path = SplitPath(TString{topicPath}); - entry.SyncVersion = true; - entry.ShowPrivatePath = true; - entry.Operation = NSchemeCache::TSchemeCacheNavigate::OpList; - - navigate->ResultSet.push_back(std::move(entry)); - //navigate->UserToken = "root@builtin"; - navigate->Cookie = 12345; - - auto& runtime = Setup->GetRuntime(); - - runtime.Send(MakeSchemeCacheID(), actorId, - new TEvTxProxySchemeCache::TEvNavigateKeySet(navigate.release()), - 0, - true); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT_VALUES_EQUAL(response->Request->Cookie, 12345); - UNIT_ASSERT_VALUES_EQUAL(response->Request->ErrorCount, 0); - - auto& front = response->Request->ResultSet.front(); - UNIT_ASSERT(front.PQGroupInfo); - UNIT_ASSERT_GT(front.PQGroupInfo->Description.PartitionsSize(), 0); - UNIT_ASSERT_LT(partition, front.PQGroupInfo->Description.PartitionsSize()); - - for (size_t i = 0; i < front.PQGroupInfo->Description.PartitionsSize(); ++i) { - auto& p = front.PQGroupInfo->Description.GetPartitions(partition); - if (p.GetPartitionId() == partition) { - return p.GetTabletId(); - } - } - - UNIT_FAIL("unknown partition"); - - return std::numeric_limits::max(); -} - -std::vector TFixture::GetTabletKeys(const TActorId& actorId, - std::uint64_t tabletId) -{ - auto request = std::make_unique(); - request->Record.SetCookie(12345); - - auto cmd = request->Record.AddCmdReadRange(); - TString from(1, '\x00'); - TString to(1, '\xFF'); - auto range = cmd->MutableRange(); - range->SetFrom(from); - range->SetIncludeFrom(true); - range->SetTo(to); - range->SetIncludeTo(true); - - auto& runtime = Setup->GetRuntime(); - - runtime.SendToPipe(tabletId, actorId, request.release()); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT(response->Record.HasCookie()); - UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); - UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadRangeResultSize(), 1); - - std::vector keys; - - auto& result = response->Record.GetReadRangeResult(0); - for (size_t i = 0; i < result.PairSize(); ++i) { - auto& kv = result.GetPair(i); - keys.emplace_back(kv.GetKey()); - } - - return keys; -} - -size_t TFixture::GetPQCacheRenameKeysCount() -{ - using namespace NKikimr::NPQ; - - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - - auto request = MakeHolder(); - - runtime.Send(MakePersQueueL2CacheID(), edge, request.Release()); - - TAutoPtr handle; - auto* result = runtime.GrabEdgeEvent(handle); - - return result->RenamedKeys; -} - -std::vector TFixture::Read_Exactly_N_Messages_From_Topic(const std::string& topicPath, - const std::string& consumerName, - std::size_t limit) -{ - std::vector result; - - while (result.size() < limit) { - auto messages = ReadFromTopic(topicPath, consumerName, TDuration::Seconds(2)); - for (auto& m : messages) { - result.push_back(std::move(m)); - } - } - - UNIT_ASSERT_VALUES_EQUAL(result.size(), limit); - - return result; -} - -void TFixture::TestWriteToTopic1() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #4", tx.get()); - - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #5", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #6", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #7", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #8", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #9", tx.get()); - - { - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); - } - - { - auto messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 4); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages[3], "message #4"); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 5); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #5"); - UNIT_ASSERT_VALUES_EQUAL(messages[4], "message #9"); - } -} - -void TFixture::TestWriteToTopic4() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - - auto session = CreateSession(); - auto tx_1 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx_1.get()); - - auto tx_2 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", tx_2.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #4", tx_2.get()); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); - - messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); - - session->CommitTx(*tx_2, EStatus::SUCCESS); - session->CommitTx(*tx_1, EStatus::ABORTED); - - messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #3"); - - messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #4"); -} - -void TFixture::TestWriteToTopic7() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #1", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #2", tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, "message #3"); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, "message #4"); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #5", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #6", tx.get()); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #3"); - UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #4"); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 4); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages[3], "message #6"); - } -} - -void TFixture::TestWriteToTopic9() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx_1 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); - - auto tx_2 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx_2.get()); - - { - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); - } - - session->CommitTx(*tx_2, EStatus::SUCCESS); - session->CommitTx(*tx_1, EStatus::ABORTED); - - { - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #2"); - } -} - -void TFixture::TestWriteToTopic10() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - - { - auto tx_1 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); - - session->CommitTx(*tx_1, EStatus::SUCCESS); - } - - { - auto tx_2 = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx_2.get()); - - session->CommitTx(*tx_2, EStatus::SUCCESS); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #2"); - } -} - -void TFixture::TestWriteToTopic11() -{ - for (auto endOfTransaction : {Commit, Rollback, CloseTableSession}) { - TestTheCompletionOfATransaction({.Topics={"topic_A"}, .EndOfTransaction = endOfTransaction}); - TestTheCompletionOfATransaction({.Topics={"topic_A", "topic_B"}, .EndOfTransaction = endOfTransaction}); - } -} - -void TFixture::TestWriteToTopic24() -{ - // - // the test verifies a transaction in which data is written to a topic and to a table - // - CreateTopic("topic_A"); - CreateTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - session->CommitTx(*tx, EStatus::SUCCESS); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - UNIT_ASSERT_VALUES_EQUAL(messages[0], MakeJsonDoc(records)); - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - - CheckTabletKeys("topic_A"); -} - -void TFixture::TestWriteToTopic26() -{ - // - // the test verifies a transaction in which data is read from a partition of one topic and written to - // another partition of this topic - // - const std::uint32_t PARTITION_0 = 0; - const std::uint32_t PARTITION_1 = 1; - - CreateTopic("topic_A", TEST_CONSUMER, 2); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", nullptr, PARTITION_0); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", nullptr, PARTITION_0); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", nullptr, PARTITION_0); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), PARTITION_0); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 3); - - for (const auto& m : messages) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, m, tx.get(), PARTITION_1); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), nullptr, PARTITION_1); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 3); -} - -void TFixture::TestWriteToTopic27() -{ - CreateTopic("topic_A", TEST_CONSUMER); - CreateTopic("topic_B", TEST_CONSUMER); - CreateTopic("topic_C", TEST_CONSUMER); - - for (size_t i = 0; i < 2; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", nullptr, 0); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", nullptr, 0); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - - WriteToTopic("topic_C", TEST_MESSAGE_GROUP_ID, messages[0], tx.get(), 0); - - messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); - - WriteToTopic("topic_C", TEST_MESSAGE_GROUP_ID, messages[0], tx.get(), 0); - - session->CommitTx(*tx, EStatus::SUCCESS); - - messages = ReadFromTopic("topic_C", TEST_CONSUMER, TDuration::Seconds(2), nullptr, 0); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); - - DumpPQTabletKeys("topic_A"); - DumpPQTabletKeys("topic_B"); - DumpPQTabletKeys("topic_C"); - } -} - -auto TFixture::GetAvgWriteBytes(const std::string& topicName, - std::uint32_t partitionId) -> TAvgWriteBytes -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partitionId); - - runtime.SendToPipe(tabletId, edge, new NKikimr::TEvPersQueue::TEvStatus()); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT_VALUES_EQUAL(tabletId, response->Record.GetTabletId()); - - TAvgWriteBytes result; - - for (std::size_t i = 0; i < response->Record.PartResultSize(); ++i) { - const auto& partition = response->Record.GetPartResult(i); - if (partition.GetPartition() == static_cast(partitionId)) { - result.PerSec = partition.GetAvgWriteSpeedPerSec(); - result.PerMin = partition.GetAvgWriteSpeedPerMin(); - result.PerHour = partition.GetAvgWriteSpeedPerHour(); - result.PerDay = partition.GetAvgWriteSpeedPerDay(); - break; - } - } - - return result; -} - -bool TFixture::GetEnableOltpSink() const -{ - return false; -} - -bool TFixture::GetEnableOlapSink() const -{ - return false; -} - -bool TFixture::GetEnableHtapTx() const -{ - return false; -} - -bool TFixture::GetAllowOlapDataQuery() const -{ - return false; -} - -NPQ::TWriteId TFixture::GetTransactionWriteId(const TActorId& actorId, - std::uint64_t tabletId) -{ - auto request = std::make_unique(); - request->Record.SetCookie(12345); - request->Record.AddCmdRead()->SetKey("_txinfo"); - - auto& runtime = Setup->GetRuntime(); - - runtime.SendToPipe(tabletId, actorId, request.release()); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT(response->Record.HasCookie()); - UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); - UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadResultSize(), 1); - - auto& read = response->Record.GetReadResult(0); - - NKikimrPQ::TTabletTxInfo info; - UNIT_ASSERT(info.ParseFromString(read.GetValue())); - - UNIT_ASSERT_VALUES_EQUAL(info.TxWritesSize(), 1); - - auto& writeInfo = info.GetTxWrites(0); - UNIT_ASSERT(writeInfo.HasWriteId()); - - return NPQ::GetWriteId(writeInfo); -} - -void TFixture::SendLongTxLockStatus(const TActorId& actorId, - std::uint64_t tabletId, - const NPQ::TWriteId& writeId, - NKikimrLongTxService::TEvLockStatus::EStatus status) -{ - auto event = - std::make_unique(writeId.KeyId, writeId.NodeId, - status); - auto& runtime = Setup->GetRuntime(); - runtime.SendToPipe(tabletId, actorId, event.release()); -} - -void TFixture::WaitForTheTabletToDeleteTheWriteInfo(const TActorId& actorId, - std::uint64_t tabletId, - const NPQ::TWriteId& writeId) -{ - while (true) { - auto request = std::make_unique(); - request->Record.SetCookie(12345); - request->Record.AddCmdRead()->SetKey("_txinfo"); - - auto& runtime = Setup->GetRuntime(); - - runtime.SendToPipe(tabletId, actorId, request.release()); - auto response = runtime.GrabEdgeEvent(); - - UNIT_ASSERT(response->Record.HasCookie()); - UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); - UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadResultSize(), 1); - - auto& read = response->Record.GetReadResult(0); - - NKikimrPQ::TTabletTxInfo info; - UNIT_ASSERT(info.ParseFromString(read.GetValue())); - - bool found = false; - - for (size_t i = 0; i < info.TxWritesSize(); ++i) { - auto& writeInfo = info.GetTxWrites(i); - UNIT_ASSERT(writeInfo.HasWriteId()); - if ((NPQ::GetWriteId(writeInfo) == writeId) && writeInfo.HasOriginalPartitionId()) { - found = true; - break; - } - } - - if (!found) { - break; - } - - std::this_thread::sleep_for(100ms); - } -} - -void TFixture::RestartPQTablet(const std::string& topicName, std::uint32_t partition) -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partition); - runtime.SendToPipe(tabletId, edge, new TEvents::TEvPoison()); - - std::this_thread::sleep_for(2s); -} - -void TFixture::DeleteSupportivePartition(const std::string& topicName, std::uint32_t partition) -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partition); - NPQ::TWriteId writeId = GetTransactionWriteId(edge, tabletId); - - SendLongTxLockStatus(edge, tabletId, writeId, NKikimrLongTxService::TEvLockStatus::STATUS_NOT_FOUND); - - WaitForTheTabletToDeleteTheWriteInfo(edge, tabletId, writeId); -} - -void TFixture::CheckTabletKeys(const std::string& topicName) -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, 0); - - const std::unordered_set types { - NPQ::TKeyPrefix::TypeInfo, - NPQ::TKeyPrefix::TypeData, - NPQ::TKeyPrefix::TypeTmpData, - NPQ::TKeyPrefix::TypeMeta, - NPQ::TKeyPrefix::TypeTxMeta, - }; - - bool found; - std::vector keys; - for (size_t i = 0; i < 20; ++i) { - keys = GetTabletKeys(edge, tabletId); - - found = false; - for (const auto& key : keys) { - UNIT_ASSERT_GT(key.size(), 0); - if (key[0] == '_') { - continue; - } - - if (types.contains(key[0])) { - found = false; - break; - } - } - - if (!found) { - break; - } - - std::this_thread::sleep_for(100ms); - } - - if (found) { - std::cerr << "keys for tablet " << tabletId << ":" << std::endl; - for (const auto& k : keys) { - std::cerr << k << std::endl; - } - std::cerr << "=============" << std::endl; - - UNIT_FAIL("unexpected keys for tablet " << tabletId); - } -} - -void TFixture::DumpPQTabletKeys(const std::string& topicName) -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, 0); - auto keys = GetTabletKeys(edge, tabletId); - for (const auto& key : keys) { - std::cerr << key << std::endl; - } -} - -void TFixture::PQTabletPrepareFromResource(const std::string& topicPath, - std::uint32_t partitionId, - const std::string& resourceName) -{ - auto& runtime = Setup->GetRuntime(); - TActorId edge = runtime.AllocateEdgeActor(); - std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicPath, partitionId); - - auto request = MakeHolder(); - std::size_t count = 0; - - for (TStringStream stream(NResource::Find(resourceName)); true; ++count) { - TString key, encoded; - - if (!stream.ReadTo(key, ' ')) { - break; - } - encoded = stream.ReadLine(); - - auto decoded = Base64Decode(encoded); - TStringInput decodedStream(decoded); - TBZipDecompress decompressor(&decodedStream); - - auto* cmd = request->Record.AddCmdWrite(); - cmd->SetKey(key); - cmd->SetValue(decompressor.ReadAll()); - } - - runtime.SendToPipe(tabletId, edge, request.Release(), 0, GetPipeConfigWithRetries()); - - TAutoPtr handle; - auto* response = runtime.GrabEdgeEvent(handle); - UNIT_ASSERT(response); - UNIT_ASSERT(response->Record.HasStatus()); - UNIT_ASSERT_EQUAL(response->Record.GetStatus(), NMsgBusProxy::MSTATUS_OK); - - UNIT_ASSERT_VALUES_EQUAL(response->Record.WriteResultSize(), count); - - for (std::size_t i = 0; i < response->Record.WriteResultSize(); ++i) { - const auto &result = response->Record.GetWriteResult(i); - UNIT_ASSERT(result.HasStatus()); - UNIT_ASSERT_EQUAL(result.GetStatus(), NKikimrProto::OK); - } -} - -void TFixture::TestTheCompletionOfATransaction(const TTransactionCompletionTestDescription& d) -{ - for (auto& topic : d.Topics) { - CreateTopic(topic); - } - - { - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (auto& topic : d.Topics) { - WriteToTopic(topic, TEST_MESSAGE_GROUP_ID, "message", tx.get()); - // TODO: нужен callback для RollbakTx - WaitForAcks(topic, TEST_MESSAGE_GROUP_ID); - } - - switch (d.EndOfTransaction) { - case Commit: - session->CommitTx(*tx, EStatus::SUCCESS); - break; - case Rollback: - session->RollbackTx(*tx, EStatus::SUCCESS); - break; - case CloseTableSession: - break; - } - } - - std::this_thread::sleep_for(std::chrono::seconds(5)); - - for (auto& topic : d.Topics) { - CheckTabletKeys(topic); - } - - for (auto& topic : d.Topics) { - CloseTopicWriteSession(topic, TEST_MESSAGE_GROUP_ID); - } -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_11_Table, TFixtureTable) -{ - TestWriteToTopic11(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_11_Query, TFixtureQuery) -{ - TestWriteToTopic11(); -} - -void TFixture::TestWriteToTopic12() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - DeleteSupportivePartition("topic_A", 0); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - WaitForSessionClose("topic_A", TEST_MESSAGE_GROUP_ID, NYdb::EStatus::PRECONDITION_FAILED); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_12_Table, TFixtureTable) -{ - TestWriteToTopic12(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_12_Query, TFixtureQuery) -{ - TestWriteToTopic12(); -} - -void TFixture::TestWriteToTopic13() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message", tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - DeleteSupportivePartition("topic_A", 0); - - session->CommitTx(*tx, EStatus::ABORTED); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_13_Table, TFixtureTable) -{ - TestWriteToTopic13(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_13_Query, TFixtureQuery) -{ - TestWriteToTopic13(); -} - -void TFixture::TestWriteToTopic14() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - DeleteSupportivePartition("topic_A", 0); - - CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - - session->CommitTx(*tx, EStatus::ABORTED); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_14_Table, TFixtureTable) -{ - TestWriteToTopic14(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_14_Query, TFixtureQuery) -{ - TestWriteToTopic14(); -} - -void TFixture::TestWriteToTopic16() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - - RestartPQTablet("topic_A", 0); - - session->CommitTx(*tx, EStatus::SUCCESS); - - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); - UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #2"); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_16_Table, TFixtureTable) -{ - TestWriteToTopic16(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_16_Query, TFixtureQuery) -{ - TestWriteToTopic16(); -} - -void TFixture::TestTxWithBigBlobs(const TTestTxWithBigBlobsParams& params) -{ - size_t oldHeadMsgCount = 0; - size_t bigBlobMsgCount = 0; - size_t newHeadMsgCount = 0; - - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (size_t i = 0; i < params.OldHeadCount; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(100'000, 'x')); - ++oldHeadMsgCount; - } - - for (size_t i = 0; i < params.BigBlobsCount; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(7'000'000, 'x'), tx.get()); - ++bigBlobMsgCount; - } - - for (size_t i = 0; i < params.NewHeadCount; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(100'000, 'x'), tx.get()); - ++newHeadMsgCount; - } - - if (params.RestartMode == ERestartBeforeCommit) { - RestartPQTablet("topic_A", 0); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - if (params.RestartMode == ERestartAfterCommit) { - RestartPQTablet("topic_A", 0); - } - - std::vector messages; - for (size_t i = 0; (i < 10) && (messages.size() < (oldHeadMsgCount + bigBlobMsgCount + newHeadMsgCount)); ++i) { - auto block = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); - for (auto& m : block) { - messages.push_back(std::move(m)); - } - } - - UNIT_ASSERT_VALUES_EQUAL(messages.size(), oldHeadMsgCount + bigBlobMsgCount + newHeadMsgCount); - - size_t start = 0; - - for (size_t i = 0; i < oldHeadMsgCount; ++i) { - UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 100'000); - } - start += oldHeadMsgCount; - - for (size_t i = 0; i < bigBlobMsgCount; ++i) { - UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 7'000'000); - } - start += bigBlobMsgCount; - - for (size_t i = 0; i < newHeadMsgCount; ++i) { - UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 100'000); - } -} - -#define Y_UNIT_TEST_WITH_REBOOTS(name, oldHeadCount, bigBlobsCount, newHeadCount) \ -Y_UNIT_TEST_F(name##_RestartNo_Table, TFixtureTable) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartNo}); \ -} \ -Y_UNIT_TEST_F(name##_RestartNo_Query, TFixtureQuery) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartNo}); \ -} \ -Y_UNIT_TEST_F(name##_RestartBeforeCommit_Table, TFixtureTable) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartBeforeCommit}); \ -} \ -Y_UNIT_TEST_F(name##_RestartBeforeCommit_Query, TFixtureQuery) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartBeforeCommit}); \ -} \ -Y_UNIT_TEST_F(name##_RestartAfterCommit_Table, TFixtureTable) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartAfterCommit}); \ -} \ -Y_UNIT_TEST_F(name##_RestartAfterCommit_Query, TFixtureQuery) { \ - TestTxWithBigBlobs({.OldHeadCount = oldHeadCount, .BigBlobsCount = bigBlobsCount, .NewHeadCount = newHeadCount, .RestartMode = ERestartAfterCommit}); \ -} - -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_18, 10, 2, 10); -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_19, 10, 0, 10); -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_20, 10, 2, 0); - -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_21, 0, 2, 10); -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_22, 0, 0, 10); -Y_UNIT_TEST_WITH_REBOOTS(WriteToTopic_Demo_23, 0, 2, 0); - -void TFixture::CreateTable(const std::string& tablePath) -{ - UNIT_ASSERT(!tablePath.empty()); - - std::string path = (tablePath[0] != '/') ? ("/Root/" + tablePath) : tablePath; - - auto createSessionResult = GetTableClient().CreateSession().ExtractValueSync(); - UNIT_ASSERT_C(createSessionResult.IsSuccess(), createSessionResult.GetIssues().ToString()); - auto session = createSessionResult.GetSession(); - - auto desc = NTable::TTableBuilder() - .AddNonNullableColumn("key", EPrimitiveType::Utf8) - .AddNonNullableColumn("value", EPrimitiveType::Utf8) - .SetPrimaryKeyColumn("key") - .Build(); - auto result = session.CreateTable(path, std::move(desc)).GetValueSync(); - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); -} - -auto TFixture::MakeTableRecords() -> std::vector -{ - std::vector records; - records.emplace_back("key-1", "value-1"); - records.emplace_back("key-2", "value-2"); - records.emplace_back("key-3", "value-3"); - records.emplace_back("key-4", "value-4"); - return records; -} - -auto TFixture::MakeJsonDoc(const std::vector& records) -> std::string -{ - auto makeJsonObject = [](const TTableRecord& r) { - return Sprintf(R"({"key":"%s", "value":"%s"})", - r.Key.data(), - r.Value.data()); - }; - - if (records.empty()) { - return "[]"; - } - - std::string s = "["; - - s += makeJsonObject(records.front()); - for (auto i = records.begin() + 1; i != records.end(); ++i) { - s += ", "; - s += makeJsonObject(*i); - } - s += "]"; - - return s; -} - -void TFixture::UpsertToTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx) -{ - auto query = Sprintf("DECLARE $key AS Utf8;" - "DECLARE $value AS Utf8;" - "UPSERT INTO `%s` (key, value) VALUES ($key, $value);", - tablePath.data()); - - for (const auto& r : records) { - auto params = TParamsBuilder() - .AddParam("$key").Utf8(r.Key).Build() - .AddParam("$value").Utf8(r.Value).Build() - .Build(); - - session.Execute(query, tx, false, params); - } -} - -void TFixture::InsertToTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx) -{ - auto query = Sprintf("DECLARE $key AS Utf8;" - "DECLARE $value AS Utf8;" - "INSERT INTO `%s` (key, value) VALUES ($key, $value);", - tablePath.data()); - - for (const auto& r : records) { - auto params = TParamsBuilder() - .AddParam("$key").Utf8(r.Key).Build() - .AddParam("$value").Utf8(r.Value).Build() - .Build(); - - session.Execute(query, tx, false, params); - } -} - -void TFixture::DeleteFromTable(const std::string& tablePath, - const std::vector& records, - ISession& session, - TTransactionBase* tx) -{ - auto query = Sprintf("DECLARE $key AS Utf8;" - "DECLARE $value AS Utf8;" - "DELETE FROM `%s` ON (key, value) VALUES ($key, $value);", - tablePath.data()); - - for (const auto& r : records) { - auto params = TParamsBuilder() - .AddParam("$key").Utf8(r.Key).Build() - .AddParam("$value").Utf8(r.Value).Build() - .Build(); - - session.Execute(query, tx, false, params); - } -} - -std::size_t TFixture::GetTableRecordsCount(const std::string& tablePath) -{ - auto query = Sprintf(R"(SELECT COUNT(*) FROM `%s`)", - tablePath.data()); - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto result = session->Execute(query, tx.get()); - - NYdb::TResultSetParser parser(result.at(0)); - UNIT_ASSERT(parser.TryNextRow()); - - return parser.ColumnParser(0).GetUint64(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_24_Table, TFixtureTable) -{ - TestWriteToTopic24(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_24_Query, TFixtureQuery) +Y_UNIT_TEST_F(WriteToTopic_Demo_24_Query, TFixtureQuery) { TestWriteToTopic24(); } -Y_UNIT_TEST_F(WriteToTopic_Demo_27_Table, TFixtureTable) -{ - TestWriteToTopic27(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_27_Query, TFixtureQuery) -{ - TestWriteToTopic27(); -} - -void TFixture::WriteMessagesInTx(std::size_t big, std::size_t small) -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (std::size_t i = 0; i < big; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(7'000'000, 'x'), tx.get(), 0); - } - - for (std::size_t i = 0; i < small; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(16'384, 'x'), tx.get(), 0); - } - - session->CommitTx(*tx, EStatus::SUCCESS); -} - -void TFixture::TestWriteToTopic38() -{ - WriteMessagesInTx(2, 202); - WriteMessagesInTx(2, 200); - WriteMessagesInTx(0, 1); - WriteMessagesInTx(4, 0); - WriteMessagesInTx(0, 1); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_38_Table, TFixtureTable) -{ - TestWriteToTopic38(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_38_Query, TFixtureQuery) -{ - TestWriteToTopic38(); -} - -void TFixture::TestWriteToTopic40() -{ - // The recording stream will run into a quota. Before the commit, the client will receive confirmations - // for some of the messages. The `CommitTx` call will wait for the rest. - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (std::size_t k = 0; k < 100; ++k) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_40_Table, TFixtureTable) -{ - TestWriteToTopic40(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_40_Query, TFixtureQuery) -{ - TestWriteToTopic40(); -} - -void TFixture::TestWriteToTopic41() -{ - // If the recording session does not wait for confirmations, the commit will fail - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (std::size_t k = 0; k < 100; ++k) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); - } - - CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID, true); // force close - - session->CommitTx(*tx, EStatus::SESSION_EXPIRED); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_41_Table, TFixtureTable) -{ - TestWriteToTopic41(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_41_Query, TFixtureQuery) -{ - TestWriteToTopic41(); -} - -void TFixture::TestWriteToTopic42() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (std::size_t k = 0; k < 100; ++k) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); - } - - CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); // gracefully close - - session->CommitTx(*tx, EStatus::SUCCESS); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_42_Table, TFixtureTable) -{ - TestWriteToTopic42(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_42_Query, TFixtureQuery) -{ - TestWriteToTopic42(); -} - -void TFixture::TestWriteToTopic43() -{ - // The recording stream will run into a quota. Before the commit, the client will receive confirmations - // for some of the messages. The `ExecuteDataQuery` call will wait for the rest. - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - for (std::size_t k = 0; k < 100; ++k) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); - } - - session->Execute("SELECT 1", tx.get()); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_43_Table, TFixtureTable) -{ - TestWriteToTopic43(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_43_Query, TFixtureQuery) -{ - TestWriteToTopic43(); -} - -void TFixture::TestWriteToTopic44() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - - auto [_, tx] = session->ExecuteInTx("SELECT 1", false); - - for (std::size_t k = 0; k < 100; ++k) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); - } - - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(60)); - UNIT_ASSERT_EQUAL(messages.size(), 0u); - - session->Execute("SELECT 2", tx.get()); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_44_Table, TFixtureTable) -{ - TestWriteToTopic44(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_44_Query, TFixtureQuery) -{ - TestWriteToTopic44(); -} - -Y_UNIT_TEST_F(ReadRuleGeneration, TFixtureNoClient) -{ - // There was a server - NotifySchemeShard({.EnablePQConfigTransactionsAtSchemeShard = false}); - - // Users have created their own topic on it - CreateTopic(TEST_TOPIC); - - // And they wrote their messages into it - WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-1"); - WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-2"); - WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-3"); - - // And he had a consumer - AddConsumer(TEST_TOPIC, {"consumer-1"}); - - // We read messages from the topic and committed offsets - Read_Exactly_N_Messages_From_Topic(TEST_TOPIC, "consumer-1", 3); - CloseTopicReadSession(TEST_TOPIC, "consumer-1"); - - // And then the Logbroker team turned on the feature flag - NotifySchemeShard({.EnablePQConfigTransactionsAtSchemeShard = true}); - - // Users continued to write to the topic - WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-4"); - - // Users have added new consumers - AddConsumer(TEST_TOPIC, {"consumer-2"}); - - // And they wanted to continue reading their messages - Read_Exactly_N_Messages_From_Topic(TEST_TOPIC, "consumer-1", 1); -} - -void TFixture::CheckAvgWriteBytes(const std::string& topicPath, - std::uint32_t partitionId, - std::size_t minSize, std::size_t maxSize) -{ -#define UNIT_ASSERT_AVGWRITEBYTES(v, minSize, maxSize) \ - UNIT_ASSERT_LE_C(minSize, v, ", actual " << minSize << " > " << v); \ - UNIT_ASSERT_LE_C(v, maxSize, ", actual " << v << " > " << maxSize); - - auto avgWriteBytes = GetAvgWriteBytes(topicPath, partitionId); - - UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerSec, minSize, maxSize); - UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerMin, minSize, maxSize); - UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerHour, minSize, maxSize); - UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerDay, minSize, maxSize); - -#undef UNIT_ASSERT_AVGWRITEBYTES -} - -void TFixture::SplitPartition(const std::string& topicName, - std::uint32_t partitionId, - const std::string& boundary) -{ - NKikimr::NPQ::NTest::SplitPartition(Setup->GetRuntime(), - ++SchemaTxId, - TString(topicName), - partitionId, - TString(boundary)); -} - -void TFixture::TestWriteToTopic45() -{ - // Writing to a topic in a transaction affects the `AvgWriteBytes` indicator - CreateTopic("topic_A", TEST_CONSUMER, 2); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - std::string message(1'000, 'x'); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); - - session->CommitTx(*tx, EStatus::SUCCESS); - - std::size_t minSize = (message.size() + TEST_MESSAGE_GROUP_ID_1.size()) * 2; - std::size_t maxSize = minSize + 200; - - CheckAvgWriteBytes("topic_A", 0, minSize, maxSize); - - minSize = (message.size() + TEST_MESSAGE_GROUP_ID_2.size()); - maxSize = minSize + 200; - - CheckAvgWriteBytes("topic_A", 1, minSize, maxSize); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_45_Table, TFixtureTable) -{ - TestWriteToTopic45(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_45_Query, TFixtureQuery) -{ - TestWriteToTopic45(); -} - -void TFixture::TestWriteToTopic46() -{ - // The `split` operation of the topic partition affects the writing in the transaction. - // The transaction commit should fail with an error - CreateTopic("topic_A", TEST_CONSUMER, 2, 10); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - std::string message(1'000, 'x'); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); - - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2); - - SplitPartition("topic_A", 1, "\xC0"); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); - - session->CommitTx(*tx, EStatus::ABORTED); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_46_Table, TFixtureTable) -{ - TestWriteToTopic46(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_46_Query, TFixtureQuery) -{ - TestWriteToTopic46(); -} - -void TFixture::TestWriteToTopic47() -{ - // The `split` operation of the topic partition does not affect the reading in the transaction. - CreateTopic("topic_A", TEST_CONSUMER, 2, 10); - - std::string message(1'000, 'x'); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1); - - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_1); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2); - - SplitPartition("topic_A", 1, "\xC0"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); - - CloseTopicReadSession("topic_A", TEST_CONSUMER); - - messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 1); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); - - session->CommitTx(*tx, EStatus::SUCCESS); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_47_Table, TFixtureTable) -{ - TestWriteToTopic47(); -} - -Y_UNIT_TEST_F(WriteToTopic_Demo_47_Query, TFixtureQuery) -{ - TestWriteToTopic47(); -} - -void TFixture::TestWriteRandomSizedMessagesInWideTransactions() -{ - // The test verifies the simultaneous execution of several transactions. There is a topic - // with PARTITIONS_COUNT partitions. In each transaction, the test writes to all the partitions. - // The size of the messages is random. Such that both large blobs in the body and small ones in - // the head of the partition are obtained. Message sizes are multiples of 500 KB. This way we - // will make sure that when committing transactions, the division into blocks is taken into account. - - const std::size_t PARTITIONS_COUNT = 20; - const std::size_t TXS_COUNT = 10; - - CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); - - SetPartitionWriteSpeed("topic_A", 50'000'000); - - std::vector> sessions; - std::vector> transactions; - - // We open TXS_COUNT transactions and write messages to the topic. - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - sessions.push_back(CreateSession()); - auto& session = sessions.back(); - - transactions.push_back(session->BeginTx()); - auto& tx = transactions.back(); - - for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { - std::string sourceId = TEST_MESSAGE_GROUP_ID; - sourceId += "_"; - sourceId += ToString(i); - sourceId += "_"; - sourceId += ToString(j); - - std::size_t count = RandomNumber(20) + 3; - WriteToTopic("topic_A", sourceId, std::string(512 * 1000 * count, 'x'), tx.get(), j); - - WaitForAcks("topic_A", sourceId); - } - } - - // We are doing an asynchronous commit of transactions. They will be executed simultaneously. - std::vector futures; - - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); - } - - // All transactions must be completed successfully. - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures[i].Wait(); - const auto& result = futures[i].GetValueSync(); - UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - } -} - -Y_UNIT_TEST_F(Write_Random_Sized_Messages_In_Wide_Transactions_Table, TFixtureTable) +Y_UNIT_TEST_F(WriteToTopic_Demo_27_Table, TFixtureTable) { - TestWriteRandomSizedMessagesInWideTransactions(); + TestWriteToTopic27(); } -Y_UNIT_TEST_F(Write_Random_Sized_Messages_In_Wide_Transactions_Query, TFixtureQuery) +Y_UNIT_TEST_F(WriteToTopic_Demo_27_Query, TFixtureQuery) { - TestWriteRandomSizedMessagesInWideTransactions(); + TestWriteToTopic27(); } -void TFixture::TestWriteOnlyBigMessagesInWideTransactions() +Y_UNIT_TEST_F(WriteToTopic_Demo_38_Table, TFixtureTable) { - // The test verifies the simultaneous execution of several transactions. There is a topic `topic_A` and - // it contains a `PARTITIONS_COUNT' of partitions. In each transaction, the test writes to all partitions. - // The size of the messages is chosen so that only large blobs are recorded in the transaction and there - // are no records in the head. Thus, we verify that transaction bundling is working correctly. - - const std::size_t PARTITIONS_COUNT = 20; - const std::size_t TXS_COUNT = 100; - - CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); - - SetPartitionWriteSpeed("topic_A", 50'000'000); - - std::vector> sessions; - std::vector> transactions; - - // We open TXS_COUNT transactions and write messages to the topic. - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - sessions.push_back(CreateSession()); - auto& session = sessions.back(); - - transactions.push_back(session->BeginTx()); - auto& tx = transactions.back(); - - for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { - std::string sourceId = TEST_MESSAGE_GROUP_ID; - sourceId += "_"; - sourceId += ToString(i); - sourceId += "_"; - sourceId += ToString(j); - - WriteToTopic("topic_A", sourceId, std::string(6'500'000, 'x'), tx.get(), j); - - WaitForAcks("topic_A", sourceId); - } - } - - // We are doing an asynchronous commit of transactions. They will be executed simultaneously. - std::vector futures; - - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); - } - - // All transactions must be completed successfully. - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures[i].Wait(); - const auto& result = futures[i].GetValueSync(); - UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - } + TestWriteToTopic38(); } -Y_UNIT_TEST_F(Write_Only_Big_Messages_In_Wide_Transactions_Table, TFixtureTable) +Y_UNIT_TEST_F(WriteToTopic_Demo_38_Query, TFixtureQuery) { - TestWriteOnlyBigMessagesInWideTransactions(); + TestWriteToTopic38(); } -Y_UNIT_TEST_F(Write_Only_Big_Messages_In_Wide_Transactions_Query, TFixtureQuery) +Y_UNIT_TEST_F(WriteToTopic_Demo_40_Table, TFixtureTable) { - TestWriteOnlyBigMessagesInWideTransactions(); + TestWriteToTopic40(); } -void TFixture::TestTransactionsConflictOnSeqNo() +Y_UNIT_TEST_F(WriteToTopic_Demo_40_Query, TFixtureQuery) { - const std::uint32_t PARTITIONS_COUNT = 20; - const std::size_t TXS_COUNT = 100; - - CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); - - SetPartitionWriteSpeed("topic_A", 50'000'000); - - auto session = CreateSession(); - std::vector> topicWriteSessions; - - for (std::uint32_t i = 0; i < PARTITIONS_COUNT; ++i) { - std::string sourceId = TEST_MESSAGE_GROUP_ID; - sourceId += "_"; - sourceId += ToString(i); - - NTopic::TTopicClient client(GetDriver()); - NTopic::TWriteSessionSettings options; - options.Path(Setup->GetTopicPath("topic_A")); - options.ProducerId(sourceId); - options.MessageGroupId(sourceId); - options.PartitionId(i); - options.Codec(ECodec::RAW); - - auto session = client.CreateSimpleBlockingWriteSession(options); - - topicWriteSessions.push_back(std::move(session)); - } - - std::vector> sessions; - std::vector> transactions; - - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - sessions.push_back(CreateSession()); - auto& session = sessions.back(); - - transactions.push_back(session->BeginTx()); - auto& tx = transactions.back(); - - for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { - std::string sourceId = TEST_MESSAGE_GROUP_ID; - sourceId += "_"; - sourceId += ToString(j); - - for (std::size_t k = 0, count = RandomNumber(20) + 1; k < count; ++k) { - const std::string data(RandomNumber(1'000) + 100, 'x'); - NTopic::TWriteMessage params(data); - params.Tx(*tx); - - topicWriteSessions[j]->Write(std::move(params)); - } - } - } - - std::vector futures; + TestWriteToTopic40(); +} - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); - } +Y_UNIT_TEST_F(WriteToTopic_Demo_41_Table, TFixtureTable) +{ + TestWriteToTopic41(); +} - // Some transactions should end with the error `ABORTED` - std::size_t successCount = 0; +Y_UNIT_TEST_F(WriteToTopic_Demo_41_Query, TFixtureQuery) +{ + TestWriteToTopic41(); +} - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures[i].Wait(); - const auto& result = futures[i].GetValueSync(); - switch (result.GetStatus()) { - case EStatus::SUCCESS: - ++successCount; - break; - case EStatus::ABORTED: - break; - default: - UNIT_FAIL("unexpected status: " + ToString(static_cast(result))); - break; - } - } +Y_UNIT_TEST_F(WriteToTopic_Demo_42_Table, TFixtureTable) +{ + TestWriteToTopic42(); +} - UNIT_ASSERT_UNEQUAL(successCount, TXS_COUNT); +Y_UNIT_TEST_F(WriteToTopic_Demo_42_Query, TFixtureQuery) +{ + TestWriteToTopic42(); } -Y_UNIT_TEST_F(Transactions_Conflict_On_SeqNo_Table, TFixtureTable) +Y_UNIT_TEST_F(WriteToTopic_Demo_43_Table, TFixtureTable) { - TestTransactionsConflictOnSeqNo(); + TestWriteToTopic43(); } -Y_UNIT_TEST_F(Transactions_Conflict_On_SeqNo_Query, TFixtureQuery) +Y_UNIT_TEST_F(WriteToTopic_Demo_43_Query, TFixtureQuery) { - TestTransactionsConflictOnSeqNo(); + TestWriteToTopic43(); } -class TFixtureSinks : public TFixture { -protected: - void CreateRowTable(const std::string& path); - void CreateColumnTable(const std::string& tablePath); +Y_UNIT_TEST_F(ReadRuleGeneration, TFixtureNoClient) +{ + // There was a server + NotifySchemeShard({.EnablePQConfigTransactionsAtSchemeShard = false}); + + // Users have created their own topic on it + CreateTopic(TEST_TOPIC); + + // And they wrote their messages into it + WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-1"); + WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-2"); + WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-3"); - bool GetEnableOltpSink() const override; - bool GetEnableOlapSink() const override; - bool GetEnableHtapTx() const override; - bool GetAllowOlapDataQuery() const override; + // And he had a consumer + AddConsumer(TEST_TOPIC, {"consumer-1"}); - void TestSinksOltpWriteToTopic5(); + // We read messages from the topic and committed offsets + Read_Exactly_N_Messages_From_Topic(TEST_TOPIC, "consumer-1", 3); + CloseTopicReadSession(TEST_TOPIC, "consumer-1"); - void TestSinksOltpWriteToTopicAndTable2(); - void TestSinksOltpWriteToTopicAndTable3(); - void TestSinksOltpWriteToTopicAndTable4(); - void TestSinksOltpWriteToTopicAndTable5(); - void TestSinksOltpWriteToTopicAndTable6(); + // And then the Logbroker team turned on the feature flag + NotifySchemeShard({.EnablePQConfigTransactionsAtSchemeShard = true}); - void TestSinksOlapWriteToTopicAndTable1(); - void TestSinksOlapWriteToTopicAndTable2(); - void TestSinksOlapWriteToTopicAndTable3(); - void TestSinksOlapWriteToTopicAndTable4(); -}; + // Users continued to write to the topic + WriteToTopic(TEST_TOPIC, TEST_MESSAGE_GROUP_ID, "message-4"); -class TFixtureSinksTable : public TFixtureSinks { -protected: - EClientType GetClientType() const override { - return EClientType::Table; - } -}; + // Users have added new consumers + AddConsumer(TEST_TOPIC, {"consumer-2"}); -class TFixtureSinksQuery : public TFixtureSinks { -protected: - EClientType GetClientType() const override { - return EClientType::Query; - } -}; + // And they wanted to continue reading their messages + Read_Exactly_N_Messages_From_Topic(TEST_TOPIC, "consumer-1", 1); +} -void TFixtureSinks::CreateRowTable(const std::string& path) +Y_UNIT_TEST_F(WriteToTopic_Demo_45_Table, TFixtureTable) { - CreateTable(path); + TestWriteToTopic45(); } -void TFixtureSinks::CreateColumnTable(const std::string& tablePath) +Y_UNIT_TEST_F(WriteToTopic_Demo_45_Query, TFixtureQuery) { - UNIT_ASSERT(!tablePath.empty()); - - std::string path = (tablePath[0] != '/') ? ("/Root/" + tablePath) : tablePath; - - auto createSessionResult = GetTableClient().CreateSession().ExtractValueSync(); - UNIT_ASSERT_C(createSessionResult.IsSuccess(), createSessionResult.GetIssues().ToString()); - auto session = createSessionResult.GetSession(); - - auto desc = NTable::TTableBuilder() - .SetStoreType(NTable::EStoreType::Column) - .AddNonNullableColumn("key", EPrimitiveType::Utf8) - .AddNonNullableColumn("value", EPrimitiveType::Utf8) - .SetPrimaryKeyColumn("key") - .Build(); - auto result = session.CreateTable(path, std::move(desc)).GetValueSync(); - UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + TestWriteToTopic45(); } -bool TFixtureSinks::GetEnableOltpSink() const +Y_UNIT_TEST_F(WriteToTopic_Demo_46_Table, TFixtureTable) { - return true; + TestWriteToTopic46(); } -bool TFixtureSinks::GetEnableOlapSink() const +Y_UNIT_TEST_F(WriteToTopic_Demo_46_Query, TFixtureQuery) { - return true; + TestWriteToTopic46(); } -bool TFixtureSinks::GetEnableHtapTx() const +Y_UNIT_TEST_F(WriteToTopic_Demo_47_Table, TFixtureTable) { - return true; + TestWriteToTopic47(); } -bool TFixtureSinks::GetAllowOlapDataQuery() const +Y_UNIT_TEST_F(WriteToTopic_Demo_47_Query, TFixtureQuery) { - return true; + TestWriteToTopic47(); } Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopic_1_Table, TFixtureSinksTable) @@ -2811,24 +256,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopic_4_Query, TFixtureSinksQuery) TestWriteToTopic9(); } -void TFixtureSinks::TestSinksOltpWriteToTopic5() -{ - CreateTopic("topic_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); - - session->RollbackTx(*tx, EStatus::SUCCESS); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); -} - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopic_5_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopic5(); @@ -2889,44 +316,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_1_Query, TFixtureSinksQuery) TestWriteToTopic24(); } -void TFixtureSinks::TestSinksOltpWriteToTopicAndTable2() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - CreateRowTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 3); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #3"); - } - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - - CheckTabletKeys("topic_A"); - CheckTabletKeys("topic_B"); -} - - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_2_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopicAndTable2(); @@ -2937,48 +326,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_2_Query, TFixtureSinksQuery) TestSinksOltpWriteToTopicAndTable2(); } -void TFixtureSinks::TestSinksOltpWriteToTopicAndTable3() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - - CreateRowTable("/Root/table_A"); - CreateRowTable("/Root/table_B"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - UpsertToTable("table_B", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - const size_t topicMsgCnt = 10; - for (size_t i = 1; i <= topicMsgCnt; ++i) { - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); - } - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), records.size()); - - CheckTabletKeys("topic_A"); - CheckTabletKeys("topic_B"); -} - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_3_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopicAndTable3(); @@ -2989,33 +336,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_3_Query, TFixtureSinksQuery) TestSinksOltpWriteToTopicAndTable3(); } -void TFixtureSinks::TestSinksOltpWriteToTopicAndTable4() -{ - CreateTopic("topic_A"); - CreateRowTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx1 = session->BeginTx(); - auto tx2 = session->BeginTx(); - - session->Execute(R"(SELECT COUNT(*) FROM `table_A`)", tx1.get(), false); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx2.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx1.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - session->CommitTx(*tx2, EStatus::SUCCESS); - session->CommitTx(*tx1, EStatus::ABORTED); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - - CheckTabletKeys("topic_A"); -} - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_4_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopicAndTable4(); @@ -3026,29 +346,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_4_Query, TFixtureSinksQuery) TestSinksOltpWriteToTopicAndTable4(); } -void TFixtureSinks::TestSinksOltpWriteToTopicAndTable5() -{ - CreateTopic("topic_A"); - CreateRowTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - session->RollbackTx(*tx, EStatus::SUCCESS); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); - - CheckTabletKeys("topic_A"); -} - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_5_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopicAndTable5(); @@ -3059,46 +356,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_5_Query, TFixtureSinksQuery) TestSinksOltpWriteToTopicAndTable5(); } -void TFixtureSinks::TestSinksOltpWriteToTopicAndTable6() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - CreateRowTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - InsertToTable("table_A", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); - - DeleteFromTable("table_A", records, *session, tx.get()); - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 3); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #3"); - } - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); - - CheckTabletKeys("topic_A"); - CheckTabletKeys("topic_B"); -} - - Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_6_Table, TFixtureSinksTable) { TestSinksOltpWriteToTopicAndTable6(); @@ -3109,28 +366,6 @@ Y_UNIT_TEST_F(Sinks_Oltp_WriteToTopicAndTable_6_Query, TFixtureSinksQuery) TestSinksOltpWriteToTopicAndTable6(); } -void TFixtureSinks::TestSinksOlapWriteToTopicAndTable1() -{ - CreateTopic("topic_A"); - CreateColumnTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - session->CommitTx(*tx, EStatus::SUCCESS); - - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - - CheckTabletKeys("topic_A"); -} - Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_1_Table, TFixtureSinksTable) { TestSinksOlapWriteToTopicAndTable1(); @@ -3141,49 +376,6 @@ Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_1_Query, TFixtureSinksQuery) TestSinksOlapWriteToTopicAndTable1(); } -void TFixtureSinks::TestSinksOlapWriteToTopicAndTable2() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - - CreateRowTable("/Root/table_A"); - CreateColumnTable("/Root/table_B"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - - UpsertToTable("table_A", records, *session, tx.get()); - UpsertToTable("table_B", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - const size_t topicMsgCnt = 10; - for (size_t i = 1; i <= topicMsgCnt; ++i) { - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); - } - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); - } - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), records.size()); - - CheckTabletKeys("topic_A"); - CheckTabletKeys("topic_B"); -} - Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_2_Table, TFixtureSinksTable) { TestSinksOlapWriteToTopicAndTable2(); @@ -3194,29 +386,6 @@ Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_2_Query, TFixtureSinksQuery) TestSinksOlapWriteToTopicAndTable2(); } -void TFixtureSinks::TestSinksOlapWriteToTopicAndTable3() -{ - CreateTopic("topic_A"); - CreateColumnTable("/Root/table_A"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - UpsertToTable("table_A", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); - - session->RollbackTx(*tx, EStatus::SUCCESS); - - Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); - - CheckTabletKeys("topic_A"); -} - Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_3_Table, TFixtureSinksTable) { TestSinksOlapWriteToTopicAndTable3(); @@ -3227,54 +396,6 @@ Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_3_Query, TFixtureSinksQuery) TestSinksOlapWriteToTopicAndTable3(); } -void TFixtureSinks::TestSinksOlapWriteToTopicAndTable4() -{ - CreateTopic("topic_A"); - CreateTopic("topic_B"); - - CreateRowTable("/Root/table_A"); - CreateColumnTable("/Root/table_B"); - CreateColumnTable("/Root/table_C"); - - auto session = CreateSession(); - auto tx = session->BeginTx(); - - auto records = MakeTableRecords(); - - InsertToTable("table_A", records, *session, tx.get()); - InsertToTable("table_B", records, *session, tx.get()); - UpsertToTable("table_C", records, *session, tx.get()); - - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); - - const size_t topicMsgCnt = 10; - for (size_t i = 1; i <= topicMsgCnt; ++i) { - WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); - } - - DeleteFromTable("table_B", records, *session, tx.get()); - - session->CommitTx(*tx, EStatus::SUCCESS); - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); - } - - { - auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); - UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); - UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); - } - - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), 0); - UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_C"), records.size()); - - CheckTabletKeys("topic_A"); - CheckTabletKeys("topic_B"); -} - Y_UNIT_TEST_F(Sinks_Olap_WriteToTopicAndTable_4_Table, TFixtureSinksTable) { TestSinksOlapWriteToTopicAndTable4(); @@ -3299,25 +420,6 @@ Y_UNIT_TEST_F(The_Transaction_Starts_On_One_Version_And_Ends_On_The_Other, TFixt RestartPQTablet("topic_A", 1); } -void TFixture::TestWriteAndReadMessages(size_t count, size_t size, bool restart) -{ - CreateTopic("topic_A"); - - SetPartitionWriteSpeed("topic_A", 50'000'000); - - for (size_t i = 0; i < count; ++i) { - WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(size, 'x'), nullptr); - } - CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); - - if (restart) { - RestartPQTablet("topic_A", 0); - } - - auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, count); - UNIT_ASSERT_VALUES_EQUAL(messages.size(), count); -} - Y_UNIT_TEST_F(Write_And_Read_Small_Messages_1, TFixtureNoClient) { TestWriteAndReadMessages(320, 64'000, false); @@ -3328,94 +430,6 @@ Y_UNIT_TEST_F(Write_And_Read_Small_Messages_2, TFixtureNoClient) TestWriteAndReadMessages(320, 64'000, true); } -Y_UNIT_TEST_F(Write_And_Read_Big_Messages_1, TFixtureNoClient) -{ - TestWriteAndReadMessages(27, 64'000 * 12, false); -} - -Y_UNIT_TEST_F(Write_And_Read_Big_Messages_2, TFixtureNoClient) -{ - TestWriteAndReadMessages(27, 64'000 * 12, true); -} - -Y_UNIT_TEST_F(Write_And_Read_Huge_Messages_1, TFixtureNoClient) -{ - TestWriteAndReadMessages(4, 9'000'000, false); -} - -Y_UNIT_TEST_F(Write_And_Read_Huge_Messages_2, TFixtureNoClient) -{ - TestWriteAndReadMessages(4, 9'000'000, true); -} - -Y_UNIT_TEST_F(Write_And_Read_Gigant_Messages_1, TFixtureNoClient) -{ - TestWriteAndReadMessages(4, 61'000'000, false); -} - -Y_UNIT_TEST_F(Write_And_Read_Gigant_Messages_2, TFixtureNoClient) -{ - TestWriteAndReadMessages(4, 61'000'000, true); -} - -Y_UNIT_TEST_F(Write_50k_100times_50tx, TFixtureTable) -{ - // 100 transactions. Write 100 50KB messages in each folder. Call the commit at the same time. - // As a result, there will be a lot of small blobs in the FastWrite zone of the main batch, - // which will be picked up by a compact. The scenario is similar to the work of Ya.Metrika. - - const std::size_t PARTITIONS_COUNT = 2; - const std::size_t TXS_COUNT = 50; - - auto makeSourceId = [](unsigned txId, unsigned partitionId) { - std::string sourceId = TEST_MESSAGE_GROUP_ID; - sourceId += "_"; - sourceId += ToString(txId); - sourceId += "_"; - sourceId += ToString(partitionId); - return sourceId; - }; - - CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); - - SetPartitionWriteSpeed("topic_A", 50'000'000); - - std::vector> sessions; - std::vector> transactions; - - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - sessions.push_back(CreateSession()); - auto& session = sessions.back(); - - transactions.push_back(session->BeginTx()); - auto& tx = transactions.back(); - - auto sourceId = makeSourceId(i, 0); - for (size_t j = 0; j < 100; ++j) { - WriteToTopic("topic_A", sourceId, std::string(50'000, 'x'), tx.get(), 0); - } - WaitForAcks("topic_A", sourceId); - - sourceId = makeSourceId(i, 1); - WriteToTopic("topic_A", sourceId, std::string(50'000, 'x'), tx.get(), 1); - WaitForAcks("topic_A", sourceId); - } - - // We are doing an asynchronous commit of transactions. They will be executed simultaneously. - std::vector futures; - - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); - } - - // All transactions must be completed successfully. - for (std::size_t i = 0; i < TXS_COUNT; ++i) { - futures[i].Wait(); - const auto& result = futures[i].GetValueSync(); - UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - } -} - } } diff --git a/src/client/topic/ut/ut_utils/txusage_fixture.cpp b/src/client/topic/ut/ut_utils/txusage_fixture.cpp new file mode 100644 index 00000000000..994e7969d2e --- /dev/null +++ b/src/client/topic/ut/ut_utils/txusage_fixture.cpp @@ -0,0 +1,2431 @@ +#include "txusage_fixture.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace std::chrono_literals; + +namespace NYdb::inline V3::NTopic::NTests::NTxUsage { + +namespace NPQ = NKikimr::NPQ; +namespace NConsole = NKikimr::NConsole; +namespace NSchemeCache = NKikimr::NSchemeCache; +namespace TEvKeyValue = NKikimr::TEvKeyValue; +namespace NMsgBusProxy = NKikimr::NMsgBusProxy; + +const auto TEST_MESSAGE_GROUP_ID_1 = TEST_MESSAGE_GROUP_ID + "_1"; +const auto TEST_MESSAGE_GROUP_ID_2 = TEST_MESSAGE_GROUP_ID + "_2"; +const auto TEST_MESSAGE_GROUP_ID_3 = TEST_MESSAGE_GROUP_ID + "_3"; +const auto TEST_MESSAGE_GROUP_ID_4 = TEST_MESSAGE_GROUP_ID + "_4"; + +TFixture::TTableRecord::TTableRecord(const std::string& key, const std::string& value) : + Key(key), + Value(value) +{ +} + +void TFixture::SetUp(NUnitTest::TTestContext&) +{ + NKikimr::Tests::TServerSettings settings = TTopicSdkTestSetup::MakeServerSettings(); + settings.SetEnableTopicServiceTx(true); + settings.SetEnableTopicSplitMerge(true); + settings.SetEnablePQConfigTransactionsAtSchemeShard(true); + settings.SetEnableOltpSink(GetEnableOltpSink()); + settings.SetEnableOlapSink(GetEnableOlapSink()); + settings.SetEnableHtapTx(GetEnableHtapTx()); + settings.SetAllowOlapDataQuery(GetAllowOlapDataQuery()); + + Setup = std::make_unique(TEST_CASE_NAME, settings); + + Driver = std::make_unique(Setup->MakeDriver()); + auto tableSettings = NTable::TClientSettings().SessionPoolSettings(NTable::TSessionPoolSettings() + .MaxActiveSessions(3000) + ); + + auto querySettings = NQuery::TClientSettings().SessionPoolSettings(NQuery::TSessionPoolSettings() + .MaxActiveSessions(3000) + ); + + TableClient = std::make_unique(*Driver, tableSettings); + QueryClient = std::make_unique(*Driver, querySettings); +} + +void TFixture::NotifySchemeShard(const TFeatureFlags& flags) +{ + auto request = std::make_unique(); + *request->Record.MutableConfig() = *Setup->GetServer().ServerSettings.AppConfig; + request->Record.MutableConfig()->MutableFeatureFlags()->SetEnablePQConfigTransactionsAtSchemeShard(flags.EnablePQConfigTransactionsAtSchemeShard); + + auto& runtime = Setup->GetRuntime(); + auto actorId = runtime.AllocateEdgeActor(); + + std::uint64_t ssId = GetSchemeShardTabletId(actorId); + + runtime.SendToPipe(ssId, actorId, request.release()); + runtime.GrabEdgeEvent(); +} + +TFixture::TTableSession::TTableSession(NTable::TTableClient& client) + : Session_(Init(client)) +{ +} + +NTable::TSession TFixture::TTableSession::Init(NTable::TTableClient& client) +{ + auto result = client.GetSession().ExtractValueSync(); + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return result.GetSession(); +} + +std::vector TFixture::TTableSession::Execute(const std::string& query, + TTransactionBase* tx, + bool commit, + const TParams& params) +{ + while (true) { + auto txTable = dynamic_cast(tx); + auto txControl = NTable::TTxControl::Tx(*txTable).CommitTx(commit); + + auto result = Session_.ExecuteDataQuery(query, txControl, params).GetValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return std::move(result).ExtractResultSets(); + } + std::this_thread::sleep_for(100ms); + } +} + +TFixture::ISession::TExecuteInTxResult TFixture::TTableSession::ExecuteInTx(const std::string& query, + bool commit, + const TParams& params) +{ + while (true) { + auto txControl = NTable::TTxControl::BeginTx().CommitTx(commit); + + auto result = Session_.ExecuteDataQuery(query, txControl, params).GetValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return {std::move(result).ExtractResultSets(), std::make_unique(*result.GetTransaction())}; + } + std::this_thread::sleep_for(100ms); + } +} + +std::unique_ptr TFixture::TTableSession::BeginTx() +{ + while (true) { + auto result = Session_.BeginTransaction().ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return std::make_unique(result.GetTransaction()); + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TTableSession::CommitTx(TTransactionBase& tx, EStatus status) +{ + auto txTable = dynamic_cast(tx); + while (true) { + auto result = txTable.Commit().ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); + return; + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TTableSession::RollbackTx(TTransactionBase& tx, EStatus status) +{ + auto txTable = dynamic_cast(tx); + while (true) { + auto result = txTable.Rollback().ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); + return; + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TTableSession::Close() +{ + Session_.Close(); +} + +TAsyncStatus TFixture::TTableSession::AsyncCommitTx(TTransactionBase& tx) +{ + auto txTable = dynamic_cast(tx); + return txTable.Commit().Apply([](auto result) { + return TStatus(result.GetValue()); + }); +} + +TFixture::TQuerySession::TQuerySession(NQuery::TQueryClient& client, + const std::string& endpoint, + const std::string& database) + : Session_(Init(client)) + , Endpoint_(endpoint) + , Database_(database) +{ +} + +NQuery::TSession TFixture::TQuerySession::Init(NQuery::TQueryClient& client) +{ + auto result = client.GetSession().ExtractValueSync(); + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return result.GetSession(); +} + +std::vector TFixture::TQuerySession::Execute(const std::string& query, + TTransactionBase* tx, + bool commit, + const TParams& params) +{ + while (true) { + auto txQuery = dynamic_cast(tx); + auto txControl = NQuery::TTxControl::Tx(*txQuery).CommitTx(commit); + + auto result = Session_.ExecuteQuery(query, txControl, params).ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return result.GetResultSets(); + } + std::this_thread::sleep_for(100ms); + } +} + +TFixture::ISession::TExecuteInTxResult TFixture::TQuerySession::ExecuteInTx(const std::string& query, + bool commit, + const TParams& params) +{ + while (true) { + auto txControl = NQuery::TTxControl::BeginTx().CommitTx(commit); + + auto result = Session_.ExecuteQuery(query, txControl, params).ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return {result.GetResultSets(), std::make_unique(*result.GetTransaction())}; + } + std::this_thread::sleep_for(100ms); + } +} + +std::unique_ptr TFixture::TQuerySession::BeginTx() +{ + while (true) { + auto result = Session_.BeginTransaction(NQuery::TTxSettings()).ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); + return std::make_unique(result.GetTransaction()); + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TQuerySession::CommitTx(TTransactionBase& tx, EStatus status) +{ + auto txQuery = dynamic_cast(tx); + while (true) { + auto result = txQuery.Commit().ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); + return; + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TQuerySession::RollbackTx(TTransactionBase& tx, EStatus status) +{ + auto txQuery = dynamic_cast(tx); + while (true) { + auto result = txQuery.Rollback().ExtractValueSync(); + if (result.GetStatus() != EStatus::SESSION_BUSY) { + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), status, result.GetIssues().ToString()); + return; + } + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::TQuerySession::Close() +{ + // SDK doesn't provide a method to close the session for Query Client, so we use grpc API directly + auto credentials = grpc::InsecureChannelCredentials(); + auto channel = grpc::CreateChannel(TStringType(Endpoint_), credentials); + auto stub = Ydb::Query::V1::QueryService::NewStub(channel); + + grpc::ClientContext context; + context.AddMetadata("x-ydb-database", TStringType(Database_)); + + Ydb::Query::DeleteSessionRequest request; + request.set_session_id(Session_.GetId()); + + Ydb::Query::DeleteSessionResponse response; + auto status = stub->DeleteSession(&context, request, &response); + + NIssue::TIssues issues; + NYdb::NIssue::IssuesFromMessage(response.issues(), issues); + UNIT_ASSERT_C(status.ok(), status.error_message()); + UNIT_ASSERT_VALUES_EQUAL_C(response.status(), Ydb::StatusIds::SUCCESS, issues.ToString()); +} + +TAsyncStatus TFixture::TQuerySession::AsyncCommitTx(TTransactionBase& tx) +{ + auto txQuery = dynamic_cast(tx); + return txQuery.Commit().Apply([](auto result) { + return TStatus(result.GetValue()); + }); +} + +std::unique_ptr TFixture::CreateSession() +{ + switch (GetClientType()) { + case EClientType::Table: { + UNIT_ASSERT_C(TableClient, "TableClient is not initialized"); + return std::make_unique(*TableClient); + } + case EClientType::Query: { + UNIT_ASSERT_C(QueryClient, "QueryClient is not initialized"); + return std::make_unique(*QueryClient, + Setup->GetEndpoint(), + Setup->GetDatabase()); + } + case EClientType::None: { + UNIT_FAIL("CreateSession is forbidden for None client type"); + } + } + + return nullptr; +} + +template +E TFixture::ReadEvent(TTopicReadSessionPtr reader, TTransactionBase& tx) +{ + NTopic::TReadSessionGetEventSettings options; + options.Block(true); + options.MaxEventsCount(1); + options.Tx(tx); + + auto event = reader->GetEvent(options); + UNIT_ASSERT(event); + + auto ev = std::get_if(&*event); + UNIT_ASSERT(ev); + + return *ev; +} + +template +E TFixture::ReadEvent(TTopicReadSessionPtr reader) +{ + auto event = reader->GetEvent(true, 1); + UNIT_ASSERT(event); + + auto ev = std::get_if(&*event); + UNIT_ASSERT(ev); + + return *ev; +} + +void TFixture::CreateTopic(const std::string& path, + const std::string& consumer, + std::size_t partitionCount, + std::optional maxPartitionCount, + const TDuration retention, + bool important) +{ + Setup->CreateTopic(path, consumer, partitionCount, maxPartitionCount, retention, important); +} + +void TFixture::AddConsumer(const std::string& topicPath, + const std::vector& consumers) +{ + NTopic::TTopicClient client(GetDriver()); + NTopic::TAlterTopicSettings settings; + + for (const auto& consumer : consumers) { + settings.BeginAddConsumer(consumer); + } + + auto result = client.AlterTopic(topicPath, settings).GetValueSync(); + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); +} + +void TFixture::SetPartitionWriteSpeed(const std::string& topicName, std::size_t bytesPerSeconds) +{ + NTopic::TTopicClient client(GetDriver()); + NTopic::TAlterTopicSettings settings; + + settings.SetPartitionWriteSpeedBytesPerSecond(bytesPerSeconds); + + auto result = client.AlterTopic(Setup->GetTopicPath(topicName), settings).GetValueSync(); + Y_ENSURE_BT(result.IsSuccess(), ToString(static_cast(result))); +} + +const TDriver& TFixture::GetDriver() const +{ + return *Driver; +} + +NTable::TTableClient& TFixture::GetTableClient() +{ + return *TableClient; +} + +auto TFixture::CreateTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + std::optional partitionId) -> TTopicWriteSessionPtr +{ + NTopic::TTopicClient client(GetDriver()); + NTopic::TWriteSessionSettings options; + options.Path(topicPath); + options.ProducerId(messageGroupId); + options.MessageGroupId(messageGroupId); + options.PartitionId(partitionId); + options.Codec(ECodec::RAW); + return client.CreateWriteSession(options); +} + +auto TFixture::GetTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + std::optional partitionId) -> TTopicWriteSessionContext& +{ + std::pair key(topicPath, messageGroupId); + auto i = TopicWriteSessions.find(key); + + if (i == TopicWriteSessions.end()) { + TTopicWriteSessionContext context; + context.Session = CreateTopicWriteSession(topicPath, messageGroupId, partitionId); + + TopicWriteSessions.emplace(key, std::move(context)); + + i = TopicWriteSessions.find(key); + } + + return i->second; +} + +NTopic::TTopicReadSettings MakeTopicReadSettings(const std::string& topicPath, + std::optional partitionId) +{ + TTopicReadSettings options; + options.Path(topicPath); + if (partitionId) { + options.AppendPartitionIds(*partitionId); + } + return options; +} + +NTopic::TReadSessionSettings MakeTopicReadSessionSettings(const std::string& topicPath, + const std::string& consumerName, + std::optional partitionId) +{ + NTopic::TReadSessionSettings options; + options.AppendTopics(MakeTopicReadSettings(topicPath, partitionId)); + options.ConsumerName(consumerName); + return options; +} + +auto TFixture::CreateTopicReadSession(const std::string& topicPath, + const std::string& consumerName, + std::optional partitionId) -> TTopicReadSessionPtr +{ + NTopic::TTopicClient client(GetDriver()); + return client.CreateReadSession(MakeTopicReadSessionSettings(topicPath, + consumerName, + partitionId)); +} + +auto TFixture::GetTopicReadSession(const std::string& topicPath, + const std::string& consumerName, + std::optional partitionId) -> TTopicReadSessionPtr +{ + TTopicReadSessionPtr session; + + if (auto i = TopicReadSessions.find(topicPath); i == TopicReadSessions.end()) { + session = CreateTopicReadSession(topicPath, consumerName, partitionId); + auto event = ReadEvent(session); + event.Confirm(); + TopicReadSessions.emplace(topicPath, session); + } else { + session = i->second; + } + + return session; +} + +void TFixture::TTopicWriteSessionContext::WaitForContinuationToken() +{ + while (!ContinuationToken.has_value()) { + WaitForEvent(); + } +} + +void TFixture::TTopicWriteSessionContext::WaitForEvent() +{ + Session->WaitEvent().Wait(); + for (auto& event : Session->GetEvents()) { + if (auto* e = std::get_if(&event)) { + ContinuationToken = std::move(e->ContinuationToken); + } else if (auto* e = std::get_if(&event)) { + for (auto& ack : e->Acks) { + switch (ack.State) { + case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN: + ++WrittenAckCount; + break; + case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN_IN_TX: + ++WrittenInTxAckCount; + break; + default: + break; + } + } + } else if ([[maybe_unused]] auto* e = std::get_if(&event)) { + UNIT_FAIL(""); + } + } +} + +void TFixture::TTopicWriteSessionContext::Write(const std::string& message, TTransactionBase* tx) +{ + NTopic::TWriteMessage params(message); + + if (tx) { + params.Tx(*tx); + } + + Session->Write(std::move(*ContinuationToken), + std::move(params)); + + ++WriteCount; + ContinuationToken = std::nullopt; +} + +void TFixture::CloseTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + bool force) +{ + std::pair key(topicPath, messageGroupId); + auto i = TopicWriteSessions.find(key); + + UNIT_ASSERT(i != TopicWriteSessions.end()); + + TTopicWriteSessionContext& context = i->second; + + context.Session->Close(force ? TDuration::MilliSeconds(0) : TDuration::Max()); + TopicWriteSessions.erase(key); +} + +void TFixture::CloseTopicReadSession(const std::string& topicPath, + const std::string& consumerName) +{ + Y_UNUSED(consumerName); + TopicReadSessions.erase(topicPath); +} + +void TFixture::WriteToTopic(const std::string& topicPath, + const std::string& messageGroupId, + const std::string& message, + TTransactionBase* tx, + std::optional partitionId) +{ + TTopicWriteSessionContext& context = GetTopicWriteSession(topicPath, messageGroupId, partitionId); + context.WaitForContinuationToken(); + UNIT_ASSERT(context.ContinuationToken.has_value()); + context.Write(message, tx); +} + +std::vector TFixture::ReadFromTopic(const std::string& topicPath, + const std::string& consumerName, + const TDuration& duration, + TTransactionBase* tx, + std::optional partitionId) +{ + std::vector messages; + + TInstant end = TInstant::Now() + duration; + TDuration remain = duration; + + auto session = GetTopicReadSession(topicPath, consumerName, partitionId); + + while (TInstant::Now() < end) { + if (!session->WaitEvent().Wait(remain)) { + return messages; + } + + NTopic::TReadSessionGetEventSettings settings; + if (tx) { + settings.Tx(*tx); + } + + for (auto& event : session->GetEvents(settings)) { + if (auto* e = std::get_if(&event)) { + std::cerr << e->HasCompressedMessages() << " " << e->GetMessagesCount() << std::endl; + for (auto& m : e->GetMessages()) { + messages.emplace_back(m.GetData()); + } + + if (!tx) { + e->Commit(); + } + } + } + + remain = end - TInstant::Now(); + } + + return messages; +} + +void TFixture::WaitForAcks(const std::string& topicPath, const std::string& messageGroupId, std::size_t writtenInTxCount) +{ + std::pair key(topicPath, messageGroupId); + auto i = TopicWriteSessions.find(key); + UNIT_ASSERT(i != TopicWriteSessions.end()); + + auto& context = i->second; + + UNIT_ASSERT(context.AckCount() <= context.WriteCount); + + while (context.AckCount() < context.WriteCount) { + context.WaitForEvent(); + } + + UNIT_ASSERT((context.WrittenAckCount + context.WrittenInTxAckCount) == context.WriteCount); + + if (writtenInTxCount != std::numeric_limits::max()) { + UNIT_ASSERT_VALUES_EQUAL(context.WrittenInTxAckCount, writtenInTxCount); + } +} + +void TFixture::WaitForSessionClose(const std::string& topicPath, + const std::string& messageGroupId, + NYdb::EStatus status) +{ + std::pair key(topicPath, messageGroupId); + auto i = TopicWriteSessions.find(key); + UNIT_ASSERT(i != TopicWriteSessions.end()); + + auto& context = i->second; + + UNIT_ASSERT(context.AckCount() <= context.WriteCount); + + for(bool stop = false; !stop; ) { + context.Session->WaitEvent().Wait(); + for (auto& event : context.Session->GetEvents()) { + if (auto* e = std::get_if(&event)) { + context.ContinuationToken = std::move(e->ContinuationToken); + } else if (auto* e = std::get_if(&event)) { + for (auto& ack : e->Acks) { + switch (ack.State) { + case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN: + ++context.WrittenAckCount; + break; + case NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN_IN_TX: + ++context.WrittenInTxAckCount; + break; + default: + break; + } + } + } else if (auto* e = std::get_if(&event)) { + UNIT_ASSERT_VALUES_EQUAL(e->GetStatus(), status); + UNIT_ASSERT_GT(e->GetIssues().Size(), 0); + stop = true; + } + } + } + + UNIT_ASSERT(context.AckCount() <= context.WriteCount); +} + +std::uint64_t TFixture::GetSchemeShardTabletId(const NActors::TActorId& actorId) +{ + auto navigate = std::make_unique(); + navigate->DatabaseName = "/Root"; + + NSchemeCache::TSchemeCacheNavigate::TEntry entry; + entry.Path = NKikimr::SplitPath("/Root"); + entry.SyncVersion = true; + entry.ShowPrivatePath = true; + entry.Operation = NSchemeCache::TSchemeCacheNavigate::OpList; + + navigate->ResultSet.push_back(std::move(entry)); + //navigate->UserToken = "root@builtin"; + navigate->Cookie = 12345; + + auto& runtime = Setup->GetRuntime(); + + runtime.Send(NKikimr::MakeSchemeCacheID(), actorId, + new NKikimr::TEvTxProxySchemeCache::TEvNavigateKeySet(navigate.release()), + 0, + true); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT_VALUES_EQUAL(response->Request->Cookie, 12345); + UNIT_ASSERT_VALUES_EQUAL(response->Request->ErrorCount, 0); + + auto& front = response->Request->ResultSet.front(); + + return front.Self->Info.GetSchemeshardId(); +} + +std::uint64_t TFixture::GetTopicTabletId(const NActors::TActorId& actorId, const std::string& topicPath, std::uint32_t partition) +{ + auto navigate = std::make_unique(); + navigate->DatabaseName = "/Root"; + + NSchemeCache::TSchemeCacheNavigate::TEntry entry; + entry.Path = NKikimr::SplitPath(TString{topicPath}); + entry.SyncVersion = true; + entry.ShowPrivatePath = true; + entry.Operation = NSchemeCache::TSchemeCacheNavigate::OpList; + + navigate->ResultSet.push_back(std::move(entry)); + //navigate->UserToken = "root@builtin"; + navigate->Cookie = 12345; + + auto& runtime = Setup->GetRuntime(); + + runtime.Send(NKikimr::MakeSchemeCacheID(), actorId, + new NKikimr::TEvTxProxySchemeCache::TEvNavigateKeySet(navigate.release()), + 0, + true); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT_VALUES_EQUAL(response->Request->Cookie, 12345); + UNIT_ASSERT_VALUES_EQUAL(response->Request->ErrorCount, 0); + + auto& front = response->Request->ResultSet.front(); + UNIT_ASSERT(front.PQGroupInfo); + UNIT_ASSERT_GT(front.PQGroupInfo->Description.PartitionsSize(), 0); + UNIT_ASSERT_LT(partition, front.PQGroupInfo->Description.PartitionsSize()); + + for (size_t i = 0; i < front.PQGroupInfo->Description.PartitionsSize(); ++i) { + auto& p = front.PQGroupInfo->Description.GetPartitions(partition); + if (p.GetPartitionId() == partition) { + return p.GetTabletId(); + } + } + + UNIT_FAIL("unknown partition"); + + return std::numeric_limits::max(); +} + +std::vector TFixture::GetTabletKeys(const NActors::TActorId& actorId, + std::uint64_t tabletId) +{ + auto request = std::make_unique(); + request->Record.SetCookie(12345); + + auto cmd = request->Record.AddCmdReadRange(); + TString from(1, '\x00'); + TString to(1, '\xFF'); + auto range = cmd->MutableRange(); + range->SetFrom(from); + range->SetIncludeFrom(true); + range->SetTo(to); + range->SetIncludeTo(true); + + auto& runtime = Setup->GetRuntime(); + + runtime.SendToPipe(tabletId, actorId, request.release()); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT(response->Record.HasCookie()); + UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); + UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadRangeResultSize(), 1); + + std::vector keys; + + auto& result = response->Record.GetReadRangeResult(0); + for (size_t i = 0; i < result.PairSize(); ++i) { + auto& kv = result.GetPair(i); + keys.emplace_back(kv.GetKey()); + } + + return keys; +} + +size_t TFixture::GetPQCacheRenameKeysCount() +{ + using namespace NKikimr::NPQ; + + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + + auto request = MakeHolder(); + + runtime.Send(MakePersQueueL2CacheID(), edge, request.Release()); + + TAutoPtr handle; + auto* result = runtime.GrabEdgeEvent(handle); + + return result->RenamedKeys; +} + +std::vector TFixture::Read_Exactly_N_Messages_From_Topic(const std::string& topicPath, + const std::string& consumerName, + std::size_t limit) +{ + std::vector result; + + while (result.size() < limit) { + auto messages = ReadFromTopic(topicPath, consumerName, TDuration::Seconds(2)); + for (auto& m : messages) { + result.push_back(std::move(m)); + } + } + + UNIT_ASSERT_VALUES_EQUAL(result.size(), limit); + + return result; +} + +void TFixture::TestWriteToTopic1() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #4", tx.get()); + + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #5", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #6", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #7", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #8", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #9", tx.get()); + + { + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); + } + + { + auto messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 4); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages[3], "message #4"); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 5); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #5"); + UNIT_ASSERT_VALUES_EQUAL(messages[4], "message #9"); + } +} + +void TFixture::TestWriteToTopic4() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + + auto session = CreateSession(); + auto tx_1 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx_1.get()); + + auto tx_2 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", tx_2.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #4", tx_2.get()); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); + + messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); + + session->CommitTx(*tx_2, EStatus::SUCCESS); + session->CommitTx(*tx_1, EStatus::ABORTED); + + messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #3"); + + messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #4"); +} + +void TFixture::TestWriteToTopic7() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #1", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #2", tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, "message #3"); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, "message #4"); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #5", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, "message #6", tx.get()); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #3"); + UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #4"); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 4); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages[3], "message #6"); + } +} + +void TFixture::TestWriteToTopic9() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx_1 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); + + auto tx_2 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx_2.get()); + + { + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 0); + } + + session->CommitTx(*tx_2, EStatus::SUCCESS); + session->CommitTx(*tx_1, EStatus::ABORTED); + + { + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #2"); + } +} + +void TFixture::TestWriteToTopic10() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + + { + auto tx_1 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx_1.get()); + + session->CommitTx(*tx_1, EStatus::SUCCESS); + } + + { + auto tx_2 = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx_2.get()); + + session->CommitTx(*tx_2, EStatus::SUCCESS); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #2"); + } +} + +void TFixture::TestWriteToTopic11() +{ + for (auto endOfTransaction : {Commit, Rollback, CloseTableSession}) { + TestTheCompletionOfATransaction({.Topics={"topic_A"}, .EndOfTransaction = endOfTransaction}); + TestTheCompletionOfATransaction({.Topics={"topic_A", "topic_B"}, .EndOfTransaction = endOfTransaction}); + } +} + +void TFixture::TestWriteToTopic24() +{ + // + // the test verifies a transaction in which data is written to a topic and to a table + // + CreateTopic("topic_A"); + CreateTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + session->CommitTx(*tx, EStatus::SUCCESS); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(messages[0], MakeJsonDoc(records)); + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + + CheckTabletKeys("topic_A"); +} + +void TFixture::TestWriteToTopic26() +{ + // + // the test verifies a transaction in which data is read from a partition of one topic and written to + // another partition of this topic + // + const std::uint32_t PARTITION_0 = 0; + const std::uint32_t PARTITION_1 = 1; + + CreateTopic("topic_A", TEST_CONSUMER, 2); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", nullptr, PARTITION_0); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", nullptr, PARTITION_0); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #3", nullptr, PARTITION_0); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), PARTITION_0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 3); + + for (const auto& m : messages) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, m, tx.get(), PARTITION_1); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), nullptr, PARTITION_1); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 3); +} + +void TFixture::TestWriteToTopic27() +{ + CreateTopic("topic_A", TEST_CONSUMER); + CreateTopic("topic_B", TEST_CONSUMER); + CreateTopic("topic_C", TEST_CONSUMER); + + for (size_t i = 0; i < 2; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", nullptr, 0); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", nullptr, 0); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + + WriteToTopic("topic_C", TEST_MESSAGE_GROUP_ID, messages[0], tx.get(), 0); + + messages = ReadFromTopic("topic_B", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 1); + + WriteToTopic("topic_C", TEST_MESSAGE_GROUP_ID, messages[0], tx.get(), 0); + + session->CommitTx(*tx, EStatus::SUCCESS); + + messages = ReadFromTopic("topic_C", TEST_CONSUMER, TDuration::Seconds(2), nullptr, 0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); + + DumpPQTabletKeys("topic_A"); + DumpPQTabletKeys("topic_B"); + DumpPQTabletKeys("topic_C"); + } +} + +auto TFixture::GetAvgWriteBytes(const std::string& topicName, + std::uint32_t partitionId) -> TAvgWriteBytes +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partitionId); + + runtime.SendToPipe(tabletId, edge, new NKikimr::TEvPersQueue::TEvStatus()); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT_VALUES_EQUAL(tabletId, response->Record.GetTabletId()); + + TAvgWriteBytes result; + + for (std::size_t i = 0; i < response->Record.PartResultSize(); ++i) { + const auto& partition = response->Record.GetPartResult(i); + if (partition.GetPartition() == static_cast(partitionId)) { + result.PerSec = partition.GetAvgWriteSpeedPerSec(); + result.PerMin = partition.GetAvgWriteSpeedPerMin(); + result.PerHour = partition.GetAvgWriteSpeedPerHour(); + result.PerDay = partition.GetAvgWriteSpeedPerDay(); + break; + } + } + + return result; +} + +bool TFixture::GetEnableOltpSink() const +{ + return false; +} + +bool TFixture::GetEnableOlapSink() const +{ + return false; +} + +bool TFixture::GetEnableHtapTx() const +{ + return false; +} + +bool TFixture::GetAllowOlapDataQuery() const +{ + return false; +} + +NPQ::TWriteId TFixture::GetTransactionWriteId(const NActors::TActorId& actorId, + std::uint64_t tabletId) +{ + auto request = std::make_unique(); + request->Record.SetCookie(12345); + request->Record.AddCmdRead()->SetKey("_txinfo"); + + auto& runtime = Setup->GetRuntime(); + + runtime.SendToPipe(tabletId, actorId, request.release()); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT(response->Record.HasCookie()); + UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); + UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadResultSize(), 1); + + auto& read = response->Record.GetReadResult(0); + + NKikimrPQ::TTabletTxInfo info; + UNIT_ASSERT(info.ParseFromString(read.GetValue())); + + UNIT_ASSERT_VALUES_EQUAL(info.TxWritesSize(), 1); + + auto& writeInfo = info.GetTxWrites(0); + UNIT_ASSERT(writeInfo.HasWriteId()); + + return NPQ::GetWriteId(writeInfo); +} + +void TFixture::SendLongTxLockStatus(const NActors::TActorId& actorId, + std::uint64_t tabletId, + const NPQ::TWriteId& writeId, + NKikimrLongTxService::TEvLockStatus::EStatus status) +{ + auto event = + std::make_unique(writeId.KeyId, writeId.NodeId, + status); + auto& runtime = Setup->GetRuntime(); + runtime.SendToPipe(tabletId, actorId, event.release()); +} + +void TFixture::WaitForTheTabletToDeleteTheWriteInfo(const NActors::TActorId& actorId, + std::uint64_t tabletId, + const NPQ::TWriteId& writeId) +{ + while (true) { + auto request = std::make_unique(); + request->Record.SetCookie(12345); + request->Record.AddCmdRead()->SetKey("_txinfo"); + + auto& runtime = Setup->GetRuntime(); + + runtime.SendToPipe(tabletId, actorId, request.release()); + auto response = runtime.GrabEdgeEvent(); + + UNIT_ASSERT(response->Record.HasCookie()); + UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), 12345); + UNIT_ASSERT_VALUES_EQUAL(response->Record.ReadResultSize(), 1); + + auto& read = response->Record.GetReadResult(0); + + NKikimrPQ::TTabletTxInfo info; + UNIT_ASSERT(info.ParseFromString(read.GetValue())); + + bool found = false; + + for (size_t i = 0; i < info.TxWritesSize(); ++i) { + auto& writeInfo = info.GetTxWrites(i); + UNIT_ASSERT(writeInfo.HasWriteId()); + if ((NPQ::GetWriteId(writeInfo) == writeId) && writeInfo.HasOriginalPartitionId()) { + found = true; + break; + } + } + + if (!found) { + break; + } + + std::this_thread::sleep_for(100ms); + } +} + +void TFixture::RestartPQTablet(const std::string& topicName, std::uint32_t partition) +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partition); + runtime.SendToPipe(tabletId, edge, new NActors::TEvents::TEvPoison()); + + std::this_thread::sleep_for(2s); +} + +void TFixture::DeleteSupportivePartition(const std::string& topicName, std::uint32_t partition) +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partition); + NPQ::TWriteId writeId = GetTransactionWriteId(edge, tabletId); + + SendLongTxLockStatus(edge, tabletId, writeId, NKikimrLongTxService::TEvLockStatus::STATUS_NOT_FOUND); + + WaitForTheTabletToDeleteTheWriteInfo(edge, tabletId, writeId); +} + +void TFixture::CheckTabletKeys(const std::string& topicName) +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, 0); + + const std::unordered_set types { + NPQ::TKeyPrefix::TypeInfo, + NPQ::TKeyPrefix::TypeData, + NPQ::TKeyPrefix::TypeTmpData, + NPQ::TKeyPrefix::TypeMeta, + NPQ::TKeyPrefix::TypeTxMeta, + }; + + bool found; + std::vector keys; + for (size_t i = 0; i < 20; ++i) { + keys = GetTabletKeys(edge, tabletId); + + found = false; + for (const auto& key : keys) { + UNIT_ASSERT_GT(key.size(), 0); + if (key[0] == '_') { + continue; + } + + if (types.contains(key[0])) { + found = false; + break; + } + } + + if (!found) { + break; + } + + std::this_thread::sleep_for(100ms); + } + + if (found) { + std::cerr << "keys for tablet " << tabletId << ":" << std::endl; + for (const auto& k : keys) { + std::cerr << k << std::endl; + } + std::cerr << "=============" << std::endl; + + UNIT_FAIL("unexpected keys for tablet " << tabletId); + } +} + +void TFixture::DumpPQTabletKeys(const std::string& topicName) +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicName, 0); + auto keys = GetTabletKeys(edge, tabletId); + for (const auto& key : keys) { + std::cerr << key << std::endl; + } +} + +void TFixture::PQTabletPrepareFromResource(const std::string& topicPath, + std::uint32_t partitionId, + const std::string& resourceName) +{ + auto& runtime = Setup->GetRuntime(); + NActors::TActorId edge = runtime.AllocateEdgeActor(); + std::uint64_t tabletId = GetTopicTabletId(edge, "/Root/" + topicPath, partitionId); + + auto request = MakeHolder(); + std::size_t count = 0; + + for (TStringStream stream(NResource::Find(resourceName)); true; ++count) { + TString key, encoded; + + if (!stream.ReadTo(key, ' ')) { + break; + } + encoded = stream.ReadLine(); + + auto decoded = Base64Decode(encoded); + TStringInput decodedStream(decoded); + TBZipDecompress decompressor(&decodedStream); + + auto* cmd = request->Record.AddCmdWrite(); + cmd->SetKey(key); + cmd->SetValue(decompressor.ReadAll()); + } + + runtime.SendToPipe(tabletId, edge, request.Release(), 0, NKikimr::GetPipeConfigWithRetries()); + + TAutoPtr handle; + auto* response = runtime.GrabEdgeEvent(handle); + UNIT_ASSERT(response); + UNIT_ASSERT(response->Record.HasStatus()); + UNIT_ASSERT_EQUAL(response->Record.GetStatus(), NMsgBusProxy::MSTATUS_OK); + + UNIT_ASSERT_VALUES_EQUAL(response->Record.WriteResultSize(), count); + + for (std::size_t i = 0; i < response->Record.WriteResultSize(); ++i) { + const auto &result = response->Record.GetWriteResult(i); + UNIT_ASSERT(result.HasStatus()); + UNIT_ASSERT_EQUAL(result.GetStatus(), NKikimrProto::OK); + } +} + +void TFixture::TestTheCompletionOfATransaction(const TTransactionCompletionTestDescription& d) +{ + for (auto& topic : d.Topics) { + CreateTopic(topic); + } + + { + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (auto& topic : d.Topics) { + WriteToTopic(topic, TEST_MESSAGE_GROUP_ID, "message", tx.get()); + // TODO: нужен callback для RollbakTx + WaitForAcks(topic, TEST_MESSAGE_GROUP_ID); + } + + switch (d.EndOfTransaction) { + case Commit: + session->CommitTx(*tx, EStatus::SUCCESS); + break; + case Rollback: + session->RollbackTx(*tx, EStatus::SUCCESS); + break; + case CloseTableSession: + break; + } + } + + std::this_thread::sleep_for(std::chrono::seconds(5)); + + for (auto& topic : d.Topics) { + CheckTabletKeys(topic); + } + + for (auto& topic : d.Topics) { + CloseTopicWriteSession(topic, TEST_MESSAGE_GROUP_ID); + } +} + +void TFixture::TestWriteToTopic12() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + DeleteSupportivePartition("topic_A", 0); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + WaitForSessionClose("topic_A", TEST_MESSAGE_GROUP_ID, NYdb::EStatus::PRECONDITION_FAILED); +} + +void TFixture::TestWriteToTopic13() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message", tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + DeleteSupportivePartition("topic_A", 0); + + session->CommitTx(*tx, EStatus::ABORTED); +} + +void TFixture::TestWriteToTopic14() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + DeleteSupportivePartition("topic_A", 0); + + CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + + session->CommitTx(*tx, EStatus::ABORTED); +} + +void TFixture::TestWriteToTopic16() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + + RestartPQTablet("topic_A", 0); + + session->CommitTx(*tx, EStatus::SUCCESS); + + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 2); + UNIT_ASSERT_VALUES_EQUAL(messages[0], "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages[1], "message #2"); +} + +void TFixture::TestTxWithBigBlobs(const TTestTxWithBigBlobsParams& params) +{ + size_t oldHeadMsgCount = 0; + size_t bigBlobMsgCount = 0; + size_t newHeadMsgCount = 0; + + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (size_t i = 0; i < params.OldHeadCount; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(100'000, 'x')); + ++oldHeadMsgCount; + } + + for (size_t i = 0; i < params.BigBlobsCount; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(7'000'000, 'x'), tx.get()); + ++bigBlobMsgCount; + } + + for (size_t i = 0; i < params.NewHeadCount; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(100'000, 'x'), tx.get()); + ++newHeadMsgCount; + } + + if (params.RestartMode == ERestartBeforeCommit) { + RestartPQTablet("topic_A", 0); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + if (params.RestartMode == ERestartAfterCommit) { + RestartPQTablet("topic_A", 0); + } + + std::vector messages; + for (size_t i = 0; (i < 10) && (messages.size() < (oldHeadMsgCount + bigBlobMsgCount + newHeadMsgCount)); ++i) { + auto block = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2)); + for (auto& m : block) { + messages.push_back(std::move(m)); + } + } + + UNIT_ASSERT_VALUES_EQUAL(messages.size(), oldHeadMsgCount + bigBlobMsgCount + newHeadMsgCount); + + size_t start = 0; + + for (size_t i = 0; i < oldHeadMsgCount; ++i) { + UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 100'000); + } + start += oldHeadMsgCount; + + for (size_t i = 0; i < bigBlobMsgCount; ++i) { + UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 7'000'000); + } + start += bigBlobMsgCount; + + for (size_t i = 0; i < newHeadMsgCount; ++i) { + UNIT_ASSERT_VALUES_EQUAL(messages[start + i].size(), 100'000); + } +} + +void TFixture::CreateTable(const std::string& tablePath) +{ + UNIT_ASSERT(!tablePath.empty()); + + std::string path = (tablePath[0] != '/') ? ("/Root/" + tablePath) : tablePath; + + auto createSessionResult = GetTableClient().CreateSession().ExtractValueSync(); + UNIT_ASSERT_C(createSessionResult.IsSuccess(), createSessionResult.GetIssues().ToString()); + auto session = createSessionResult.GetSession(); + + auto desc = NTable::TTableBuilder() + .AddNonNullableColumn("key", EPrimitiveType::Utf8) + .AddNonNullableColumn("value", EPrimitiveType::Utf8) + .SetPrimaryKeyColumn("key") + .Build(); + auto result = session.CreateTable(path, std::move(desc)).GetValueSync(); + UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString()); +} + +auto TFixture::MakeTableRecords() -> std::vector +{ + std::vector records; + records.emplace_back("key-1", "value-1"); + records.emplace_back("key-2", "value-2"); + records.emplace_back("key-3", "value-3"); + records.emplace_back("key-4", "value-4"); + return records; +} + +auto TFixture::MakeJsonDoc(const std::vector& records) -> std::string +{ + auto makeJsonObject = [](const TTableRecord& r) { + return Sprintf(R"({"key":"%s", "value":"%s"})", + r.Key.data(), + r.Value.data()); + }; + + if (records.empty()) { + return "[]"; + } + + std::string s = "["; + + s += makeJsonObject(records.front()); + for (auto i = records.begin() + 1; i != records.end(); ++i) { + s += ", "; + s += makeJsonObject(*i); + } + s += "]"; + + return s; +} + +void TFixture::UpsertToTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx) +{ + auto query = Sprintf("DECLARE $key AS Utf8;" + "DECLARE $value AS Utf8;" + "UPSERT INTO `%s` (key, value) VALUES ($key, $value);", + tablePath.data()); + + for (const auto& r : records) { + auto params = TParamsBuilder() + .AddParam("$key").Utf8(r.Key).Build() + .AddParam("$value").Utf8(r.Value).Build() + .Build(); + + session.Execute(query, tx, false, params); + } +} + +void TFixture::InsertToTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx) +{ + auto query = Sprintf("DECLARE $key AS Utf8;" + "DECLARE $value AS Utf8;" + "INSERT INTO `%s` (key, value) VALUES ($key, $value);", + tablePath.data()); + + for (const auto& r : records) { + auto params = TParamsBuilder() + .AddParam("$key").Utf8(r.Key).Build() + .AddParam("$value").Utf8(r.Value).Build() + .Build(); + + session.Execute(query, tx, false, params); + } +} + +void TFixture::DeleteFromTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx) +{ + auto query = Sprintf("DECLARE $key AS Utf8;" + "DECLARE $value AS Utf8;" + "DELETE FROM `%s` ON (key, value) VALUES ($key, $value);", + tablePath.data()); + + for (const auto& r : records) { + auto params = TParamsBuilder() + .AddParam("$key").Utf8(r.Key).Build() + .AddParam("$value").Utf8(r.Value).Build() + .Build(); + + session.Execute(query, tx, false, params); + } +} + +std::size_t TFixture::GetTableRecordsCount(const std::string& tablePath) +{ + auto query = Sprintf(R"(SELECT COUNT(*) FROM `%s`)", + tablePath.data()); + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto result = session->Execute(query, tx.get()); + + NYdb::TResultSetParser parser(result.at(0)); + UNIT_ASSERT(parser.TryNextRow()); + + return parser.ColumnParser(0).GetUint64(); +} + +void TFixture::WriteMessagesInTx(std::size_t big, std::size_t small) +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (std::size_t i = 0; i < big; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(7'000'000, 'x'), tx.get(), 0); + } + + for (std::size_t i = 0; i < small; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(16'384, 'x'), tx.get(), 0); + } + + session->CommitTx(*tx, EStatus::SUCCESS); +} + +void TFixture::TestWriteToTopic38() +{ + WriteMessagesInTx(2, 202); + WriteMessagesInTx(2, 200); + WriteMessagesInTx(0, 1); + WriteMessagesInTx(4, 0); + WriteMessagesInTx(0, 1); +} + +void TFixture::TestWriteToTopic40() +{ + // The recording stream will run into a quota. Before the commit, the client will receive confirmations + // for some of the messages. The `CommitTx` call will wait for the rest. + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (std::size_t k = 0; k < 100; ++k) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); +} + +void TFixture::TestWriteToTopic41() +{ + // If the recording session does not wait for confirmations, the commit will fail + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (std::size_t k = 0; k < 100; ++k) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); + } + + CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID, true); // force close + + session->CommitTx(*tx, EStatus::SESSION_EXPIRED); +} + +void TFixture::TestWriteToTopic42() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (std::size_t k = 0; k < 100; ++k) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); + } + + CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); // gracefully close + + session->CommitTx(*tx, EStatus::SUCCESS); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); +} + +void TFixture::TestWriteToTopic43() +{ + // The recording stream will run into a quota. Before the commit, the client will receive confirmations + // for some of the messages. The `ExecuteDataQuery` call will wait for the rest. + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + for (std::size_t k = 0; k < 100; ++k) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); + } + + session->Execute("SELECT 1", tx.get()); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); +} + +void TFixture::TestWriteToTopic44() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + + auto [_, tx] = session->ExecuteInTx("SELECT 1", false); + + for (std::size_t k = 0; k < 100; ++k) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(1'000'000, 'a'), tx.get()); + } + + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(60)); + UNIT_ASSERT_EQUAL(messages.size(), 0u); + + session->Execute("SELECT 2", tx.get()); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100); +} + +void TFixture::CheckAvgWriteBytes(const std::string& topicPath, + std::uint32_t partitionId, + std::size_t minSize, std::size_t maxSize) +{ +#define UNIT_ASSERT_AVGWRITEBYTES(v, minSize, maxSize) \ + UNIT_ASSERT_LE_C(minSize, v, ", actual " << minSize << " > " << v); \ + UNIT_ASSERT_LE_C(v, maxSize, ", actual " << v << " > " << maxSize); + + auto avgWriteBytes = GetAvgWriteBytes(topicPath, partitionId); + + UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerSec, minSize, maxSize); + UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerMin, minSize, maxSize); + UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerHour, minSize, maxSize); + UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerDay, minSize, maxSize); + +#undef UNIT_ASSERT_AVGWRITEBYTES +} + +void TFixture::SplitPartition(const std::string& topicName, + std::uint32_t partitionId, + const std::string& boundary) +{ + NKikimr::NPQ::NTest::SplitPartition(Setup->GetRuntime(), + ++SchemaTxId, + TString(topicName), + partitionId, + TString(boundary)); +} + +void TFixture::TestWriteToTopic45() +{ + // Writing to a topic in a transaction affects the `AvgWriteBytes` indicator + CreateTopic("topic_A", TEST_CONSUMER, 2); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + std::string message(1'000, 'x'); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); + + session->CommitTx(*tx, EStatus::SUCCESS); + + std::size_t minSize = (message.size() + TEST_MESSAGE_GROUP_ID_1.size()) * 2; + std::size_t maxSize = minSize + 200; + + CheckAvgWriteBytes("topic_A", 0, minSize, maxSize); + + minSize = (message.size() + TEST_MESSAGE_GROUP_ID_2.size()); + maxSize = minSize + 200; + + CheckAvgWriteBytes("topic_A", 1, minSize, maxSize); +} + +void TFixture::TestWriteToTopic46() +{ + // The `split` operation of the topic partition affects the writing in the transaction. + // The transaction commit should fail with an error + CreateTopic("topic_A", TEST_CONSUMER, 2, 10); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + std::string message(1'000, 'x'); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, tx.get(), 0); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); + + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2); + + SplitPartition("topic_A", 1, "\xC0"); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, tx.get(), 1); + + session->CommitTx(*tx, EStatus::ABORTED); +} + +void TFixture::TestWriteToTopic47() +{ + // The `split` operation of the topic partition does not affect the reading in the transaction. + CreateTopic("topic_A", TEST_CONSUMER, 2, 10); + + std::string message(1'000, 'x'); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1); + + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_1); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2); + + SplitPartition("topic_A", 1, "\xC0"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); + + CloseTopicReadSession("topic_A", TEST_CONSUMER); + + messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), tx.get(), 1); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); + + session->CommitTx(*tx, EStatus::SUCCESS); +} + +void TFixture::TestWriteRandomSizedMessagesInWideTransactions() +{ + // The test verifies the simultaneous execution of several transactions. There is a topic + // with PARTITIONS_COUNT partitions. In each transaction, the test writes to all the partitions. + // The size of the messages is random. Such that both large blobs in the body and small ones in + // the head of the partition are obtained. Message sizes are multiples of 500 KB. This way we + // will make sure that when committing transactions, the division into blocks is taken into account. + + const std::size_t PARTITIONS_COUNT = 20; + const std::size_t TXS_COUNT = 10; + + CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); + + SetPartitionWriteSpeed("topic_A", 50'000'000); + + std::vector> sessions; + std::vector> transactions; + + // We open TXS_COUNT transactions and write messages to the topic. + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + sessions.push_back(CreateSession()); + auto& session = sessions.back(); + + transactions.push_back(session->BeginTx()); + auto& tx = transactions.back(); + + for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { + std::string sourceId = TEST_MESSAGE_GROUP_ID; + sourceId += "_"; + sourceId += ToString(i); + sourceId += "_"; + sourceId += ToString(j); + + std::size_t count = RandomNumber(20) + 3; + WriteToTopic("topic_A", sourceId, std::string(512 * 1000 * count, 'x'), tx.get(), j); + + WaitForAcks("topic_A", sourceId); + } + } + + // We are doing an asynchronous commit of transactions. They will be executed simultaneously. + std::vector futures; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); + } + + // All transactions must be completed successfully. + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures[i].Wait(); + const auto& result = futures[i].GetValueSync(); + UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); + } +} + +void TFixture::TestWriteOnlyBigMessagesInWideTransactions() +{ + // The test verifies the simultaneous execution of several transactions. There is a topic `topic_A` and + // it contains a `PARTITIONS_COUNT' of partitions. In each transaction, the test writes to all partitions. + // The size of the messages is chosen so that only large blobs are recorded in the transaction and there + // are no records in the head. Thus, we verify that transaction bundling is working correctly. + + const std::size_t PARTITIONS_COUNT = 20; + const std::size_t TXS_COUNT = 100; + + CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); + + SetPartitionWriteSpeed("topic_A", 50'000'000); + + std::vector> sessions; + std::vector> transactions; + + // We open TXS_COUNT transactions and write messages to the topic. + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + sessions.push_back(CreateSession()); + auto& session = sessions.back(); + + transactions.push_back(session->BeginTx()); + auto& tx = transactions.back(); + + for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { + std::string sourceId = TEST_MESSAGE_GROUP_ID; + sourceId += "_"; + sourceId += ToString(i); + sourceId += "_"; + sourceId += ToString(j); + + WriteToTopic("topic_A", sourceId, std::string(6'500'000, 'x'), tx.get(), j); + + WaitForAcks("topic_A", sourceId); + } + } + + // We are doing an asynchronous commit of transactions. They will be executed simultaneously. + std::vector futures; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); + } + + // All transactions must be completed successfully. + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures[i].Wait(); + const auto& result = futures[i].GetValueSync(); + UNIT_ASSERT_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); + } +} + +void TFixture::TestTransactionsConflictOnSeqNo() +{ + const std::uint32_t PARTITIONS_COUNT = 20; + const std::size_t TXS_COUNT = 100; + + CreateTopic("topic_A", TEST_CONSUMER, PARTITIONS_COUNT); + + SetPartitionWriteSpeed("topic_A", 50'000'000); + + auto session = CreateSession(); + std::vector> topicWriteSessions; + + for (std::uint32_t i = 0; i < PARTITIONS_COUNT; ++i) { + std::string sourceId = TEST_MESSAGE_GROUP_ID; + sourceId += "_"; + sourceId += ToString(i); + + NTopic::TTopicClient client(GetDriver()); + NTopic::TWriteSessionSettings options; + options.Path(Setup->GetTopicPath("topic_A")); + options.ProducerId(sourceId); + options.MessageGroupId(sourceId); + options.PartitionId(i); + options.Codec(ECodec::RAW); + + auto session = client.CreateSimpleBlockingWriteSession(options); + + topicWriteSessions.push_back(std::move(session)); + } + + std::vector> sessions; + std::vector> transactions; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + sessions.push_back(CreateSession()); + auto& session = sessions.back(); + + transactions.push_back(session->BeginTx()); + auto& tx = transactions.back(); + + for (std::size_t j = 0; j < PARTITIONS_COUNT; ++j) { + std::string sourceId = TEST_MESSAGE_GROUP_ID; + sourceId += "_"; + sourceId += ToString(j); + + for (std::size_t k = 0, count = RandomNumber(20) + 1; k < count; ++k) { + const std::string data(RandomNumber(1'000) + 100, 'x'); + NTopic::TWriteMessage params(data); + params.Tx(*tx); + + topicWriteSessions[j]->Write(std::move(params)); + } + } + } + + std::vector futures; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures.push_back(sessions[i]->AsyncCommitTx(*transactions[i])); + } + + // Some transactions should end with the error `ABORTED` + std::size_t successCount = 0; + + for (std::size_t i = 0; i < TXS_COUNT; ++i) { + futures[i].Wait(); + const auto& result = futures[i].GetValueSync(); + switch (result.GetStatus()) { + case EStatus::SUCCESS: + ++successCount; + break; + case EStatus::ABORTED: + break; + default: + UNIT_FAIL("unexpected status: " + ToString(static_cast(result))); + break; + } + } + + UNIT_ASSERT_UNEQUAL(successCount, TXS_COUNT); +} + +void TFixtureSinks::CreateRowTable(const std::string& path) +{ + CreateTable(path); +} + +void TFixtureSinks::CreateColumnTable(const std::string& tablePath) +{ + UNIT_ASSERT(!tablePath.empty()); + + std::string path = (tablePath[0] != '/') ? ("/Root/" + tablePath) : tablePath; + + auto createSessionResult = GetTableClient().CreateSession().ExtractValueSync(); + UNIT_ASSERT_C(createSessionResult.IsSuccess(), createSessionResult.GetIssues().ToString()); + auto session = createSessionResult.GetSession(); + + auto desc = NTable::TTableBuilder() + .SetStoreType(NTable::EStoreType::Column) + .AddNonNullableColumn("key", EPrimitiveType::Utf8) + .AddNonNullableColumn("value", EPrimitiveType::Utf8) + .SetPrimaryKeyColumn("key") + .Build(); + auto result = session.CreateTable(path, std::move(desc)).GetValueSync(); + 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; +} + +bool TFixtureSinks::GetAllowOlapDataQuery() const +{ + return true; +} + +void TFixtureSinks::TestSinksOltpWriteToTopic5() +{ + CreateTopic("topic_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); + + session->RollbackTx(*tx, EStatus::SUCCESS); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); +} + +void TFixtureSinks::TestSinksOltpWriteToTopicAndTable2() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + CreateRowTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 3); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #3"); + } + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + + CheckTabletKeys("topic_A"); + CheckTabletKeys("topic_B"); +} + +void TFixtureSinks::TestSinksOltpWriteToTopicAndTable3() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + + CreateRowTable("/Root/table_A"); + CreateRowTable("/Root/table_B"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + UpsertToTable("table_B", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + const size_t topicMsgCnt = 10; + for (size_t i = 1; i <= topicMsgCnt; ++i) { + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); + } + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), records.size()); + + CheckTabletKeys("topic_A"); + CheckTabletKeys("topic_B"); +} + +void TFixtureSinks::TestSinksOltpWriteToTopicAndTable4() +{ + CreateTopic("topic_A"); + CreateRowTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx1 = session->BeginTx(); + auto tx2 = session->BeginTx(); + + session->Execute(R"(SELECT COUNT(*) FROM `table_A`)", tx1.get(), false); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx2.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx1.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + session->CommitTx(*tx2, EStatus::SUCCESS); + session->CommitTx(*tx1, EStatus::ABORTED); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + + CheckTabletKeys("topic_A"); +} + +void TFixtureSinks::TestSinksOltpWriteToTopicAndTable5() +{ + CreateTopic("topic_A"); + CreateRowTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + session->RollbackTx(*tx, EStatus::SUCCESS); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); + + CheckTabletKeys("topic_A"); +} + +void TFixtureSinks::TestSinksOltpWriteToTopicAndTable6() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + CreateRowTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + InsertToTable("table_A", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #1", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #2", tx.get()); + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #3", tx.get()); + + DeleteFromTable("table_A", records, *session, tx.get()); + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, 3); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #3"); + } + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); + + CheckTabletKeys("topic_A"); + CheckTabletKeys("topic_B"); +} + +void TFixtureSinks::TestSinksOlapWriteToTopicAndTable1() +{ + CreateTopic("topic_A"); + CreateColumnTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + session->CommitTx(*tx, EStatus::SUCCESS); + + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + + CheckTabletKeys("topic_A"); +} + +void TFixtureSinks::TestSinksOlapWriteToTopicAndTable2() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + + CreateRowTable("/Root/table_A"); + CreateColumnTable("/Root/table_B"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + + UpsertToTable("table_A", records, *session, tx.get()); + UpsertToTable("table_B", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + const size_t topicMsgCnt = 10; + for (size_t i = 1; i <= topicMsgCnt; ++i) { + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); + } + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); + } + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), records.size()); + + CheckTabletKeys("topic_A"); + CheckTabletKeys("topic_B"); +} + +void TFixtureSinks::TestSinksOlapWriteToTopicAndTable3() +{ + CreateTopic("topic_A"); + CreateColumnTable("/Root/table_A"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + UpsertToTable("table_A", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID); + + session->RollbackTx(*tx, EStatus::SUCCESS); + + Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 0); + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), 0); + + CheckTabletKeys("topic_A"); +} + +void TFixtureSinks::TestSinksOlapWriteToTopicAndTable4() +{ + CreateTopic("topic_A"); + CreateTopic("topic_B"); + + CreateRowTable("/Root/table_A"); + CreateColumnTable("/Root/table_B"); + CreateColumnTable("/Root/table_C"); + + auto session = CreateSession(); + auto tx = session->BeginTx(); + + auto records = MakeTableRecords(); + + InsertToTable("table_A", records, *session, tx.get()); + InsertToTable("table_B", records, *session, tx.get()); + UpsertToTable("table_C", records, *session, tx.get()); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, MakeJsonDoc(records), tx.get()); + + const size_t topicMsgCnt = 10; + for (size_t i = 1; i <= topicMsgCnt; ++i) { + WriteToTopic("topic_B", TEST_MESSAGE_GROUP_ID, "message #" + std::to_string(i), tx.get()); + } + + DeleteFromTable("table_B", records, *session, tx.get()); + + session->CommitTx(*tx, EStatus::SUCCESS); + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 1); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), MakeJsonDoc(records)); + } + + { + auto messages = Read_Exactly_N_Messages_From_Topic("topic_B", TEST_CONSUMER, topicMsgCnt); + UNIT_ASSERT_VALUES_EQUAL(messages.front(), "message #1"); + UNIT_ASSERT_VALUES_EQUAL(messages.back(), "message #" + std::to_string(topicMsgCnt)); + } + + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_A"), records.size()); + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_B"), 0); + UNIT_ASSERT_VALUES_EQUAL(GetTableRecordsCount("table_C"), records.size()); + + CheckTabletKeys("topic_A"); + CheckTabletKeys("topic_B"); +} + +void TFixture::TestWriteAndReadMessages(size_t count, size_t size, bool restart) +{ + CreateTopic("topic_A"); + + SetPartitionWriteSpeed("topic_A", 50'000'000); + + for (size_t i = 0; i < count; ++i) { + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID, std::string(size, 'x'), nullptr); + } + CloseTopicWriteSession("topic_A", TEST_MESSAGE_GROUP_ID); + + if (restart) { + RestartPQTablet("topic_A", 0); + } + + auto messages = Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, count); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), count); +} + +} diff --git a/src/client/topic/ut/ut_utils/txusage_fixture.h b/src/client/topic/ut/ut_utils/txusage_fixture.h new file mode 100644 index 00000000000..1c953bc6cc3 --- /dev/null +++ b/src/client/topic/ut/ut_utils/txusage_fixture.h @@ -0,0 +1,439 @@ +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include + +#include + +namespace NYdb::inline V3::NTopic::NTests::NTxUsage { + +class TFixture : public NUnitTest::TBaseFixture { +protected: + using TTopicReadSession = NTopic::IReadSession; + using TTopicReadSessionPtr = std::shared_ptr; + using TTopicWriteSession = NTopic::IWriteSession; + using TTopicWriteSessionPtr = std::shared_ptr; + + struct TTopicWriteSessionContext { + TTopicWriteSessionPtr Session; + std::optional ContinuationToken; + size_t WriteCount = 0; + size_t WrittenAckCount = 0; + size_t WrittenInTxAckCount = 0; + + void WaitForContinuationToken(); + void Write(const std::string& message, TTransactionBase* tx = nullptr); + + size_t AckCount() const { return WrittenAckCount + WrittenInTxAckCount; } + + void WaitForEvent(); + }; + + struct TFeatureFlags { + bool EnablePQConfigTransactionsAtSchemeShard = true; + }; + + class ISession { + public: + using TExecuteInTxResult = std::pair, std::unique_ptr>; + + virtual std::vector Execute(const std::string& query, + TTransactionBase* tx, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) = 0; + + virtual TExecuteInTxResult ExecuteInTx(const std::string& query, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) = 0; + + virtual std::unique_ptr BeginTx() = 0; + virtual void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) = 0; + virtual void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) = 0; + + virtual void Close() = 0; + + virtual TAsyncStatus AsyncCommitTx(TTransactionBase& tx) = 0; + + virtual ~ISession() = default; + }; + + void SetUp(NUnitTest::TTestContext&) override; + + void NotifySchemeShard(const TFeatureFlags& flags); + + std::unique_ptr CreateSession(); + + struct TReadMessageSettings { + TTransactionBase& Tx; + bool CommitOffsets = false; + std::optional Offset; + }; + + void CreateTopic(const std::string& path = TEST_TOPIC, + const std::string& consumer = TEST_CONSUMER, + std::size_t partitionCount = 1, + std::optional maxPartitionCount = std::nullopt, + const TDuration retention = TDuration::Hours(1), + bool important = false); + + void AddConsumer(const std::string& topicPath, const std::vector& consumers); + + void SetPartitionWriteSpeed(const std::string& topicName, std::size_t bytesPerSeconds); + + TTopicWriteSessionPtr CreateTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + std::optional partitionId); + TTopicWriteSessionContext& GetTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + std::optional partitionId); + + TTopicReadSessionPtr CreateTopicReadSession(const std::string& topicPath, + const std::string& consumerName, + std::optional partitionId); + TTopicReadSessionPtr GetTopicReadSession(const std::string& topicPath, + const std::string& consumerName, + std::optional partitionId); + + void WriteToTopic(const std::string& topicPath, + const std::string& messageGroupId, + const std::string& message, + TTransactionBase* tx = nullptr, + std::optional partitionId = std::nullopt); + std::vector ReadFromTopic(const std::string& topicPath, + const std::string& consumerName, + const TDuration& duration, + TTransactionBase* tx = nullptr, + std::optional partitionId = std::nullopt); + void WaitForAcks(const std::string& topicPath, + const std::string& messageGroupId, + std::size_t writtenInTxCount = std::numeric_limits::max()); + void WaitForSessionClose(const std::string& topicPath, + const std::string& messageGroupId, + NYdb::EStatus status); + void CloseTopicWriteSession(const std::string& topicPath, + const std::string& messageGroupId, + bool force = false); + void CloseTopicReadSession(const std::string& topicPath, + const std::string& consumerName); + + enum EEndOfTransaction { + Commit, + Rollback, + CloseTableSession + }; + + struct TTransactionCompletionTestDescription { + std::vector Topics; + EEndOfTransaction EndOfTransaction = Commit; + }; + + void TestTheCompletionOfATransaction(const TTransactionCompletionTestDescription& d); + void RestartPQTablet(const std::string& topicPath, std::uint32_t partition); + void DumpPQTabletKeys(const std::string& topicName, std::uint32_t partition); + void PQTabletPrepareFromResource(const std::string& topicPath, + std::uint32_t partitionId, + const std::string& resourceName); + + void DeleteSupportivePartition(const std::string& topicName, + std::uint32_t partition); + + struct TTableRecord { + TTableRecord() = default; + TTableRecord(const std::string& key, const std::string& value); + + std::string Key; + std::string Value; + }; + + std::vector MakeTableRecords(); + std::string MakeJsonDoc(const std::vector& records); + + void CreateTable(const std::string& path); + void UpsertToTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx); + void InsertToTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx); + void DeleteFromTable(const std::string& tablePath, + const std::vector& records, + ISession& session, + TTransactionBase* tx); + size_t GetTableRecordsCount(const std::string& tablePath); + + enum ERestartPQTabletMode { + ERestartNo, + ERestartBeforeCommit, + ERestartAfterCommit, + }; + + struct TTestTxWithBigBlobsParams { + size_t OldHeadCount = 0; + size_t BigBlobsCount = 2; + size_t NewHeadCount = 0; + ERestartPQTabletMode RestartMode = ERestartNo; + }; + + void TestTxWithBigBlobs(const TTestTxWithBigBlobsParams& params); + + void WriteMessagesInTx(std::size_t big, size_t small); + + const TDriver& GetDriver() const; + NTable::TTableClient& GetTableClient(); + + void CheckTabletKeys(const std::string& topicName); + void DumpPQTabletKeys(const std::string& topicName); + + std::vector Read_Exactly_N_Messages_From_Topic(const std::string& topicPath, + const std::string& consumerName, + size_t count); + + void TestWriteRandomSizedMessagesInWideTransactions(); + + void TestWriteOnlyBigMessagesInWideTransactions(); + + void TestTransactionsConflictOnSeqNo(); + + void TestWriteToTopic1(); + + void TestWriteToTopic4(); + + void TestWriteToTopic7(); + + void TestWriteToTopic9(); + + void TestWriteToTopic10(); + + void TestWriteToTopic11(); + + void TestWriteToTopic12(); + + void TestWriteToTopic13(); + + void TestWriteToTopic14(); + + void TestWriteToTopic16(); + + void TestWriteToTopic24(); + + void TestWriteToTopic26(); + + void TestWriteToTopic27(); + + void TestWriteToTopic38(); + + void TestWriteToTopic40(); + + void TestWriteToTopic41(); + + void TestWriteToTopic42(); + + void TestWriteToTopic43(); + + void TestWriteToTopic44(); + + void TestWriteToTopic45(); + + void TestWriteToTopic46(); + + void TestWriteToTopic47(); + + void TestWriteToTopic50(); + + struct TAvgWriteBytes { + std::uint64_t PerSec = 0; + std::uint64_t PerMin = 0; + std::uint64_t PerHour = 0; + std::uint64_t PerDay = 0; + }; + + TAvgWriteBytes GetAvgWriteBytes(const std::string& topicPath, + std::uint32_t partitionId); + + void CheckAvgWriteBytes(const std::string& topicPath, + std::uint32_t partitionId, + std::size_t minSize, std::size_t maxSize); + + void SplitPartition(const std::string& topicPath, + std::uint32_t partitionId, + const std::string& boundary); + + virtual bool GetEnableOltpSink() const; + virtual bool GetEnableOlapSink() const; + virtual bool GetEnableHtapTx() const; + virtual bool GetAllowOlapDataQuery() const; + + size_t GetPQCacheRenameKeysCount(); + + enum class EClientType { + Table, + Query, + None + }; + + virtual EClientType GetClientType() const = 0; + virtual ~TFixture() = default; + + void TestWriteAndReadMessages(size_t count, size_t size, bool restart); + +private: + class TTableSession : public ISession { + public: + TTableSession(NTable::TTableClient& client); + + std::vector Execute(const std::string& query, + TTransactionBase* tx, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) override; + + TExecuteInTxResult ExecuteInTx(const std::string& query, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) override; + + std::unique_ptr BeginTx() override; + void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; + void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; + + TAsyncStatus AsyncCommitTx(TTransactionBase& tx) override; + + void Close() override; + + private: + NTable::TSession Init(NTable::TTableClient& client); + + NTable::TSession Session_; + }; + + class TQuerySession : public ISession { + public: + TQuerySession(NQuery::TQueryClient& client, + const std::string& endpoint, + const std::string& database); + + std::vector Execute(const std::string& query, + TTransactionBase* tx, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) override; + + TExecuteInTxResult ExecuteInTx(const std::string& query, + bool commit = true, + const TParams& params = TParamsBuilder().Build()) override; + + std::unique_ptr BeginTx() override; + void CommitTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; + void RollbackTx(TTransactionBase& tx, EStatus status = EStatus::SUCCESS) override; + + TAsyncStatus AsyncCommitTx(TTransactionBase& tx) override; + + void Close() override; + + private: + NQuery::TSession Init(NQuery::TQueryClient& client); + + NQuery::TSession Session_; + std::string Endpoint_; + std::string Database_; + }; + + template + E ReadEvent(TTopicReadSessionPtr reader, TTransactionBase& tx); + template + E ReadEvent(TTopicReadSessionPtr reader); + + std::uint64_t GetTopicTabletId(const NActors::TActorId& actorId, + const std::string& topicPath, + std::uint32_t partition); + std::vector GetTabletKeys(const NActors::TActorId& actorId, + std::uint64_t tabletId); + NKikimr::NPQ::TWriteId GetTransactionWriteId(const NActors::TActorId& actorId, + std::uint64_t tabletId); + void SendLongTxLockStatus(const NActors::TActorId& actorId, + std::uint64_t tabletId, + const NKikimr::NPQ::TWriteId& writeId, + NKikimrLongTxService::TEvLockStatus::EStatus status); + void WaitForTheTabletToDeleteTheWriteInfo(const NActors::TActorId& actorId, + std::uint64_t tabletId, + const NKikimr::NPQ::TWriteId& writeId); + + std::uint64_t GetSchemeShardTabletId(const NActors::TActorId& actorId); + + std::unique_ptr Setup; + std::unique_ptr Driver; + std::unique_ptr TableClient; + std::unique_ptr QueryClient; + + std::unordered_map, TTopicWriteSessionContext> TopicWriteSessions; + std::unordered_map TopicReadSessions; + + ui64 SchemaTxId = 1000; +}; + +class TFixtureTable : public TFixture { +protected: + EClientType GetClientType() const override { + return EClientType::Table; + } +}; + +class TFixtureQuery : public TFixture { +protected: + EClientType GetClientType() const override { + return EClientType::Query; + } +}; + +class TFixtureNoClient : public TFixture { +protected: + EClientType GetClientType() const override { + return EClientType::None; + } +}; + +class TFixtureSinks : public TFixture { +protected: + 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; + + void TestSinksOltpWriteToTopic5(); + + void TestSinksOltpWriteToTopicAndTable2(); + void TestSinksOltpWriteToTopicAndTable3(); + void TestSinksOltpWriteToTopicAndTable4(); + void TestSinksOltpWriteToTopicAndTable5(); + void TestSinksOltpWriteToTopicAndTable6(); + + void TestSinksOlapWriteToTopicAndTable1(); + void TestSinksOlapWriteToTopicAndTable2(); + void TestSinksOlapWriteToTopicAndTable3(); + void TestSinksOlapWriteToTopicAndTable4(); +}; + +class TFixtureSinksTable : public TFixtureSinks { +protected: + EClientType GetClientType() const override { + return EClientType::Table; + } +}; + +class TFixtureSinksQuery : public TFixtureSinks { +protected: + EClientType GetClientType() const override { + return EClientType::Query; + } +}; + +} diff --git a/src/client/types/core_facility/core_facility.h b/src/client/types/core_facility/core_facility.h index 772512bdb6b..6d3a3035a58 100644 --- a/src/client/types/core_facility/core_facility.h +++ b/src/client/types/core_facility/core_facility.h @@ -3,6 +3,8 @@ #include #include +#include + namespace NYdb::inline V3 { using TPeriodicCb = std::function; @@ -13,7 +15,7 @@ class ICoreFacility { virtual ~ICoreFacility() = default; // Add task to execute periodicaly // Task should return false to stop execution - virtual void AddPeriodicTask(TPeriodicCb&& cb, TDuration period) = 0; + virtual void AddPeriodicTask(TPeriodicCb&& cb, TDeadline::Duration period) = 0; }; } // namespace NYdb diff --git a/src/client/types/credentials/login/login.cpp b/src/client/types/credentials/login/login.cpp index a2a089a7bd2..b0c3bd64333 100644 --- a/src/client/types/credentials/login/login.cpp +++ b/src/client/types/credentials/login/login.cpp @@ -9,6 +9,8 @@ #include +using namespace std::chrono_literals; + namespace NYdb::inline V3 { namespace { @@ -97,7 +99,7 @@ TLoginCredentialsProvider::TLoginCredentialsProvider(std::weak_ptrAddPeriodicTask(std::move(periodicTask), TDuration::Minutes(1)); + strongFacility->AddPeriodicTask(std::move(periodicTask), 1min); } } @@ -137,7 +139,7 @@ void TLoginCredentialsProvider::RequestToken() { request.set_user(TStringType{Params_.User}); request.set_password(TStringType{Params_.Password}); TRpcRequestSettings rpcSettings; - rpcSettings.ClientTimeout = TDuration::Seconds(60); + rpcSettings.Deadline = TDeadline::AfterDuration(60s); TGRpcConnectionsImpl::RunOnDiscoveryEndpoint( strongFacility, std::move(request), std::move(responseCb), &Ydb::Auth::V1::AuthService::Stub::AsyncLogin, diff --git a/src/library/CMakeLists.txt b/src/library/CMakeLists.txt index 33f373a28d9..3d3d0a24c60 100644 --- a/src/library/CMakeLists.txt +++ b/src/library/CMakeLists.txt @@ -15,4 +15,5 @@ add_subdirectory(retry) add_subdirectory(string_utils/base64) add_subdirectory(string_utils/helpers) add_subdirectory(string_utils/misc) +add_subdirectory(time) add_subdirectory(uuid) diff --git a/src/library/grpc/client/CMakeLists.txt b/src/library/grpc/client/CMakeLists.txt index aef53794904..d6543cebbaa 100644 --- a/src/library/grpc/client/CMakeLists.txt +++ b/src/library/grpc/client/CMakeLists.txt @@ -3,6 +3,7 @@ _ydb_sdk_add_library(grpc-client) target_link_libraries(grpc-client PUBLIC yutil + ydb-time gRPC::grpc++ ) diff --git a/src/library/grpc/client/grpc_client_low.cpp b/src/library/grpc/client/grpc_client_low.cpp index b4421b2f1c0..d23ef4d3087 100644 --- a/src/library/grpc/client/grpc_client_low.cpp +++ b/src/library/grpc/client/grpc_client_low.cpp @@ -20,8 +20,7 @@ #include -namespace NYdbGrpc { -inline namespace V3 { +namespace NYdbGrpc::inline V3 { void EnableGRpcTracing() { grpc_tracer_set_enabled("tcp", true); @@ -111,6 +110,29 @@ grpc_socket_mutator_vtable TGRpcKeepAliveSocketMutator::VTable = }; #endif +void TGRpcRequestProcessorCommon::ApplyMeta(const TCallMeta& meta) { + for (const auto& rec : meta.Aux) { + Context.AddMetadata(NYdb::TStringType{rec.first}, NYdb::TStringType{rec.second}); + } + if (meta.CallCredentials) { + Context.set_credentials(meta.CallCredentials); + } + if (const NYdb::TDeadline::Duration* timeout = std::get_if(&meta.Timeout)) { + Context.set_deadline(NYdb::TDeadline::AfterDuration(*timeout)); + } else if (const NYdb::TDeadline* deadline = std::get_if(&meta.Timeout)) { + Context.set_deadline(*deadline); + } +} + +void TGRpcRequestProcessorCommon::GetInitialMetadata(std::unordered_multimap* metadata) { + for (const auto& [key, value] : Context.GetServerInitialMetadata()) { + metadata->emplace( + std::string(key.begin(), key.end()), + std::string(value.begin(), value.end()) + ); + } +} + TChannelPool::TChannelPool(const TTcpKeepAliveSettings& tcpKeepAliveSettings, const TDuration& expireTime) : TcpKeepAliveSettings_(tcpKeepAliveSettings) , ExpireTime_(expireTime) @@ -604,5 +626,27 @@ grpc_socket_mutator* NImpl::CreateGRpcKeepAliveSocketMutator(const TTcpKeepAlive return nullptr; } +gpr_timespec DurationToTimespec(const NYdb::TDeadline::Duration& duration) noexcept { + const auto secs = std::chrono::floor(duration); + if (duration == NYdb::TDeadline::Duration::max() || secs.count() >= gpr_inf_future(GPR_CLOCK_MONOTONIC).tv_sec) { + return gpr_inf_future(GPR_TIMESPAN); + } + if (secs.count() < 0) { + return gpr_inf_past(GPR_TIMESPAN); + } + const auto nsecs = std::chrono::duration_cast(duration - secs); + Y_ASSERT(0 <= nsecs.count() && nsecs.count() < GPR_NS_PER_SEC); + gpr_timespec t; + t.tv_sec = static_cast(secs.count()); + t.tv_nsec = static_cast(nsecs.count()); + t.clock_type = GPR_TIMESPAN; + return t; } + +gpr_timespec DeadlineToTimespec(const NYdb::TDeadline& deadline) { + gpr_timespec t = + DurationToTimespec(deadline.GetTimePoint() != NYdb::TDeadline::TimePoint::max() ? deadline.GetTimePoint() - NYdb::TDeadline::Clock::now() : NYdb::TDeadline::Duration::max()); + return gpr_convert_clock_type(t, GPR_CLOCK_MONOTONIC); +} + } diff --git a/src/library/grpc/client/grpc_client_low.h b/src/library/grpc/client/grpc_client_low.h index f1e2156e8af..1b190087f86 100644 --- a/src/library/grpc/client/grpc_client_low.h +++ b/src/library/grpc/client/grpc_client_low.h @@ -4,6 +4,8 @@ #include +#include + #include #include #include @@ -24,8 +26,7 @@ * This file contains low level logic for grpc * This file should not be used in high level code without special reason */ -namespace NYdbGrpc { -inline namespace V3 { +namespace NYdbGrpc::inline V3 { const size_t DEFAULT_NUM_THREADS = 2; @@ -198,40 +199,14 @@ using TAdvancedResponseCallback = std::function CallCredentials; std::vector> Aux; - std::variant Timeout; // timeout as duration from now or time point in future + std::variant Timeout; // timeout as duration from now or time point in future }; class TGRpcRequestProcessorCommon { protected: - void ApplyMeta(const TCallMeta& meta) { - for (const auto& rec : meta.Aux) { - Context.AddMetadata(NYdb::TStringType{rec.first}, NYdb::TStringType{rec.second}); - } - if (meta.CallCredentials) { - Context.set_credentials(meta.CallCredentials); - } - if (const TDuration* timeout = std::get_if(&meta.Timeout)) { - if (*timeout) { - auto deadline = gpr_time_add( - gpr_now(GPR_CLOCK_MONOTONIC), - gpr_time_from_micros(timeout->MicroSeconds(), GPR_TIMESPAN)); - Context.set_deadline(deadline); - } - } else if (const TInstant* deadline = std::get_if(&meta.Timeout)) { - if (*deadline) { - Context.set_deadline(gpr_time_from_micros(deadline->MicroSeconds(), GPR_CLOCK_MONOTONIC)); - } - } - } + void ApplyMeta(const TCallMeta& meta); - void GetInitialMetadata(std::unordered_multimap* metadata) { - for (const auto& [key, value] : Context.GetServerInitialMetadata()) { - metadata->emplace( - std::string(key.begin(), key.end()), - std::string(value.begin(), value.end()) - ); - } - } + void GetInitialMetadata(std::unordered_multimap* metadata); grpc::Status Status; grpc::ClientContext Context; @@ -1437,5 +1412,23 @@ class TGRpcClientLow std::mutex JoinMutex_; }; +gpr_timespec DurationToTimespec(const NYdb::TDeadline::Duration& duration) noexcept; + +gpr_timespec DeadlineToTimespec(const NYdb::TDeadline& deadline); + } -} + +template <> +class grpc::TimePoint { +public: + TimePoint(const NYdb::TDeadline& deadline) + : time_(NYdbGrpc::DeadlineToTimespec(deadline)) + {} + + gpr_timespec raw_time() const noexcept { + return time_; + } + +private: + gpr_timespec time_; +}; diff --git a/src/library/time/CMakeLists.txt b/src/library/time/CMakeLists.txt new file mode 100644 index 00000000000..5fafcbf4fdf --- /dev/null +++ b/src/library/time/CMakeLists.txt @@ -0,0 +1,12 @@ +_ydb_sdk_add_library(ydb-time) + +target_link_libraries(ydb-time + PUBLIC + yutil +) + +target_sources(ydb-time PRIVATE + time.cpp +) + +_ydb_sdk_install_targets(TARGETS ydb-time) diff --git a/src/library/time/time.cpp b/src/library/time/time.cpp new file mode 100644 index 00000000000..eba86aa66fd --- /dev/null +++ b/src/library/time/time.cpp @@ -0,0 +1,50 @@ +#include "time.h" + +namespace NYdb::inline V3 { + +namespace { + TDeadline::TimePoint SafeSum(TDeadline::TimePoint timePoint, TDeadline::Duration duration) { + if (duration > TDeadline::TimePoint::max() - timePoint) { + return TDeadline::TimePoint::max(); + } + return timePoint + duration; + } +} + +TDeadline TDeadline::Now() { + return TDeadline(Clock::now()); +} + +TDeadline TDeadline::Max() { + return TDeadline(TimePoint::max()); +} + +TDeadline TDeadline::AfterDuration(const TDuration& after) { + return AfterDuration(SafeDurationCast(after)); +} + +TDeadline TDeadline::AfterDuration(const Duration& after) { + return TDeadline(SafeSum(Clock::now(), after)); +} + +TDeadline::Duration TDeadline::SafeDurationCast(const TDuration& after) { + if (after.MicroSeconds() > static_cast(std::chrono::microseconds::max().count())) { + return Duration::max(); + } + std::chrono::microseconds afterMicros(after.MicroSeconds()); + return SafeDurationCast(afterMicros); +} + +TDeadline::TimePoint TDeadline::GetTimePoint() const noexcept { + return TimePoint_; +} + +bool TDeadline::operator<(const TDeadline& other) const noexcept { + return TimePoint_ < other.TimePoint_; +} + +bool TDeadline::operator==(const TDeadline& other) const noexcept { + return TimePoint_ == other.TimePoint_; +} + +} // namespace NYdb diff --git a/src/library/time/time.h b/src/library/time/time.h new file mode 100644 index 00000000000..a4ece9db770 --- /dev/null +++ b/src/library/time/time.h @@ -0,0 +1,66 @@ +#pragma once + +#include + +#include + +namespace NYdb::inline V3 { + +class TDeadline { +public: + using Clock = std::chrono::steady_clock; + using TimePoint = Clock::time_point; + using Duration = Clock::duration; + + TDeadline() = default; + + static TDeadline Now(); + + static TDeadline Max(); + + static TDeadline AfterDuration(const TDuration& after); + + static TDeadline AfterDuration(const Duration& after); + + template + static TDeadline AfterDuration(const std::chrono::duration& after) { + return AfterDuration(SafeDurationCast(after)); + } + + static Duration SafeDurationCast(const TDuration& duration); + + template + static Duration SafeDurationCast(const std::chrono::duration& duration) { + using FromDuration = std::chrono::duration; + + // Require resolution of 'Duration' higher than 'FromDuration', + // to safely cast 'Duration::max' value to 'FromDuration' + static_assert(std::is_constructible_v); + + if (std::chrono::duration_cast(Duration::max()) < duration) { + return Duration::max(); + } + + if constexpr (std::is_signed_v) { + if (duration < std::chrono::duration_cast(Duration::min())) { + return Duration::min(); + } + } + + return std::chrono::duration_cast(duration); + } + + TimePoint GetTimePoint() const noexcept; + + bool operator<(const TDeadline& other) const noexcept; + bool operator==(const TDeadline& other) const noexcept; + +private: + explicit TDeadline(TimePoint timePoint) noexcept + : TimePoint_(timePoint) + {} + + TimePoint TimePoint_; +}; + +} // namespace NYdb diff --git a/tests/integration/basic_example/basic_example.cpp b/tests/integration/basic_example/basic_example.cpp index 14e6c347d41..4ec19ac46a2 100644 --- a/tests/integration/basic_example/basic_example.cpp +++ b/tests/integration/basic_example/basic_example.cpp @@ -35,6 +35,7 @@ TRunArgs GetRunArgs() { //! Creates sample tables with CrateTable API. void CreateTables(TTableClient client, const std::string& path) { + std::cout << "CreateTables step" << std::endl; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path](TSession session) { auto seriesDesc = TTableBuilder() .AddNullableColumn("series_id", NYdb::EPrimitiveType::Uint64) @@ -459,6 +460,7 @@ static NYdb::TStatus ScanQuerySelect(TTableClient client, const std::string& pat /////////////////////////////////////////////////////////////////////////////// NYdb::TResultSet SelectSimple(TTableClient client, const std::string& path) { + std::cout << "SelectSimple step" << std::endl; std::optional resultSet; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path, &resultSet](TSession session) { return SelectSimpleTransaction(session, path, resultSet); @@ -468,12 +470,14 @@ NYdb::TResultSet SelectSimple(TTableClient client, const std::string& path) { } void UpsertSimple(TTableClient client, const std::string& path) { + std::cout << "UpsertSimple step" << std::endl; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path](TSession session) { return UpsertSimpleTransaction(session, path); })); } NYdb::TResultSet SelectWithParams(TTableClient client, const std::string& path) { + std::cout << "SelectWithParams step" << std::endl; std::optional resultSet; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path, &resultSet](TSession session) { return SelectWithParamsTransaction(session, path, 2, 3, resultSet); @@ -483,6 +487,7 @@ NYdb::TResultSet SelectWithParams(TTableClient client, const std::string& path) } NYdb::TResultSet PreparedSelect(TTableClient client, const std::string& path, ui32 seriesId, ui32 seasonId, ui32 episodeId) { + std::cout << "PreparedSelect step" << std::endl; std::optional resultSet; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path, seriesId, seasonId, episodeId, &resultSet](TSession session) { return PreparedSelectTransaction(session, path, seriesId, seasonId, episodeId, resultSet); @@ -492,6 +497,7 @@ NYdb::TResultSet PreparedSelect(TTableClient client, const std::string& path, ui } NYdb::TResultSet MultiStep(TTableClient client, const std::string& path) { + std::cout << "MultiStep step" << std::endl; std::optional resultSet; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path, &resultSet](TSession session) { return MultiStepTransaction(session, path, 2, 5, resultSet); @@ -501,12 +507,14 @@ NYdb::TResultSet MultiStep(TTableClient client, const std::string& path) { } void ExplicitTcl(TTableClient client, const std::string& path) { + std::cout << "ExplicitTcl step" << std::endl; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path](TSession session) { return ExplicitTclTransaction(session, path, TInstant()); })); } std::vector ScanQuerySelect(TTableClient client, const std::string& path) { + std::cout << "ScanQuerySelect step" << std::endl; std::vector resultSets; NYdb::NStatusHelpers::ThrowOnError(client.RetryOperationSync([path, &resultSets](TTableClient& client) { return ScanQuerySelect(client, path, resultSets); diff --git a/tests/slo_workloads/key_value/generate.cpp b/tests/slo_workloads/key_value/generate.cpp index 1486c9649aa..4e6b8bb3d9f 100644 --- a/tests/slo_workloads/key_value/generate.cpp +++ b/tests/slo_workloads/key_value/generate.cpp @@ -8,7 +8,7 @@ using namespace NYdb::NTable; TGenerateInitialContentJob::TGenerateInitialContentJob(const TCreateOptions& createOpts, std::uint32_t maxId) - : TThreadJob(createOpts.CommonOptions) + : TThreadJob(createOpts.CommonOptions, "generate") , Executor(createOpts.CommonOptions, Stats, TExecutor::ModeBlocking) , PackGenerator( createOpts.CommonOptions @@ -81,5 +81,4 @@ UPSERT INTO `%s` SELECT * FROM AS_TABLE($items); void TGenerateInitialContentJob::OnFinish() { Executor.Finish(); Executor.Wait(); - Stats.Flush(); } diff --git a/tests/slo_workloads/key_value/key_value.cpp b/tests/slo_workloads/key_value/key_value.cpp index 9f48597019f..1cb40f3620f 100644 --- a/tests/slo_workloads/key_value/key_value.cpp +++ b/tests/slo_workloads/key_value/key_value.cpp @@ -70,12 +70,12 @@ int DoRun(TDatabaseOptions& dbOptions, int argc, char** argv) { if (!runOptions.DontRunA) { runOptions.CommonOptions.Rps = runOptions.Read_rps; - runOptions.CommonOptions.ReactionTime = TDuration::MilliSeconds(runOptions.CommonOptions.A_ReactionTime); + runOptions.CommonOptions.ReactionTime = runOptions.ReadTimeout; jobs->Add(new TReadJob(runOptions.CommonOptions, maxId)); } if (!runOptions.DontRunB) { runOptions.CommonOptions.Rps = runOptions.Write_rps; - runOptions.CommonOptions.ReactionTime = DefaultReactionTime; + runOptions.CommonOptions.ReactionTime = runOptions.WriteTimeout; jobs->Add(new TWriteJob(runOptions.CommonOptions, maxId)); } diff --git a/tests/slo_workloads/key_value/key_value.h b/tests/slo_workloads/key_value/key_value.h index 10c0e876925..4238f01c735 100644 --- a/tests/slo_workloads/key_value/key_value.h +++ b/tests/slo_workloads/key_value/key_value.h @@ -48,7 +48,6 @@ class TReadJob : public TThreadJob { private: std::unique_ptr Executor; std::uint32_t ObjectIdRange; - bool SaveResult; }; int CreateTable(TDatabaseOptions& dbOptions); diff --git a/tests/slo_workloads/key_value/run.cpp b/tests/slo_workloads/key_value/run.cpp index fa2393c6b95..671947ccf01 100644 --- a/tests/slo_workloads/key_value/run.cpp +++ b/tests/slo_workloads/key_value/run.cpp @@ -7,7 +7,7 @@ using namespace NYdb; using namespace NYdb::NTable; TWriteJob::TWriteJob(const TCommonOptions& opts, std::uint32_t maxId) - : TThreadJob(opts) + : TThreadJob(opts, "write") , Executor(opts, Stats, TExecutor::ModeNonBlocking) , Generator(opts, maxId) {} @@ -77,15 +77,13 @@ UPSERT INTO `%s` SELECT * FROM AS_TABLE($items); void TWriteJob::OnFinish() { Executor.Finish(); Executor.Wait(); - Stats.Flush(); } // Implementation of TReadJob TReadJob::TReadJob(const TCommonOptions& opts, std::uint32_t maxId) - : TThreadJob(opts) - , Executor(opts.RetryMode ? new TExecutorWithRetry(opts, Stats) : new TExecutor(opts, Stats)) + : TThreadJob(opts, "read") + , Executor(std::make_unique(opts, Stats)) , ObjectIdRange(static_cast(maxId * 1.25)) // 20% of requests with no result - , SaveResult(opts.SaveResult) {} void TReadJob::ShowProgress(TStringBuilder& report) { @@ -154,8 +152,4 @@ void TReadJob::OnFinish() { if (infly) { Cerr << "Warning: thread A finished while having " << infly << " infly requests." << Endl; } - Stats.Flush(); - if (SaveResult) { - Stats.SaveResult(); - } } diff --git a/tests/slo_workloads/utils/executor.cpp b/tests/slo_workloads/utils/executor.cpp index 5e45cef31ca..28add35c98e 100644 --- a/tests/slo_workloads/utils/executor.cpp +++ b/tests/slo_workloads/utils/executor.cpp @@ -274,13 +274,17 @@ bool TExecutor::Execute(const NYdb::NTable::TTableClient::TOperationFunc& func) } } - TStatUnit stat = Stats.CreateStatUnit(); + auto stat = Stats.StartRequest(); + + auto future = InsistentClient.ExecuteWithRetry([func, stat](NYdb::NTable::TSession session) { + auto result = func(session); + return result; + }); - auto future = InsistentClient.ExecuteWithRetry(func); future.Subscribe([this, stat, SemaphoreWrapper](const TAsyncFinalStatus& future) mutable { Y_ABORT_UNLESS(future.HasValue()); TFinalStatus resultStatus = future.GetValue(); - Stats.Report(stat, resultStatus); + Stats.FinishRequest(stat, resultStatus); if (resultStatus) { CheckForError(*resultStatus); } @@ -351,7 +355,6 @@ bool TExecutor::IsStopped() { } void TExecutor::Finish() { - // Stats.UpdateSessionStats(InsistentClient.GetSessionStats()); with_lock(Lock) { if (!AllJobsLaunched) { AllJobsLaunched = true; @@ -361,9 +364,6 @@ void TExecutor::Finish() { } void TExecutor::UpdateStats() { - if (Infly > MaxSecInfly) { - MaxSecInfly = Infly; - } std::uint64_t activeSessions = InsistentClient.GetActiveSessions(); if (activeSessions > MaxSecSessions) { MaxSecSessions = activeSessions; @@ -382,10 +382,10 @@ void TExecutor::UpdateStats() { } void TExecutor::ReportStats() { + Stats.ReportStats(MaxSecSessions, MaxSecReadPromises, MaxSecExecutorPromises); TInstant now = TInstant::Now(); if (now.Seconds() > LastReportSec) { - Stats.ReportStats(MaxSecInfly, MaxSecSessions, MaxSecReadPromises, MaxSecExecutorPromises); - MaxSecInfly = 0; + Stats.ReportStats(MaxSecSessions, MaxSecReadPromises, MaxSecExecutorPromises); MaxSecSessions = 0; MaxSecReadPromises = 0; MaxSecExecutorPromises = 0; @@ -442,93 +442,3 @@ void TExecutor::Report(TStringBuilder& out) const { } } } - - -TExecutorWithRetry::TExecutorWithRetry(const TCommonOptions& opts, TStat& stats) - : TExecutor(opts, stats) -{} - -bool TExecutorWithRetry::Execute(const NYdb::NTable::TTableClient::TOperationFunc& func) { - auto threadFunc = [this, func]() { - if (IsStopped()) { - DecrementWaiting(); - return; - } - - with_lock(Lock) { - --Waiting; - if (Infly < Opts.MaxInfly) { - ++Infly; - if (Infly > MaxInfly) { - MaxInfly = Infly; - } - UpdateStats(); - } else { - Stats.ReportMaxInfly(); - UpdateStats(); - return; - } - } - - std::shared_ptr context = std::make_shared(Stats); - - auto executeOperation = [this, func]() { - return InsistentClient.ExecuteWithRetry(func); - }; - - context->HandleStatusFunc = std::make_unique>( - [this, executeOperation, context](const TAsyncFinalStatus& future) mutable { - Y_ABORT_UNLESS(future.HasValue()); - TFinalStatus resultStatus = future.GetValue(); - if (resultStatus) { - // Reply received - CheckForError(*resultStatus); - if (resultStatus->IsSuccess()) { - //Ok received - Stats.Report(context->LifeTimeStat, resultStatus->GetStatus()); - DecrementInfly(); - context->HandleStatusFunc.reset(); - return; - } - } - if (IsStopped() || TInstant::Now() - context->LifeTimeStat.Start > GlobalTimeout) { - // Application stopped working or global timeout reached. Ok reply hasn't received yet - Stats.Report(context->LifeTimeStat, TInnerStatus::StatusNotFinished); - DecrementInfly(); - context->HandleStatusFunc.reset(); - return; - } - Stats.Report(context->PerRequestStat, resultStatus); - context->PerRequestStat = Stats.CreateStatUnit(); - // Retrying: - executeOperation().Subscribe(*context->HandleStatusFunc); - }); - - context->Retries.fetch_add(1); - Y_ABORT_UNLESS(context->Retries.load() < 500, "Too much retries"); - - executeOperation().Subscribe(*context->HandleStatusFunc); - }; - - if (IsStopped()) { - return false; - } - - bool CanLaunchJob = false; - - with_lock(Lock) { - if (!AllJobsLaunched) { - CanLaunchJob = true; - ++Waiting; - } - } - - if (CanLaunchJob) { - if (!InputQueue->AddFunc(threadFunc)) { - DecrementWaiting(); - } - } - ++InProgressCount; - InProgressSum += InputQueue->Size(); - return true; -} diff --git a/tests/slo_workloads/utils/executor.h b/tests/slo_workloads/utils/executor.h index 5dddb336b6d..69ed261a352 100644 --- a/tests/slo_workloads/utils/executor.h +++ b/tests/slo_workloads/utils/executor.h @@ -1,5 +1,6 @@ #pragma once +#include "statistics.h" #include "utils.h" #include @@ -56,7 +57,7 @@ class TInsistentClient { bool Valid = true; }; - struct TOperationContext : public TThrRefBase { + struct TOperationContext { bool Finished = false; TAdaptiveLock Lock; TCheckedIterator RetryIter; @@ -131,7 +132,6 @@ class TExecutor { std::uint32_t Wait(TDuration waitTimeout); bool IsStopped(); void Finish(); - std::uint32_t GetTotal() const; void Report(TStringBuilder& out) const; protected: @@ -167,8 +167,6 @@ class TExecutor { TInstant Deadline; // Last second we reported Infly std::uint64_t LastReportSec = 0; - // Max infly for current second - std::uint64_t MaxSecInfly = 0; // Max Active sessions for current second std::uint64_t MaxSecSessions = 0; @@ -178,21 +176,3 @@ class TExecutor { std::uint64_t MaxSecReadPromises = 0; std::uint64_t MaxSecExecutorPromises = 0; }; - -class TExecutorWithRetry : public TExecutor { -public: - struct TRetryContext { - TRetryContext(TStat& stat) - : LifeTimeStat(stat.CreateStatUnit()) - , PerRequestStat(stat.CreateStatUnit()) - {} - - TStatUnit LifeTimeStat; - TStatUnit PerRequestStat; - std::unique_ptr> HandleStatusFunc; - std::atomic Retries = 0; - }; - - TExecutorWithRetry(const TCommonOptions& opts, TStat& stats); - bool Execute(const NYdb::NTable::TTableClient::TOperationFunc& func) override; -}; diff --git a/tests/slo_workloads/utils/job.cpp b/tests/slo_workloads/utils/job.cpp index 23421bcb747..cc627b2b707 100644 --- a/tests/slo_workloads/utils/job.cpp +++ b/tests/slo_workloads/utils/job.cpp @@ -5,18 +5,12 @@ const std::string LogFileName = "benchmark.log"; -TThreadJob::TThreadJob(const TCommonOptions& opts) +TThreadJob::TThreadJob(const TCommonOptions& opts, const std::string& operationType) : RpsProvider(opts.Rps) , Prefix(opts.DatabaseOptions.Prefix) - , Stats( - opts.ReactionTime, - opts.ResultFileName, - !opts.DontPushMetrics, - opts.RetryMode - ) + , Stats(opts.DontPushMetrics ? std::nullopt : std::make_optional(opts.MetricsPushUrl), operationType) , StopOnError(opts.StopOnError) , MaxDelay(opts.ReactionTime) - , UseFollowers(opts.UseFollowers) { } @@ -28,7 +22,7 @@ void TThreadJob::Start(TInstant deadline) { void TThreadJob::StartThread() { auto threadFunc = [this]() { - Stats.Reset(); + Stats.Start(); RpsProvider.Reset(); DoJob(); Stats.Finish(); diff --git a/tests/slo_workloads/utils/job.h b/tests/slo_workloads/utils/job.h index 7e6e48bede5..a57a790b432 100644 --- a/tests/slo_workloads/utils/job.h +++ b/tests/slo_workloads/utils/job.h @@ -1,5 +1,6 @@ #pragma once +#include "statistics.h" #include "utils.h" #include @@ -7,7 +8,7 @@ class TThreadJob { public: - TThreadJob(const TCommonOptions& opts); + TThreadJob(const TCommonOptions& opts, const std::string& operationType); virtual ~TThreadJob() = default; virtual void Start(TInstant deadline); @@ -31,10 +32,9 @@ class TThreadJob { TStat Stats; bool StopOnError; TDuration MaxDelay; - bool UseFollowers; }; -class TJobContainer : public TThrRefBase { +class TJobContainer { public: void Add(TThreadJob* job); void Start(TInstant deadline = TInstant()); diff --git a/tests/slo_workloads/utils/metrics.cpp b/tests/slo_workloads/utils/metrics.cpp new file mode 100644 index 00000000000..48476712ff6 --- /dev/null +++ b/tests/slo_workloads/utils/metrics.cpp @@ -0,0 +1,150 @@ +#include "metrics.h" + +#include "utils.h" + +#include + +#include +#include + +#include + +using namespace std::chrono_literals; + +class TOtelMetricsPusher : public IMetricsPusher { +public: + TOtelMetricsPusher(const std::string& metricsPushUrl, const std::string& operationType) + : OperationType_(operationType) + { + auto exporterOptions = opentelemetry::exporter::otlp::OtlpHttpMetricExporterOptions(); + exporterOptions.url = metricsPushUrl; + + auto exporter = opentelemetry::exporter::otlp::OtlpHttpMetricExporterFactory::Create(exporterOptions); + + opentelemetry::sdk::metrics::PeriodicExportingMetricReaderOptions readerOptions; + readerOptions.export_interval_millis = 1000ms; + readerOptions.export_timeout_millis = 500ms; + + auto metricReader = opentelemetry::sdk::metrics::PeriodicExportingMetricReaderFactory::Create(std::move(exporter), readerOptions); + + MeterProvider_ = opentelemetry::sdk::metrics::MeterProviderFactory::Create(); + MeterProvider_->AddMetricReader(std::move(metricReader)); + + Meter_ = MeterProvider_->GetMeter("slo_workloads", NYdb::GetSdkSemver()); + + InitMetrics(); + } + + void PushRequestData(const TRequestData& requestData) override { + if (requestData.Status == NYdb::EStatus::SUCCESS) { + OperationsSuccessTotal_->Add(1, {{"operation_type", OperationType_}}); + } else { + ErrorsTotal_->Add(1, {{"status", YdbStatusToString(requestData.Status)}}); + OperationsFailureTotal_->Add(1, {{"operation_type", OperationType_}}); + } + OperationsTotal_->Add(1, {{"operation_type", OperationType_}}); + OperationLatencySeconds_->Record(requestData.Delay.SecondsFloat(), {{"operation_type", OperationType_}, {"status", YdbStatusToString(requestData.Status)}}); + RetryAttempts_->Record(requestData.RetryAttempts, {{"operation_type", OperationType_}}); + } + +private: + void InitMetrics() { + ErrorsTotal_ = Meter_->CreateUInt64Counter("sdk_errors_total", + "Total number of errors encountered, categorized by error type." + ); + + OperationsTotal_ = Meter_->CreateUInt64Counter("sdk_operations_total", + "Total number of operations, categorized by type attempted by the SDK." + ); + + OperationsSuccessTotal_ = Meter_->CreateUInt64Counter("sdk_operations_success_total", + "Total number of successful operations, categorized by type." + ); + + OperationsFailureTotal_ = Meter_->CreateUInt64Counter("sdk_operations_failure_total", + "Total number of failed operations, categorized by type." + ); + + OperationLatencySeconds_ = CreateDoubleHistogram("sdk_operation_latency_seconds", + "Latency of operations performed by the SDK in seconds, categorized by type and status.", + { + 0.001, // 1 ms + 0.002, // 2 ms + 0.003, // 3 ms + 0.004, // 4 ms + 0.005, // 5 ms + 0.0075, // 7.5 ms + 0.010, // 10 ms + 0.020, // 20 ms + 0.050, // 50 ms + 0.100, // 100 ms + 0.200, // 200 ms + 0.500, // 500 ms + 1.000, // 1 s + }, + "s" + ); + + RetryAttempts_ = Meter_->CreateInt64Gauge("sdk_retry_attempts", + "Current retry attempts, categorized by operation type." + ); + } + + std::unique_ptr> CreateDoubleHistogram( + const std::string& name, + const std::string& description, + const std::vector& buckets, + const std::string& unit = {}) + { + auto selector = std::make_unique( + opentelemetry::sdk::metrics::InstrumentType::kHistogram, + name, + unit + ); + + auto meterSelector = std::make_unique( + "slo_workloads", + NYdb::GetSdkSemver(), + "" + ); + + auto histogramConfig = std::make_shared(); + histogramConfig->boundaries_ = buckets; + + auto view = std::make_unique( + "", + "", + opentelemetry::sdk::metrics::AggregationType::kHistogram, + histogramConfig + ); + + MeterProvider_->AddView(std::move(selector), std::move(meterSelector), std::move(view)); + + return Meter_->CreateDoubleHistogram(name, description, unit); + } + + std::string OperationType_; + + std::unique_ptr MeterProvider_; + std::shared_ptr Meter_; + + std::unique_ptr> ErrorsTotal_; + std::unique_ptr> OperationsTotal_; + std::unique_ptr> OperationsSuccessTotal_; + std::unique_ptr> OperationsFailureTotal_; + std::unique_ptr> OperationLatencySeconds_; + std::unique_ptr> RetryAttempts_; +}; + +class TNoopMetricsPusher : public IMetricsPusher { +public: + void PushRequestData([[maybe_unused]] const TRequestData& requestData) override {} +}; + +std::unique_ptr CreateOtelMetricsPusher(const std::string& metricsPushUrl, const std::string& operationType) { + return std::make_unique(metricsPushUrl, operationType); +} + +std::unique_ptr CreateNoopMetricsPusher() { + return std::make_unique(); +} diff --git a/tests/slo_workloads/utils/metrics.h b/tests/slo_workloads/utils/metrics.h new file mode 100644 index 00000000000..cf6aac7be43 --- /dev/null +++ b/tests/slo_workloads/utils/metrics.h @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include + + +struct TRequestData { + TDuration Delay; + NYdb::EStatus Status; + std::uint64_t RetryAttempts; +}; + +class IMetricsPusher { +public: + virtual ~IMetricsPusher() = default; + + virtual void PushRequestData(const TRequestData& requestData) = 0; +}; + +std::unique_ptr CreateOtelMetricsPusher(const std::string& metricsPushUrl, const std::string& operationType); +std::unique_ptr CreateNoopMetricsPusher(); diff --git a/tests/slo_workloads/utils/statistics.cpp b/tests/slo_workloads/utils/statistics.cpp index 4b5d8e4855f..15789e76206 100644 --- a/tests/slo_workloads/utils/statistics.cpp +++ b/tests/slo_workloads/utils/statistics.cpp @@ -1,114 +1,20 @@ #include "statistics.h" -#include -#include +#include "metrics.h" +#include "utils.h" -#include - - -TStatUnit::TStatUnit(const std::shared_ptr& periodData, TInstant startTime) - : PeriodData(periodData) - , Start(startTime) -{ -} - -void TStatUnit::Report(const TInnerStatus& status) { - PeriodData->AddStat(Delay(), status); -} - -TDuration TStatUnit::Delay() const { - return End - Start; -} - -TPeriodData::TPeriodData( - TStat* stats, - bool retryMode, - const TDuration maxDelay, - std::uint64_t currentSecond, - std::uint64_t currInfly, - std::uint64_t currSessions, - std::uint64_t currPromises, - std::uint64_t currExecutorPromises -) - : Stats(stats) - , RetryMode(retryMode) - , CurrentSecond(currentSecond) - , MaxDelay(maxDelay) -{ - Counters.Infly = currInfly; - Counters.ActiveSessions = currSessions; - Counters.ReadPromises = currPromises; - Counters.ExecutorPromises = currExecutorPromises; -} - -TPeriodData::~TPeriodData() { - Stats->ReportLatencyData(CurrentSecond, std::move(Counters), std::move(Replies), std::move(OkDelays), NotOkDelays); -} - -TStatUnit TPeriodData::CreateStatUnit(TInstant startTime) { - return TStatUnit(shared_from_this(), startTime); -} - -std::uint64_t TPeriodData::GetCurrentSecond() const { - return CurrentSecond; -} - -void TPeriodData::ReportMaxInfly() { - ++Replies.CountMaxInfly; -} - -void TPeriodData::ReportInfly(std::uint64_t infly) { - if (infly > Counters.Infly) { - Counters.Infly = infly; - } -} - -void TPeriodData::ReportActiveSessions(std::uint64_t sessions) { - if (sessions > Counters.ActiveSessions) { - Counters.ActiveSessions = sessions; - } -} - -// Debug use only: -void TPeriodData::ReportReadPromises(std::uint64_t promises) { - if (promises > Counters.ReadPromises) { - Counters.ReadPromises = promises; - } -} -void TPeriodData::ReportExecutorPromises(std::uint64_t promises) { - if (promises > Counters.ExecutorPromises) { - Counters.ExecutorPromises = promises; - } -} - -void TPeriodData::AddStat(TDuration delay, const TInnerStatus& status) { - if (status.InnerStatus == TInnerStatus::StatusReceived && status.YdbStatus == NYdb::EStatus::SUCCESS) { - OkDelays.push_back(delay); - } else { - NotOkDelays.push_back(delay); - } - - switch (status.InnerStatus) { - case TInnerStatus::StatusReceived: - if (status.YdbStatus != NYdb::EStatus::SUCCESS || !RetryMode || delay <= MaxDelay) { - ++Replies.Statuses[status.YdbStatus]; - } else { - ++Replies.CountHighLatency; - } - break; - case TInnerStatus::StatusApplicationTimeout: - ++Replies.ApplicationTimeout; - break; - case TInnerStatus::StatusNotFinished: - ++Replies.NotFinished; - break; - default: - Y_ABORT_UNLESS(status.InnerStatus); - break; - } -} namespace { + // Calculated percentiles for a period of time + struct TPercentile { + TDuration P50; + TDuration P90; + TDuration P95; + TDuration P99; + TDuration P99_9; + TDuration P100; + }; + void CalculatePercentiles(TPercentile& p, std::vector& delays) { size_t count = delays.size(); if (count) { @@ -123,41 +29,14 @@ namespace { } } -TStat::TStat( - TDuration maxDelay, - const std::string& resultFileName, - bool pushMetrics, - bool retryMode -) - : MaxDelay(maxDelay) - , StartTime(TInstant::Now()) - , PushMetrics(pushMetrics) - , RetryMode(retryMode) - , ResultFileName(resultFileName) +TStat::TStat(const std::optional& metricsPushUrl, const std::string& operationType) + : StartTime(TInstant::Now()) + , MetricsPusher(metricsPushUrl ? CreateOtelMetricsPusher(*metricsPushUrl, operationType) : CreateNoopMetricsPusher()) { MetricsPushQueue.Start(20); } -TStat::~TStat() { - Flush(); -} - -void TStat::Flush() { - std::uint64_t lastSecMeasured = TInstant::Now().Seconds(); - if (ActivePeriod) { - lastSecMeasured = ActivePeriod->GetCurrentSecond(); - ActivePeriod.reset(); - } - if (PushMetrics) { - ResetMetricsPusher(lastSecMeasured + 1); - MetricsPushQueue.Stop(); - } -} - -void TStat::Reset() { - if (PushMetrics) { - ResetMetricsPusher(TInstant::Now().Seconds() - 1); - } +void TStat::Start() { StartTime = TInstant::Now(); } @@ -165,356 +44,104 @@ void TStat::Finish() { FinishTime = TInstant::Now(); } -TStatUnit TStat::CreateStatUnit() { +std::shared_ptr TStat::StartRequest() { std::lock_guard lock(Mutex); - TInstant Now = TInstant::Now(); - CheckCurrentSecond(Now); - return ActivePeriod->CreateStatUnit(Now); -} - -void TStat::Report(TStatUnit& unit, const TInnerStatus& status) { - unit.End = TInstant::Now(); - OnReport(unit, status); -} - -void TStat::Report(TStatUnit& unit, TInnerStatus::EInnerStatus innerStatus) { - Report(unit, TInnerStatus(innerStatus)); -} - -void TStat::Report(TStatUnit& unit, const TFinalStatus& status) { - Report( - unit, - status - ? TInnerStatus(TInnerStatus::StatusReceived, status->GetStatus()) - : TInnerStatus(TInnerStatus::StatusApplicationTimeout) - ); + ++Infly; + return std::make_shared(TInstant::Now()); } -void TStat::Report(TStatUnit& unit, NYdb::EStatus status) { - Report(unit, TInnerStatus(TInnerStatus::StatusReceived, status)); -} - -void TStat::ReportMaxInfly() { +void TStat::FinishRequest(const std::shared_ptr& unit, const TFinalStatus& status) { std::lock_guard lock(Mutex); - ++Replies.CountMaxInfly; - TInstant Now = TInstant::Now(); - CheckCurrentSecond(Now); - ActivePeriod->ReportMaxInfly(); -} + unit->End = TInstant::Now(); -void TStat::ReportStats(std::uint64_t infly, std::uint64_t sessions, std::uint64_t readPromises, std::uint64_t executorPromises) { - std::lock_guard lock(Mutex); + auto delay = unit->End - unit->Start; - Counters.Infly = infly; - Counters.ActiveSessions = sessions; - Counters.ReadPromises = readPromises; - Counters.ExecutorPromises = executorPromises; - TInstant Now = TInstant::Now(); - CheckCurrentSecond(Now); - ActivePeriod->ReportInfly(infly); - ActivePeriod->ReportActiveSessions(sessions); - ActivePeriod->ReportReadPromises(readPromises); - ActivePeriod->ReportExecutorPromises(executorPromises); -} + --Infly; -void TStat::Reserve(size_t size) { - std::lock_guard lock(Mutex); + if (status) { + ++Statuses[status->GetStatus()]; + } else { + ++ApplicationTimeout; + } + + if (status && status->GetStatus() == NYdb::EStatus::SUCCESS) { + OkDelays.push_back(delay); + } - LatencyStats.reserve(size); - LatencyData.reserve(size); + ScheduleMetricsPush([this, delay, status, unit]() { + MetricsPusher->PushRequestData({ + .Delay = delay, + .Status = status->GetStatus(), + .RetryAttempts = unit->RetryAttempts + }); + }); } -void TStat::ReportLatencyData( - std::uint64_t currSecond, - TCounters&& counters, - TReplies&& replies, - std::vector&& oks, - std::vector& notOks -) { +void TStat::ReportMaxInfly() { std::lock_guard lock(Mutex); - LatencyStats.emplace_back(); - TPeriodStat& p = LatencyStats.back(); - p.Seconds = currSecond; - std::swap(p.Counters, counters); - std::swap(p.Replies, replies); - CalculatePercentiles(p.Oks, oks); - CalculatePercentiles(p.NotOks, notOks); - LatencyData.emplace_back(std::move(oks)); - PushMetricsData(p); + ++CountMaxInfly; } -void TStat::UpdateSessionStats( - const std::unordered_map& sessionStats -) { +void TStat::ReportStats(std::uint64_t sessions, std::uint64_t readPromises, std::uint64_t executorPromises) { std::lock_guard lock(Mutex); - SessionStats = sessionStats; + ActiveSessions = sessions; + ReadPromises = readPromises; + ExecutorPromises = executorPromises; } void TStat::PrintStatistics(TStringBuilder& out) { std::lock_guard lock(Mutex); - TInstant now = TInstant::Now(); - CheckCurrentSecond(now); - std::uint64_t total = GetTotal(); + std::uint64_t total = CountMaxInfly; + for (const auto& [status, counter] : Statuses) { + total += counter; + } + total += ApplicationTimeout; TDuration timePassed; if (FinishTime < StartTime) { // If we ask for current progress - timePassed = now - StartTime; + timePassed = TInstant::Now() - StartTime; } else { timePassed = FinishTime - StartTime; } std::uint64_t rps = total * 1000000 / timePassed.MicroSeconds(); out << total << " requests total" << Endl - << Replies.Statuses[NYdb::EStatus::SUCCESS] << " succeeded"; + << Statuses[NYdb::EStatus::SUCCESS] << " succeeded"; if (total) { - out << " (" << Replies.Statuses[NYdb::EStatus::SUCCESS] * 100 / total << "%)"; + out << " (" << Statuses[NYdb::EStatus::SUCCESS] * 100 / total << "%)"; } - for (const auto&[status, counter] : Replies.Statuses) { + for (const auto&[status, counter] : Statuses) { out << Endl << counter << " replies with status " << YdbStatusToString(status) << Endl; } - out << Endl << Replies.CountMaxInfly << " failed due to max infly" << Endl - << Replies.CountHighLatency << " OK results exceeded latency limit of " << MaxDelay << Endl - << Replies.ApplicationTimeout << " application timeouts" << Endl - << Replies.NotFinished << " requests not finished within program lifetime" << Endl + out << Endl << CountMaxInfly << " failed due to max infly" << Endl + << ApplicationTimeout << " application timeouts" << Endl << "Time passed: " << timePassed.ToString() << Endl << "Real rps: " << rps << Endl; - if (LatencyData.size()) { - CalculateGlobalPercentile(); - TPercentile& p = *GlobalPercentile; - out << "Global latency percentiles (" << LatencyData.size() << " seconds measured):" << Endl + if (OkDelays.size()) { + TPercentile p; + CalculatePercentiles(p, OkDelays); + + out << "Global latency percentiles:" << Endl << "P50: " << p.P50 << "\tP90: " << p.P90 << "\tP95: " << p.P95 << "\tP99: " << p.P99 << "\tP99.9: " << p.P99_9 << "\tP100: " << p.P100 << Endl; - CalculateFailSeconds(); - out << *FailSeconds << " seconds where p99 reached max delay of " << MaxDelay << Endl; } else { out << "Can't calculate latency percentiles: No data (zero requests measured)" << Endl; } } -void TStat::SaveResult() { - std::lock_guard lock(Mutex); - - if (LatencyData.size()) { - CalculateGlobalPercentile(); - CalculateFailSeconds(); - NJson::TJsonValue root; - root["Oks"] = Replies.Statuses[NYdb::EStatus::SUCCESS]; - root["Total"] = GetTotal(); - root["P99"] = GlobalPercentile->P99.MilliSeconds(); - root["FailSeconds"] = *FailSeconds; - root["StartTime"] = StartTime.Seconds(); - root["FinishTime"] = FinishTime.Seconds(); - NJson::TJsonValue& items = root["SessionCountsAtFinish"]; - items.SetType(NJson::JSON_ARRAY); - for (const auto& s : SessionStats) { - items.AppendValue(s.second); - } - TFileOutput resultFile(ResultFileName); - NJsonWriter::TBuf buf; - buf.WriteJsonValue(&root); - resultFile << buf.Str(); - Cout << "Result saved to file " << ResultFileName << Endl; - } -} - -void TStat::CalculateGlobalPercentile() { - if (GlobalPercentile) { - return; - } - std::vector fullData; - size_t totalSize = 0; - for (auto& periodData : LatencyData) { - totalSize += periodData.size(); - } - fullData.reserve(totalSize); - for (auto& periodData : LatencyData) { - fullData.insert(fullData.end(), periodData.begin(), periodData.end()); - } - GlobalPercentile = std::make_unique(); - CalculatePercentiles(*GlobalPercentile, fullData); -} - -namespace { - bool IsGoodInterval(const TPeriodStat& stat, const TDuration& maxDelay) { - for (const auto& [status, counter] : stat.Replies.Statuses) { - if (status != NYdb::EStatus::SUCCESS && counter) { - return false; - } - } - return stat.Oks.P99 <= maxDelay && !stat.Replies.CountMaxInfly && !stat.Replies.ApplicationTimeout; - } -} - -void TStat::CalculateFailSeconds() { - if (FailSeconds) { - return; - } - std::sort(LatencyStats.begin(), LatencyStats.end(), [&](const TPeriodStat& a, const TPeriodStat& b) { - return a.Seconds < b.Seconds; - }); - FailSeconds = std::make_unique(0); - std::uint64_t& failSeconds = *FailSeconds; - std::uint64_t lastSecChecked = LatencyStats[0].Seconds - 1; - for (auto& stat : LatencyStats) { - failSeconds += stat.Seconds - lastSecChecked - 1; - lastSecChecked = stat.Seconds; - if (!IsGoodInterval(stat, MaxDelay)) { - ++failSeconds; - } - } -} - -std::uint64_t TStat::GetTotal() { - std::uint64_t total = Replies.CountMaxInfly + Replies.CountHighLatency + Replies.NotFinished; - if (!RetryMode) { - for (const auto& [status, counter] : Replies.Statuses) { - total += counter; - } - total += Replies.ApplicationTimeout; - } else { - total += Replies.Statuses[NYdb::EStatus::SUCCESS]; - } - return total; -} - TInstant TStat::GetStartTime() const { return StartTime; } -void TStat::OnReport(TStatUnit& unit, const TInnerStatus& status) { - std::lock_guard lock(Mutex); - - TDuration delay = unit.Delay(); - switch (status.InnerStatus) { - case TInnerStatus::StatusReceived: - if (status.YdbStatus != NYdb::EStatus::SUCCESS || !RetryMode || delay <= MaxDelay) { - ++Replies.Statuses[status.YdbStatus]; - } else { - ++Replies.CountHighLatency; - } - break; - case TInnerStatus::StatusApplicationTimeout: - ++Replies.ApplicationTimeout; - break; - case TInnerStatus::StatusNotFinished: - ++Replies.NotFinished; - break; - default: - Y_ABORT_UNLESS(status.InnerStatus); - break; - } - // Saving delay stats - unit.Report(status); -} - -void TStat::CheckCurrentSecond(TInstant now) { - std::uint64_t currSecond = now.Seconds(); - if (currSecond > CurrentSecond) { - CurrentSecond = currSecond; - ActivePeriod = std::make_shared( - this, - RetryMode, - MaxDelay, - currSecond, - Counters.Infly, - Counters.ActiveSessions, - Counters.ReadPromises, - Counters.ExecutorPromises - ); - } -} - -void TStat::PushMetricsData(const TPeriodStat& p) { - if (!PushMetrics) { - return; - } - auto threadFunc = [this, p]() { - MetricsPusher->PushData(p); - }; - if (!MetricsPushQueue.AddFunc(threadFunc)) { - Cerr << TInstant::Now().ToRfc822StringLocal() << ": Failed to push data to solomon" << Endl; - } -} - -void TStat::ResetMetricsPusher(std::uint64_t timestamp) { - while (timestamp >= TInstant::Now().Seconds()) { - Sleep(TDuration::Seconds(1)); - } - TPeriodStat pStat; - pStat.Seconds = timestamp; - PushMetricsData(pStat); -} - -std::string YdbStatusToString(NYdb::EStatus status) { - switch (status) { - case NYdb::EStatus::SUCCESS: - return "SUCCESS"; - case NYdb::EStatus::BAD_REQUEST: - return "BAD_REQUEST"; - case NYdb::EStatus::UNAUTHORIZED: - return "UNAUTHORIZED"; - case NYdb::EStatus::INTERNAL_ERROR: - return "INTERNAL_ERROR"; - case NYdb::EStatus::ABORTED: - return "ABORTED"; - case NYdb::EStatus::UNAVAILABLE: - return "UNAVAILABLE"; - case NYdb::EStatus::OVERLOADED: - return "OVERLOADED"; - case NYdb::EStatus::SCHEME_ERROR: - return "SCHEME_ERROR"; - case NYdb::EStatus::GENERIC_ERROR: - return "GENERIC_ERROR"; - case NYdb::EStatus::TIMEOUT: - return "TIMEOUT"; - case NYdb::EStatus::BAD_SESSION: - return "BAD_SESSION"; - case NYdb::EStatus::PRECONDITION_FAILED: - return "PRECONDITION_FAILED"; - case NYdb::EStatus::ALREADY_EXISTS: - return "ALREADY_EXISTS"; - case NYdb::EStatus::NOT_FOUND: - return "NOT_FOUND"; - case NYdb::EStatus::SESSION_EXPIRED: - return "SESSION_EXPIRED"; - case NYdb::EStatus::CANCELLED: - return "CANCELLED"; - case NYdb::EStatus::UNDETERMINED: - return "UNDETERMINED"; - case NYdb::EStatus::UNSUPPORTED: - return "UNSUPPORTED"; - case NYdb::EStatus::SESSION_BUSY: - return "SESSION_BUSY"; - case NYdb::EStatus::EXTERNAL_ERROR: - return "EXTERNAL_ERROR"; - case NYdb::EStatus::STATUS_UNDEFINED: - return "STATUS_UNDEFINED"; - case NYdb::EStatus::TRANSPORT_UNAVAILABLE: - return "TRANSPORT_UNAVAILABLE"; - case NYdb::EStatus::CLIENT_RESOURCE_EXHAUSTED: - return "CLIENT_RESOURCE_EXHAUSTED"; - case NYdb::EStatus::CLIENT_DEADLINE_EXCEEDED: - return "CLIENT_DEADLINE_EXCEEDED"; - case NYdb::EStatus::CLIENT_INTERNAL_ERROR: - return "CLIENT_INTERNAL_ERROR"; - case NYdb::EStatus::CLIENT_CANCELLED: - return "CLIENT_CANCELLED"; - case NYdb::EStatus::CLIENT_UNAUTHENTICATED: - return "CLIENT_UNAUTHENTICATED"; - case NYdb::EStatus::CLIENT_CALL_UNIMPLEMENTED: - return "CLIENT_CALL_UNIMPLEMENTED"; - case NYdb::EStatus::CLIENT_OUT_OF_RANGE: - return "CLIENT_OUT_OF_RANGE"; - case NYdb::EStatus::CLIENT_DISCOVERY_FAILED: - return "CLIENT_DISCOVERY_FAILED"; - case NYdb::EStatus::CLIENT_LIMITS_REACHED: - return "CLIENT_LIMITS_REACHED"; +void TStat::ScheduleMetricsPush(std::function func) { + if (!MetricsPushQueue.AddFunc(func)) { + Cerr << TInstant::Now().ToRfc822StringLocal() << ": Failed to push metrics" << Endl; } } diff --git a/tests/slo_workloads/utils/statistics.h b/tests/slo_workloads/utils/statistics.h index e3c08ee6109..9e7a092430f 100644 --- a/tests/slo_workloads/utils/statistics.h +++ b/tests/slo_workloads/utils/statistics.h @@ -1,5 +1,7 @@ #pragma once +#include "metrics.h" + #include #include @@ -20,181 +22,67 @@ inline double GetMillisecondsDouble(const TDuration& d) { using TFinalStatus = std::optional; using TAsyncFinalStatus = NThreading::TFuture; -struct TInnerStatus { - enum EInnerStatus { - StatusUnknown, - StatusReceived, - StatusApplicationTimeout, - StatusNotFinished - }; - - TInnerStatus(EInnerStatus innerStatus = StatusUnknown, NYdb::EStatus ydbStatus = NYdb::EStatus::STATUS_UNDEFINED) - : InnerStatus(innerStatus) - , YdbStatus(ydbStatus) +// Request unit +struct TStatUnit { + TStatUnit(TInstant start) + : Start(start) + , RetryAttempts(0) + , IsFirstAttempt(true) {} - EInnerStatus InnerStatus; - NYdb::EStatus YdbStatus; -}; - -class TPeriodData; + void IncRetryAttempts() { + if (IsFirstAttempt) { + IsFirstAttempt = false; + } else { + ++RetryAttempts; + } + } -// Primitive latency unit -struct TStatUnit { - TStatUnit(const std::shared_ptr& periodData, TInstant startTime); - void Report(const TInnerStatus& status); - TDuration Delay() const; - - std::shared_ptr PeriodData; TInstant Start; TInstant End; -}; - -struct TCounters { - std::uint64_t Infly = 0; - std::uint64_t ActiveSessions = 0; - - // Debug use only: - std::uint64_t ReadPromises = 0; - std::uint64_t ExecutorPromises = 0; -}; - -struct TReplies { - std::map Statuses; - std::uint64_t CountMaxInfly = 0; - std::uint64_t CountHighLatency = 0; - std::uint64_t ApplicationTimeout = 0; - std::uint64_t NotFinished = 0; -}; -// Calculated percentiles for a period of time -struct TPercentile { - TDuration P50; - TDuration P90; - TDuration P95; - TDuration P99; - TDuration P99_9; - TDuration P100; + std::uint64_t RetryAttempts; + bool IsFirstAttempt; }; -// Accumulated statistics for a period of time -struct TPeriodStat { - std::uint64_t Seconds = 0; - TCounters Counters; - TReplies Replies; - TPercentile Oks; - TPercentile NotOks; -}; - -class TStat; - -// Full latency data for a period of time -class TPeriodData : public std::enable_shared_from_this { +class TStat { public: - TPeriodData( - TStat* stats, - bool RetryMode, - const TDuration maxDelay, - std::uint64_t currentSecond, - std::uint64_t currInfly, - std::uint64_t currSessions, - std::uint64_t currPromises, - std::uint64_t currExecutorPromises - ); - ~TPeriodData(); - TStatUnit CreateStatUnit(TInstant startTime); - void AddStat(TDuration delay, const TInnerStatus& status); - std::uint64_t GetCurrentSecond() const; - void ReportMaxInfly(); - void ReportInfly(std::uint64_t infly); - void ReportActiveSessions(std::uint64_t sessions); - - // Debug use only: - void ReportReadPromises(std::uint64_t promises); - void ReportExecutorPromises(std::uint64_t promises); + explicit TStat(const std::optional& metricsPushUrl, const std::string& operationType); -private: - TStat* Stats; - bool RetryMode; - std::uint64_t CurrentSecond; - std::vector OkDelays; - std::vector NotOkDelays; - TCounters Counters; - TReplies Replies; - TDuration MaxDelay; -}; + void Start(); + void Finish(); -class IMetricsPusher { -public: - virtual ~IMetricsPusher() = default; - virtual void PushData(const TPeriodStat& p) = 0; -}; + std::shared_ptr StartRequest(); + void FinishRequest(const std::shared_ptr& unit, const TFinalStatus& status); -class TStat { - friend class TPeriodData; -public: - TStat( - TDuration maxDelay, - const std::string& resultFileName, - bool pushMetrics, - bool retryMode - ); - ~TStat(); - void Reset(); - void Finish(); - void Flush(); - TStatUnit CreateStatUnit(); - void Report(TStatUnit& unit, const TInnerStatus& status); - void Report(TStatUnit& unit, TInnerStatus::EInnerStatus innerStatus); - void Report(TStatUnit& unit, const TFinalStatus& status); - void Report(TStatUnit& unit, NYdb::EStatus status); void ReportMaxInfly(); - void ReportStats(std::uint64_t infly, std::uint64_t sessions, std::uint64_t readPromises, std::uint64_t executorPromises); - void Reserve(size_t size); - void ReportLatencyData( - std::uint64_t currSecond, - TCounters&& counters, - TReplies&& replies, - std::vector&& oks, - std::vector& notOks - ); - void UpdateSessionStats(const std::unordered_map& sessionStats); + void ReportStats(std::uint64_t sessions, std::uint64_t readPromises, std::uint64_t executorPromises); void PrintStatistics(TStringBuilder& out); - void SaveResult(); TInstant GetStartTime() const; private: - void OnReport(TStatUnit& unit, const TInnerStatus& status); - void CheckCurrentSecond(TInstant now); - void PushMetricsData(const TPeriodStat& p); - void ResetMetricsPusher(std::uint64_t timestamp); - void CalculateGlobalPercentile(); - void CalculateFailSeconds(); - std::uint64_t GetTotal(); + void ScheduleMetricsPush(std::function func); std::recursive_mutex Mutex; TThreadPool MetricsPushQueue; - // Current period we collect stats for - std::shared_ptr ActivePeriod; - std::vector LatencyStats; - std::vector> LatencyData; - TDuration MaxDelay; + TInstant StartTime; TInstant FinishTime; - std::unordered_map SessionStats; - std::uint64_t CurrentSecond = 0; + // program lifetime - TCounters Counters; - TReplies Replies; - bool PushMetrics; - bool RetryMode; - std::string ResultFileName; - std::unique_ptr MetricsPusher; + std::uint64_t Infly = 0; + std::uint64_t ActiveSessions = 0; - std::unique_ptr GlobalPercentile; - std::unique_ptr FailSeconds; -}; + std::map Statuses; + std::uint64_t CountMaxInfly = 0; + std::uint64_t ApplicationTimeout = 0; + std::vector OkDelays; -std::string YdbStatusToString(NYdb::EStatus status); + // Debug use only: + std::uint64_t ReadPromises = 0; + std::uint64_t ExecutorPromises = 0; + + std::unique_ptr MetricsPusher; +}; diff --git a/tests/slo_workloads/utils/utils.cpp b/tests/slo_workloads/utils/utils.cpp index 2388bf4caca..b7572c29811 100644 --- a/tests/slo_workloads/utils/utils.cpp +++ b/tests/slo_workloads/utils/utils.cpp @@ -16,7 +16,6 @@ using namespace NYdb; const TDuration DefaultReactionTime = TDuration::Minutes(2); const TDuration ReactionTimeDelay = TDuration::MilliSeconds(5); -const TDuration GlobalTimeout = TDuration::Minutes(2); const std::uint64_t PartitionsCount = 64; Y_DECLARE_OUT_SPEC(, NYdb::TStatus, stream, value) { @@ -289,6 +288,73 @@ std::uint32_t GetHash(std::uint32_t value) { return result; } +std::string YdbStatusToString(NYdb::EStatus status) { + switch (status) { + case NYdb::EStatus::SUCCESS: + return "SUCCESS"; + case NYdb::EStatus::BAD_REQUEST: + return "BAD_REQUEST"; + case NYdb::EStatus::UNAUTHORIZED: + return "UNAUTHORIZED"; + case NYdb::EStatus::INTERNAL_ERROR: + return "INTERNAL_ERROR"; + case NYdb::EStatus::ABORTED: + return "ABORTED"; + case NYdb::EStatus::UNAVAILABLE: + return "UNAVAILABLE"; + case NYdb::EStatus::OVERLOADED: + return "OVERLOADED"; + case NYdb::EStatus::SCHEME_ERROR: + return "SCHEME_ERROR"; + case NYdb::EStatus::GENERIC_ERROR: + return "GENERIC_ERROR"; + case NYdb::EStatus::TIMEOUT: + return "TIMEOUT"; + case NYdb::EStatus::BAD_SESSION: + return "BAD_SESSION"; + case NYdb::EStatus::PRECONDITION_FAILED: + return "PRECONDITION_FAILED"; + case NYdb::EStatus::ALREADY_EXISTS: + return "ALREADY_EXISTS"; + case NYdb::EStatus::NOT_FOUND: + return "NOT_FOUND"; + case NYdb::EStatus::SESSION_EXPIRED: + return "SESSION_EXPIRED"; + case NYdb::EStatus::CANCELLED: + return "CANCELLED"; + case NYdb::EStatus::UNDETERMINED: + return "UNDETERMINED"; + case NYdb::EStatus::UNSUPPORTED: + return "UNSUPPORTED"; + case NYdb::EStatus::SESSION_BUSY: + return "SESSION_BUSY"; + case NYdb::EStatus::EXTERNAL_ERROR: + return "EXTERNAL_ERROR"; + case NYdb::EStatus::STATUS_UNDEFINED: + return "STATUS_UNDEFINED"; + case NYdb::EStatus::TRANSPORT_UNAVAILABLE: + return "TRANSPORT_UNAVAILABLE"; + case NYdb::EStatus::CLIENT_RESOURCE_EXHAUSTED: + return "CLIENT_RESOURCE_EXHAUSTED"; + case NYdb::EStatus::CLIENT_DEADLINE_EXCEEDED: + return "CLIENT_DEADLINE_EXCEEDED"; + case NYdb::EStatus::CLIENT_INTERNAL_ERROR: + return "CLIENT_INTERNAL_ERROR"; + case NYdb::EStatus::CLIENT_CANCELLED: + return "CLIENT_CANCELLED"; + case NYdb::EStatus::CLIENT_UNAUTHENTICATED: + return "CLIENT_UNAUTHENTICATED"; + case NYdb::EStatus::CLIENT_CALL_UNIMPLEMENTED: + return "CLIENT_CALL_UNIMPLEMENTED"; + case NYdb::EStatus::CLIENT_OUT_OF_RANGE: + return "CLIENT_OUT_OF_RANGE"; + case NYdb::EStatus::CLIENT_DISCOVERY_FAILED: + return "CLIENT_DISCOVERY_FAILED"; + case NYdb::EStatus::CLIENT_LIMITS_REACHED: + return "CLIENT_LIMITS_REACHED"; + } +} + TTableStats GetTableStats(TDatabaseOptions& dbOptions, const std::string& tableName) { Cout << TInstant::Now().ToRfc822StringLocal() << " Getting table stats (maxId and count of rows) with ReadTable... " << Endl; @@ -358,33 +424,25 @@ TTableStats GetTableStats(TDatabaseOptions& dbOptions, const std::string& tableN return result; } -void ParseOptionsCommon(TOpts& opts, TCommonOptions& options, bool followers) { +void ParseOptionsCommon(TOpts& opts, TCommonOptions& options) { opts.AddLongOption("threads", "Number of threads to use").RequiredArgument("NUM") .DefaultValue(options.MaxInputThreads).StoreResult(&options.MaxInputThreads); - opts.AddLongOption("stop_on_error", "Stop thread if an error occured").NoArgument() + opts.AddLongOption("stop-on-error", "Stop thread if an error occured").NoArgument() .SetFlag(&options.StopOnError).DefaultValue(options.StopOnError); opts.AddLongOption("payload-min", "Minimum length of payload string").RequiredArgument("NUM") .DefaultValue(options.MinLength).StoreResult(&options.MinLength); opts.AddLongOption("payload-max", "Maximum length of payload string").RequiredArgument("NUM") .DefaultValue(options.MaxLength).StoreResult(&options.MaxLength); - opts.AddLongOption("timeout", "Read requests execution timeout [ms]").RequiredArgument("NUM") - .DefaultValue(options.A_ReactionTime).StoreResult(&options.A_ReactionTime); opts.AddLongOption("dont-push", "Do not push metrics").NoArgument() .SetFlag(&options.DontPushMetrics).DefaultValue(options.DontPushMetrics); - opts.AddLongOption("retry", "Retry each request until Ok reply or global timeout").NoArgument() - .SetFlag(&options.RetryMode).DefaultValue(options.RetryMode); - opts.AddLongOption("save-result", "Save result to file").NoArgument() - .SetFlag(&options.SaveResult).DefaultValue(options.SaveResult); - opts.AddLongOption("result-file-name", "Set result json file name").RequiredArgument("String") - .DefaultValue(options.ResultFileName).StoreResult(&options.ResultFileName); + opts.AddLongOption("metrics-push-url", "URL to push metrics").RequiredArgument("URL") + .DefaultValue(options.MetricsPushUrl).StoreResult(&options.MetricsPushUrl); opts.AddLongOption("app-timeout", "Use application timeout (over SDK)").NoArgument() .SetFlag(&options.UseApplicationTimeout).DefaultValue(options.UseApplicationTimeout); opts.AddLongOption("prevention-request", "Send prevention request at 1/2 of timeout").NoArgument() .SetFlag(&options.SendPreventiveRequest).DefaultValue(options.SendPreventiveRequest); - if (followers) { - opts.AddLongOption("followers", "Use followers").NoArgument() - .SetFlag(&options.UseFollowers).DefaultValue(options.UseFollowers); - } + + opts.MutuallyExclusive("dont-push", "metrics-push-url"); } bool CheckOptionsCommon(TCommonOptions& options) { @@ -396,16 +454,13 @@ bool CheckOptionsCommon(TCommonOptions& options) { Cerr << "--threads should be more than 0" << Endl; return false; } - if (!options.DontPushMetrics) { - Cerr << "Push metrics is not supported yet" << Endl; - return false; - } + return true; } -bool ParseOptionsCreate(int argc, char** argv, TCreateOptions& createOptions, bool followers) { +bool ParseOptionsCreate(int argc, char** argv, TCreateOptions& createOptions) { TOpts opts = TOpts::Default(); - ParseOptionsCommon(opts, createOptions.CommonOptions, followers); + ParseOptionsCommon(opts, createOptions.CommonOptions); opts.AddLongOption("count", "Total number of records to generate").RequiredArgument("NUM") .DefaultValue(createOptions.Count).StoreResult(&createOptions.Count); opts.AddLongOption("pack-size", "Number of new records in each create request").RequiredArgument("NUM") @@ -427,9 +482,9 @@ bool ParseOptionsCreate(int argc, char** argv, TCreateOptions& createOptions, bo return true; } -bool ParseOptionsRun(int argc, char** argv, TRunOptions& runOptions, bool followers) { +bool ParseOptionsRun(int argc, char** argv, TRunOptions& runOptions) { TOpts opts = TOpts::Default(); - ParseOptionsCommon(opts, runOptions.CommonOptions, followers); + ParseOptionsCommon(opts, runOptions.CommonOptions); opts.AddLongOption("time", "Time to run (Seconds)").RequiredArgument("Seconds") .DefaultValue(runOptions.CommonOptions.SecondsToRun).StoreResult(&runOptions.CommonOptions.SecondsToRun); opts.AddLongOption("read-rps", "Request generation rate for read requests (Thread A)").RequiredArgument("NUM") @@ -440,10 +495,12 @@ bool ParseOptionsRun(int argc, char** argv, TRunOptions& runOptions, bool follow .SetFlag(&runOptions.DontRunA).DefaultValue(runOptions.DontRunA); opts.AddLongOption("no-write", "Do not run writing requests (thread B)").NoArgument() .SetFlag(&runOptions.DontRunB).DefaultValue(runOptions.DontRunB); - opts.AddLongOption("no-c", "Do not run thread C").NoArgument() - .SetFlag(&runOptions.DontRunC).DefaultValue(runOptions.DontRunC); opts.AddLongOption("infly", "Maximum number of running jobs").RequiredArgument("NUM") .DefaultValue(runOptions.CommonOptions.MaxInfly).StoreResult(&runOptions.CommonOptions.MaxInfly); + opts.AddLongOption("read-timeout", "Read requests execution timeout [ms]").RequiredArgument("NUM") + .DefaultValue(runOptions.ReadTimeout).StoreResult(&runOptions.ReadTimeout); + opts.AddLongOption("write-timeout", "Write requests execution timeout [ms]").RequiredArgument("NUM") + .DefaultValue(runOptions.WriteTimeout).StoreResult(&runOptions.WriteTimeout); TOptsParseResult res(&opts, argc, argv); if (!CheckOptionsCommon(runOptions.CommonOptions)) { diff --git a/tests/slo_workloads/utils/utils.h b/tests/slo_workloads/utils/utils.h index 32b1ea3b9cb..65be9f48913 100644 --- a/tests/slo_workloads/utils/utils.h +++ b/tests/slo_workloads/utils/utils.h @@ -1,7 +1,5 @@ #pragma once -#include "statistics.h" - #include #include #include @@ -12,7 +10,6 @@ extern const TDuration DefaultReactionTime; extern const TDuration ReactionTimeDelay; -extern const TDuration GlobalTimeout; extern const std::uint64_t PartitionsCount; struct TRecordData { @@ -50,23 +47,18 @@ struct TCommonOptions { std::uint32_t MaxCallbackThreads = 50; std::uint32_t MaxInfly = 500; std::uint32_t MaxRetries = 50; - std::uint64_t A_ReactionTime = 70; //ms TDuration ReactionTime = DefaultReactionTime; bool StopOnError = false; bool UseApplicationTimeout = false; bool SendPreventiveRequest = false; - //Generator options: + // Generator options: std::uint32_t MinLength = 20; std::uint32_t MaxLength = 200; - //Output options: - bool DontPushMetrics = true; - std::string ResultFileName = "slo_result.json"; - - bool UseFollowers = false; - bool RetryMode = false; - bool SaveResult = false; + // Output options: + bool DontPushMetrics = false; + std::string MetricsPushUrl = "http://localhost:9090/api/v1/otlp/v1/metrics"; }; struct TCreateOptions { @@ -79,9 +71,10 @@ struct TRunOptions { TCommonOptions CommonOptions; bool DontRunA = false; bool DontRunB = false; - bool DontRunC = false; std::uint32_t Read_rps = 1000; std::uint32_t Write_rps = 10; + TDuration ReadTimeout = DefaultReactionTime; + TDuration WriteTimeout = DefaultReactionTime; }; class TRpsProvider { @@ -159,7 +152,9 @@ std::uint32_t GetShardSpecialId(std::uint64_t shardNo); std::uint32_t GetHash(std::uint32_t value); +std::string YdbStatusToString(NYdb::EStatus status); + TTableStats GetTableStats(TDatabaseOptions& dbOptions, const std::string& tableName); -bool ParseOptionsCreate(int argc, char** argv, TCreateOptions& createOptions, bool followers = false); -bool ParseOptionsRun(int argc, char** argv, TRunOptions& runOptions, bool followers = false); +bool ParseOptionsCreate(int argc, char** argv, TCreateOptions& createOptions); +bool ParseOptionsRun(int argc, char** argv, TRunOptions& runOptions); diff --git a/tests/unit/library/CMakeLists.txt b/tests/unit/library/CMakeLists.txt index bb3ffee9300..11d92442d65 100644 --- a/tests/unit/library/CMakeLists.txt +++ b/tests/unit/library/CMakeLists.txt @@ -2,3 +2,4 @@ add_subdirectory(decimal) add_subdirectory(grpc_client) add_subdirectory(issue) add_subdirectory(operation_id) +add_subdirectory(time) diff --git a/tests/unit/library/time/CMakeLists.txt b/tests/unit/library/time/CMakeLists.txt new file mode 100644 index 00000000000..9b78ff4d875 --- /dev/null +++ b/tests/unit/library/time/CMakeLists.txt @@ -0,0 +1,8 @@ +add_ydb_test(NAME time_ut GTEST + SOURCES + time_ut.cpp + LINK_LIBRARIES + ydb-time + LABELS + unit +) diff --git a/tests/unit/library/time/time_ut.cpp b/tests/unit/library/time/time_ut.cpp new file mode 100644 index 00000000000..5c7890adef3 --- /dev/null +++ b/tests/unit/library/time/time_ut.cpp @@ -0,0 +1,55 @@ +#include + +#include + +#include + +using namespace std::chrono_literals; +using namespace NYdb; + +TEST(DurationTest, Simple) { + ASSERT_EQ(TDeadline::SafeDurationCast(1s), 1s); + ASSERT_EQ(TDeadline::SafeDurationCast(1ms), 1ms); + ASSERT_EQ(TDeadline::SafeDurationCast(1us), 1us); + ASSERT_EQ(TDeadline::SafeDurationCast(1ns), 1ns); + + ASSERT_EQ(TDeadline::SafeDurationCast(TDuration::Seconds(1)), 1s); + ASSERT_EQ(TDeadline::SafeDurationCast(TDuration::MilliSeconds(1)), 1ms); + ASSERT_EQ(TDeadline::SafeDurationCast(TDuration::MicroSeconds(1)), 1us); +} + +TEST(DurationTest, CornerValues) { + ASSERT_EQ(TDeadline::SafeDurationCast(0s), TDeadline::Duration::zero()); + ASSERT_EQ(TDeadline::SafeDurationCast(0ms), TDeadline::Duration::zero()); + ASSERT_EQ(TDeadline::SafeDurationCast(0us), TDeadline::Duration::zero()); + ASSERT_EQ(TDeadline::SafeDurationCast(0ns), TDeadline::Duration::zero()); + + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::seconds::max()), TDeadline::Duration::max()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::milliseconds::max()), TDeadline::Duration::max()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::microseconds::max()), TDeadline::Duration::max()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::nanoseconds::max()), TDeadline::Duration::max()); + + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::seconds::min()), TDeadline::Duration::min()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::milliseconds::min()), TDeadline::Duration::min()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::microseconds::min()), TDeadline::Duration::min()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::nanoseconds::min()), TDeadline::Duration::min()); + + ASSERT_EQ(TDeadline::SafeDurationCast(TDuration::Zero()), TDeadline::Duration::zero()); + ASSERT_EQ(TDeadline::SafeDurationCast(TDuration::Max()), TDeadline::Duration::max()); + + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::seconds::max()), TDeadline::Duration::max()); + ASSERT_EQ(TDeadline::SafeDurationCast(std::chrono::milliseconds::max()), TDeadline::Duration::max()); +} + +TEST(DeadlineTest, Compare) { + ASSERT_EQ(TDeadline::AfterDuration(TDeadline::Duration::max()), TDeadline::Max()); + ASSERT_EQ(TDeadline::AfterDuration(std::chrono::seconds::max()), TDeadline::Max()); + ASSERT_EQ(TDeadline::AfterDuration(std::chrono::milliseconds::max()), TDeadline::Max()); + ASSERT_EQ(TDeadline::AfterDuration(std::chrono::microseconds::max()), TDeadline::Max()); + ASSERT_EQ(TDeadline::AfterDuration(std::chrono::nanoseconds::max()), TDeadline::Max()); + + ASSERT_EQ(TDeadline::AfterDuration(TDuration::Max()), TDeadline::Max()); + + ASSERT_LT(TDeadline::Now(), TDeadline::Max()); + ASSERT_LT(TDeadline::Now(), TDeadline::AfterDuration(100s)); +}