From 1a775f72f401c08cfb85f00feb0993c3c0b384bc Mon Sep 17 00:00:00 2001 From: Bulat Gayazov Date: Fri, 23 Aug 2024 17:14:05 +0000 Subject: [PATCH 01/35] Returned federated_topic tests --- .../federated_topic/ut/basic_usage_ut.cpp | 1406 +++++++++++++++++ src/client/federated_topic/ut/fds_mock.h | 178 +++ src/client/federated_topic/ut/ya.make | 38 + 3 files changed, 1622 insertions(+) create mode 100644 src/client/federated_topic/ut/basic_usage_ut.cpp create mode 100644 src/client/federated_topic/ut/fds_mock.h create mode 100644 src/client/federated_topic/ut/ya.make diff --git a/src/client/federated_topic/ut/basic_usage_ut.cpp b/src/client/federated_topic/ut/basic_usage_ut.cpp new file mode 100644 index 0000000000..1a98937b28 --- /dev/null +++ b/src/client/federated_topic/ut/basic_usage_ut.cpp @@ -0,0 +1,1406 @@ +#include +#include + +#include + +#include + +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +#include + +namespace NYdb::NFederatedTopic::NTests { + +Y_UNIT_TEST_SUITE(BasicUsage) { + + Y_UNIT_TEST(GetAllStartPartitionSessions) { + size_t partitionsCount = 5; + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2, partitionsCount); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + ReadSession->WaitEvent().Wait(TDuration::Seconds(1)); + TMaybe event = ReadSession->GetEvent(false); + Y_ASSERT(!event); + + auto fdsRequest = fdsMock.GetNextPendingRequest(); + Y_ASSERT(fdsRequest.has_value()); + // TODO check fdsRequest->Req db header + + Ydb::FederationDiscovery::ListFederationDatabasesResponse response; + + auto op = response.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + response.mutable_operation()->set_ready(true); + response.mutable_operation()->set_id("12345"); + + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("fancy_datacenter"); + auto c1 = mockResult.add_federation_databases(); + c1->set_name("dc1"); + c1->set_path("/Root"); + c1->set_id("account-dc1"); + c1->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c1->set_location("dc1"); + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c1->set_weight(1000); + auto c2 = mockResult.add_federation_databases(); + c2->set_name("dc2"); + c2->set_path("/Root"); + c2->set_id("account-dc2"); + c2->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c2->set_location("dc2"); + c2->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c2->set_weight(500); + + op->mutable_result()->PackFrom(mockResult); + + fdsRequest->Result.SetValue({std::move(response), grpc::Status::OK}); + + for (size_t i = 0; i < partitionsCount; ++i) { + ReadSession->WaitEvent().Wait(); + // Get event + TMaybe event = ReadSession->GetEvent(true/*block - will block if no event received yet*/); + Cerr << "Got new read session event: " << DebugString(*event) << Endl; + + auto* startPartitionSessionEvent = std::get_if(&*event); + Y_ASSERT(startPartitionSessionEvent); + startPartitionSessionEvent->Confirm(); + } + + ReadSession->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(WaitEventBlocksBeforeDiscovery) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + Cerr << "Before ReadSession was created" << Endl; + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + auto f = ReadSession->WaitEvent(); + Cerr << "Returned from WaitEvent" << Endl; + // observer asyncInit should respect client/session timeouts + UNIT_ASSERT(!f.Wait(TDuration::Seconds(1))); + + Cerr << "Session blocked successfully" << Endl; + + UNIT_ASSERT(ReadSession->Close(TDuration::MilliSeconds(10))); + Cerr << "Session closed gracefully" << Endl; + } + + Y_UNIT_TEST(RetryDiscoveryWithCancel) { + // TODO register requests in mock, compare time for retries, reschedules + + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings() + .RetryPolicy(NTopic::IRetryPolicy::GetFixedIntervalPolicy( + TDuration::Seconds(10), + TDuration::Seconds(10) + )); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + bool answerOk = false; + + for (int i = 0; i < 6; ++i) { + std::optional fdsRequest; + do { + Sleep(TDuration::MilliSeconds(50)); + fdsRequest = fdsMock.GetNextPendingRequest(); + + } while (!fdsRequest.has_value()); + + if (answerOk) { + fdsRequest->Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + } else { + fdsRequest->Result.SetValue({{}, grpc::Status(grpc::StatusCode::UNAVAILABLE, "mock 'unavailable'")}); + } + + answerOk = !answerOk; + } + } + + Y_UNIT_TEST(PropagateSessionClosed) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings() + .RetryPolicy(NTopic::IRetryPolicy::GetFixedIntervalPolicy( + TDuration::Seconds(10), + TDuration::Seconds(10) + )); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + Sleep(TDuration::MilliSeconds(50)); + + auto events = ReadSession->GetEvents(false); + UNIT_ASSERT(events.empty()); + + Ydb::FederationDiscovery::ListFederationDatabasesResponse Response; + + auto op = Response.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + Response.mutable_operation()->set_ready(true); + Response.mutable_operation()->set_id("12345"); + + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("fancy_datacenter"); + auto c1 = mockResult.add_federation_databases(); + c1->set_name("dc1"); + c1->set_path("/Root"); + c1->set_id("account-dc1"); + c1->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c1->set_location("dc1"); + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c1->set_weight(1000); + auto c2 = mockResult.add_federation_databases(); + c2->set_name("dc2"); + c2->set_path("/Invalid"); + c2->set_id("account-dc2"); + c2->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c2->set_location("dc2"); + c2->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c2->set_weight(500); + + op->mutable_result()->PackFrom(mockResult); + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue({Response, grpc::Status::OK}); + + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic()).MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + NPersQueue::IExecutor::TPtr executor = new NTopic::TSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& client = setup->GetPersQueueClient(); + auto session = client.CreateSimpleBlockingWriteSession(writeSettings); + TString messageBase = "message----"; + + + ui64 count = 100u; + for (auto i = 0u; i < count; i++) { + auto res = session->Write(messageBase * (200 * 1024) + ToString(i)); + UNIT_ASSERT(res); + + events = ReadSession->GetEvents(true); + UNIT_ASSERT(!events.empty()); + + for (auto& e : events) { + Cerr << ">>> Got event: " << DebugString(e) << Endl; + if (auto* dataEvent = std::get_if(&e)) { + dataEvent->Commit(); + } else if (auto* startPartitionSessionEvent = std::get_if(&e)) { + startPartitionSessionEvent->Confirm(); + } else if (auto* stopPartitionSessionEvent = std::get_if(&e)) { + stopPartitionSessionEvent->Confirm(); + } + } + } + + session->Close(); + } + + Y_UNIT_TEST(RecreateObserver) { + // TODO register requests in mock, compare time for retries, reschedules + + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings() + .RetryPolicy(NTopic::IRetryPolicy::GetNoRetryPolicy()); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + ReadSession->WaitEvent().Wait(TDuration::Seconds(1)); + auto event = ReadSession->GetEvent(false); + UNIT_ASSERT(!event.Defined()); + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue({{}, grpc::Status(grpc::StatusCode::UNAVAILABLE, "mock 'unavailable'")}); + + ReadSession->WaitEvent().Wait(); + event = ReadSession->GetEvent(false); + UNIT_ASSERT(event.Defined()); + Cerr << ">>> Got event: " << DebugString(*event) << Endl; + UNIT_ASSERT(std::holds_alternative(*event)); + + auto ReadSession2 = topicClient.CreateReadSession(readSettings); + Cerr << "Session2 was created" << Endl; + + ReadSession2->WaitEvent().Wait(TDuration::Seconds(1)); + event = ReadSession2->GetEvent(false); + UNIT_ASSERT(!event.Defined()); + + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + event = ReadSession2->GetEvent(true); + UNIT_ASSERT(event.Defined()); + Cerr << ">>> Got event: " << DebugString(*event) << Endl; + UNIT_ASSERT(std::holds_alternative(*event)); + + // Cerr << ">>> Got event: " << DebugString(*event) << Endl; + // UNIT_ASSERT(std::holds_alternative(*event)); + } + + Y_UNIT_TEST(FallbackToSingleDb) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + + setup->Start(true, true); + + std::shared_ptr ReadSession; + + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(setup->GetDriver()); + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + ReadSession->WaitEvent().Wait(TDuration::Seconds(1)); + TMaybe event = ReadSession->GetEvent(false); + Y_ASSERT(event); + Cerr << "Got new read session event: " << DebugString(*event) << Endl; + + auto* startPartitionSessionEvent = std::get_if(&*event); + Y_ASSERT(startPartitionSessionEvent); + startPartitionSessionEvent->Confirm(); + + ReadSession->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(FallbackToSingleDbAfterBadRequest) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + TFederationDiscoveryServiceMock fdsMock; + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + fdsMock.Port = setup->GetGrpcPort(); + + Cerr << "PORTS " << fdsMock.Port << " " << newServicePort << Endl; + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings(); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + auto ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << "Session was created" << Endl; + + ReadSession->WaitEvent().Wait(TDuration::Seconds(1)); + TMaybe event = ReadSession->GetEvent(false); + Y_ASSERT(!event); + + { + auto fdsRequest = fdsMock.GetNextPendingRequest(); + Y_ASSERT(fdsRequest.has_value()); + Ydb::FederationDiscovery::ListFederationDatabasesResponse response; + auto op = response.mutable_operation(); + op->set_status(Ydb::StatusIds::BAD_REQUEST); + response.mutable_operation()->set_ready(true); + response.mutable_operation()->set_id("12345"); + fdsRequest->Result.SetValue({std::move(response), grpc::Status::OK}); + } + + ReadSession->WaitEvent().Wait(); + TMaybe event2 = ReadSession->GetEvent(true); + Cerr << "Got new read session event: " << DebugString(*event2) << Endl; + + auto* sessionEvent = std::get_if(&*event2); + // At this point the SDK should connect to Topic API, but in this test we have it on a separate port, + // so SDK connects back to FederationDiscovery service and the CLIENT_CALL_UNIMPLEMENTED status is expected. + UNIT_ASSERT_EQUAL(sessionEvent->GetStatus(), EStatus::CLIENT_CALL_UNIMPLEMENTED); + ReadSession->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(SimpleHandlers) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings() + .RetryPolicy(NTopic::IRetryPolicy::GetFixedIntervalPolicy( + TDuration::Seconds(10), + TDuration::Seconds(10) + )); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + ui64 count = 300u; + + TString messageBase = "message----"; + TVector sentMessages; + + for (auto i = 0u; i < count; i++) { + // sentMessages.emplace_back(messageBase * (i+1) + ToString(i)); + sentMessages.emplace_back(messageBase * (10 * i + 1)); + } + + NThreading::TPromise checkedPromise = NThreading::NewPromise(); + auto totalReceived = 0u; + + auto f = checkedPromise.GetFuture(); + TAtomic check = 1; + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(1_MB) + .AppendTopics(setup->GetTestTopic()); + + readSettings.FederatedEventHandlers_.SimpleDataHandlers([&](TReadSessionEvent::TDataReceivedEvent& ev) mutable { + Cerr << ">>> event from dataHandler: " << DebugString(ev) << Endl; + Y_VERIFY_S(AtomicGet(check) != 0, "check is false"); + auto& messages = ev.GetMessages(); + for (size_t i = 0u; i < messages.size(); ++i) { + auto& message = messages[i]; + UNIT_ASSERT_VALUES_EQUAL(message.GetData(), sentMessages[totalReceived]); + totalReceived++; + } + if (totalReceived == sentMessages.size()) + checkedPromise.SetValue(); + }); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << ">>> Session was created" << Endl; + + Sleep(TDuration::MilliSeconds(50)); + + auto events = ReadSession->GetEvents(false); + UNIT_ASSERT(events.empty()); + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic()).MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + NPersQueue::IExecutor::TPtr executor = new NTopic::TSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& client = setup->GetPersQueueClient(); + auto session = client.CreateSimpleBlockingWriteSession(writeSettings); + + for (auto i = 0u; i < count; i++) { + auto res = session->Write(sentMessages[i]); + UNIT_ASSERT(res); + } + + f.GetValueSync(); + ReadSession->Close(); + AtomicSet(check, 0); + } + + Y_UNIT_TEST(ReadMirrored) { + auto setup = std::make_shared(TEST_CASE_NAME, false); + setup->Start(true, true); + setup->CreateTopic(setup->GetTestTopic() + "-mirrored-from-dc2", setup->GetLocalCluster()); + setup->CreateTopic(setup->GetTestTopic() + "-mirrored-from-dc3", setup->GetLocalCluster()); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr ReadSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + auto clientSettings = TFederatedTopicClientSettings() + .RetryPolicy(NTopic::IRetryPolicy::GetFixedIntervalPolicy( + TDuration::Seconds(10), + TDuration::Seconds(10) + )); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + ui64 count = 5u; + + TString messageBase = "message----"; + TVector sentMessages; + std::unordered_set sentSet; + + for (auto i = 0u; i < count; i++) { + sentMessages.emplace_back(messageBase * (10 * i + 1)); + sentSet.emplace(sentMessages.back() + "-from-dc1"); + sentSet.emplace(sentMessages.back() + "-from-dc2"); + sentSet.emplace(sentMessages.back() + "-from-dc3"); + } + + NThreading::TPromise checkedPromise = NThreading::NewPromise(); + auto totalReceived = 0u; + + auto f = checkedPromise.GetFuture(); + TAtomic check = 1; + + // Create read session. + NYdb::NFederatedTopic::TFederatedReadSessionSettings readSettings; + readSettings + .ReadMirrored("dc1") + .ConsumerName("shared/user") + .MaxMemoryUsageBytes(16_MB) + .AppendTopics(setup->GetTestTopic()); + + readSettings.FederatedEventHandlers_.SimpleDataHandlers([&](TReadSessionEvent::TDataReceivedEvent& ev) mutable { + Cerr << ">>> event from dataHandler: " << DebugString(ev) << Endl; + Y_VERIFY_S(AtomicGet(check) != 0, "check is false"); + auto& messages = ev.GetMessages(); + Cerr << ">>> get " << messages.size() << " messages in this event" << Endl; + for (size_t i = 0u; i < messages.size(); ++i) { + auto& message = messages[i]; + UNIT_ASSERT(message.GetFederatedPartitionSession()->GetReadSourceDatabaseName() == "dc1"); + UNIT_ASSERT(message.GetFederatedPartitionSession()->GetTopicPath() == setup->GetTestTopic()); + UNIT_ASSERT(message.GetData().EndsWith(message.GetFederatedPartitionSession()->GetTopicOriginDatabaseName())); + + UNIT_ASSERT(!sentSet.empty()); + UNIT_ASSERT_C(sentSet.erase(message.GetData()), "no such element is sentSet: " + message.GetData()); + totalReceived++; + } + if (totalReceived == 3 * sentMessages.size()) { + UNIT_ASSERT(sentSet.empty()); + checkedPromise.SetValue(); + } + }); + + ReadSession = topicClient.CreateReadSession(readSettings); + Cerr << ">>> Session was created" << Endl; + + Sleep(TDuration::MilliSeconds(50)); + + auto events = ReadSession->GetEvents(false); + UNIT_ASSERT(events.empty()); + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + { + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic()).MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + NPersQueue::IExecutor::TPtr executor = new NTopic::TSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& client = setup->GetPersQueueClient(); + auto session = client.CreateSimpleBlockingWriteSession(writeSettings); + + for (auto i = 0u; i < count; i++) { + auto res = session->Write(sentMessages[i] + "-from-dc1"); + UNIT_ASSERT(res); + } + + session->Close(); + + Cerr << ">>> Writes to test-topic successful" << Endl; + } + + { + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic() + "-mirrored-from-dc2").MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + NPersQueue::IExecutor::TPtr executor = new NTopic::TSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& client = setup->GetPersQueueClient(); + auto session = client.CreateSimpleBlockingWriteSession(writeSettings); + + for (auto i = 0u; i < count; i++) { + auto res = session->Write(sentMessages[i] + "-from-dc2"); + UNIT_ASSERT(res); + } + + session->Close(); + + Cerr << ">>> Writes to test-topic-mirrored-from-dc2 successful" << Endl; + } + + { + NPersQueue::TWriteSessionSettings writeSettings; + writeSettings.Path(setup->GetTestTopic() + "-mirrored-from-dc3").MessageGroupId("src_id"); + writeSettings.Codec(NPersQueue::ECodec::RAW); + NPersQueue::IExecutor::TPtr executor = new NTopic::TSyncExecutor(); + writeSettings.CompressionExecutor(executor); + + auto& client = setup->GetPersQueueClient(); + auto session = client.CreateSimpleBlockingWriteSession(writeSettings); + + for (auto i = 0u; i < count; i++) { + auto res = session->Write(sentMessages[i] + "-from-dc3"); + UNIT_ASSERT(res); + } + + session->Close(); + + Cerr << ">>> Writes to test-topic-mirrored-from-dc3 successful" << Endl; + } + + f.GetValueSync(); + ReadSession->Close(); + AtomicSet(check, 0); + } + + Y_UNIT_TEST(BasicWriteSession) { + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr WriteSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + // Create write session. + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + WriteSession = topicClient.CreateWriteSession(writeSettings); + Cerr << "Session was created" << Endl; + + WriteSession->WaitEvent().Wait(TDuration::Seconds(1)); + auto event = WriteSession->GetEvent(false); + Y_ASSERT(event); + Cerr << "Got new write session event: " << DebugString(*event) << Endl; + auto* readyToAcceptEvent = std::get_if(&*event); + Y_ASSERT(readyToAcceptEvent); + WriteSession->Write(std::move(readyToAcceptEvent->ContinuationToken), NTopic::TWriteMessage("hello")); + + WriteSession->WaitEvent().Wait(TDuration::Seconds(1)); + event = WriteSession->GetEvent(false); + Y_ASSERT(event); + Cerr << "Got new write session event: " << DebugString(*event) << Endl; + + readyToAcceptEvent = std::get_if(&*event); + Y_ASSERT(readyToAcceptEvent); + + std::optional fdsRequest; + do { + fdsRequest = fdsMock.GetNextPendingRequest(); + if (!fdsRequest.has_value()) { + Sleep(TDuration::MilliSeconds(50)); + } + } while (!fdsRequest.has_value()); + + fdsRequest->Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + WriteSession->WaitEvent().Wait(TDuration::Seconds(1)); + event = WriteSession->GetEvent(false); + Y_ASSERT(event); + Cerr << "Got new write session event: " << DebugString(*event) << Endl; + + auto* acksEvent = std::get_if(&*event); + Y_ASSERT(acksEvent); + + WriteSession->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(CloseWriteSessionImmediately) { + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + std::shared_ptr WriteSession; + + // Create topic client. + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + // Create write session. + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + WriteSession = topicClient.CreateWriteSession(writeSettings); + Cerr << "Session was created" << Endl; + + WriteSession->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(WriteSessionCloseWaitsForWrites) { + // Write a bunch of messages before the federation observer initialization. + // Then as soon as the federation discovery service responds to the observer, close the write session. + // The federated write session must wait for acks of all written messages. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + // Create write session. + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + int acks = 0; + int messageCount = 100; + + auto gotAllAcks = NThreading::NewPromise(); + writeSettings.EventHandlers_.AcksHandler([&](NTopic::TWriteSessionEvent::TAcksEvent& ev) { + acks += ev.Acks.size(); + if (acks == messageCount) { + gotAllAcks.SetValue(); + } + }); + + auto WriteSession = topicClient.CreateWriteSession(writeSettings); + Cerr << "Session was created" << Endl; + + for (int i = 0; i < messageCount; ++i) { + auto event = WriteSession->GetEvent(true); + auto* readyToAcceptEvent = std::get_if(&*event); + WriteSession->Write(std::move(readyToAcceptEvent->ContinuationToken), NTopic::TWriteMessage("hello-" + ToString(i))); + } + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + // Close method should wait until all messages have been acked. + WriteSession->Close(); + gotAllAcks.GetFuture().Wait(); + UNIT_ASSERT_VALUES_EQUAL(acks, messageCount); + } + + Y_UNIT_TEST(WriteSessionCloseIgnoresWrites) { + // Create a federated topic client with NoRetryPolicy. + // Make federation discovery service to respond with an UNAVAILABLE status. + // It makes a federation observer object to become stale and the write session to close. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + TFederatedTopicClientSettings clientSettings; + clientSettings.RetryPolicy(NPersQueue::IRetryPolicy::GetNoRetryPolicy()); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + // Create write session. + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + int acks = 0; + writeSettings.EventHandlers_.AcksHandler([&acks](NTopic::TWriteSessionEvent::TAcksEvent& ev) { + acks += ev.Acks.size(); + }); + + auto WriteSession = topicClient.CreateWriteSession(writeSettings); + Cerr << "Session was created" << Endl; + + int messageCount = 100; + for (int i = 0; i < messageCount; ++i) { + auto event = WriteSession->GetEvent(true); + auto* readyToAcceptEvent = std::get_if(&*event); + WriteSession->Write(std::move(readyToAcceptEvent->ContinuationToken), NTopic::TWriteMessage("hello-" + ToString(i))); + } + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeUnavailableResult()); + + // At this point the observer that federated write session works with should become stale, and the session closes. + // No messages we have written and no federation discovery requests should be sent. + + Sleep(TDuration::Seconds(3)); + UNIT_ASSERT(!fdsMock.GetNextPendingRequest().has_value()); + WriteSession->Close(); + UNIT_ASSERT_VALUES_EQUAL(acks, 0); + } + + Y_UNIT_TEST(PreferredDatabaseNoFallback) { + // The test checks that the session keeps trying to connect to the preferred database + // and does not fall back to other databases. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver); + + auto retryPolicy = std::make_shared(); + + auto writeSettings = TFederatedWriteSessionSettings() + .AllowFallback(false) + .PreferredDatabase("dc2"); + + writeSettings + .RetryPolicy(retryPolicy) + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + retryPolicy->Initialize(); + retryPolicy->ExpectBreakDown(); + + auto writer = topicClient.CreateWriteSession(writeSettings); + + Ydb::FederationDiscovery::ListFederationDatabasesResponse response; + auto op = response.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + response.mutable_operation()->set_ready(true); + response.mutable_operation()->set_id("12345"); + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("fancy_datacenter"); + auto c1 = mockResult.add_federation_databases(); + c1->set_name("dc1"); + c1->set_path("/Root"); + c1->set_id("account-dc1"); + c1->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c1->set_location("dc1"); + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + c1->set_weight(1000); + auto c2 = mockResult.add_federation_databases(); + c2->set_name("dc2"); + c2->set_path("/Root"); + c2->set_id("account-dc2"); + c2->set_endpoint("localhost:" + ToString(fdsMock.Port)); + c2->set_location("dc2"); + c2->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_UNAVAILABLE); + c2->set_weight(500); + op->mutable_result()->PackFrom(mockResult); + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue({std::move(response), grpc::Status::OK}); + + Cerr << "=== Session was created, waiting for retries" << Endl; + retryPolicy->WaitForRetriesSync(3); + + Cerr << "=== In the next federation discovery response dc2 will be available" << Endl; + // fdsMock.PreparedResponse.Clear(); + // std::optional fdsRequest; + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + Cerr << "=== Waiting for repair" << Endl; + retryPolicy->WaitForRepairSync(); + + Cerr << "=== Closing the session" << Endl; + writer->Close(TDuration::MilliSeconds(10)); + } + + Y_UNIT_TEST(WriteSessionNoAvailableDatabase) { + // Create a federated write session with NoRetryPolicy, PreferredDatabase set and AllowFallback=false. + // Make federation discovery service to respond with an UNAVAILABLE status for the preferred database. + // It makes the write session to close itself. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true, true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << newServicePort); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + TFederatedTopicClientSettings clientSettings; + NYdb::NFederatedTopic::TFederatedTopicClient topicClient(driver, clientSettings); + + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .RetryPolicy(NPersQueue::IRetryPolicy::GetNoRetryPolicy()) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + auto WriteSession = topicClient.CreateWriteSession(writeSettings); + Cerr << "Session was created" << Endl; + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultUnavailableDatabases()); + + { + auto e = WriteSession->GetEvent(true); + UNIT_ASSERT(e.Defined()); + Cerr << ">>> Got event: " << DebugString(*e) << Endl; + UNIT_ASSERT(std::holds_alternative(*e)); + } + { + auto e = WriteSession->GetEvent(true); + UNIT_ASSERT(e.Defined()); + Cerr << ">>> Got event: " << DebugString(*e) << Endl; + UNIT_ASSERT(std::holds_alternative(*e)); + } + + WriteSession->Close(); + } + + NTopic::TContinuationToken GetToken(std::shared_ptr writer) { + auto e = writer->GetEvent(true); + UNIT_ASSERT(e.Defined()); + Cerr << ">>> Got event: " << DebugString(*e) << Endl; + auto* readyToAcceptEvent = std::get_if(&*e); + UNIT_ASSERT(readyToAcceptEvent); + return std::move(readyToAcceptEvent->ContinuationToken); + } + + Y_UNIT_TEST(WriteSessionSwitchDatabases) { + // Test that the federated write session doesn't deadlock when reconnecting to another database, + // if the updated state of the federation is different from the previous one. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + + setup->Start(true); + + TFederationDiscoveryServiceMock fdsMock; + fdsMock.Port = setup->GetGrpcPort(); + ui16 newServicePort = setup->GetPortManager()->GetPort(4285); + auto grpcServer = setup->StartGrpcService(newServicePort, &fdsMock); + + auto driverConfig = NYdb::TDriverConfig() + .SetEndpoint(TStringBuilder() << "localhost:" << newServicePort) + .SetDatabase("/Root") + .SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + auto driver = NYdb::TDriver(driverConfig); + auto topicClient = NYdb::NFederatedTopic::TFederatedTopicClient(driver); + + auto writeSettings = NTopic::TFederatedWriteSessionSettings() + .PreferredDatabase("dc1") + .AllowFallback(true) + .DirectWriteToPartition(false) + .RetryPolicy(NPersQueue::IRetryPolicy::GetNoRetryPolicy()) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + size_t successfulSessionClosedEvents = 0; + size_t otherSessionClosedEvents = 0; + + writeSettings + .EventHandlers_.SessionClosedHandler([&](const NTopic::TSessionClosedEvent &ev) { + ++(ev.IsSuccess() ? successfulSessionClosedEvents : otherSessionClosedEvents); + }); + + writeSettings.EventHandlers_.HandlersExecutor(NTopic::CreateSyncExecutor()); + + auto WriteSession = topicClient.CreateWriteSession(writeSettings); + + TMaybe token; + + auto fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + { + WriteSession->Write(GetToken(WriteSession), NTopic::TWriteMessage("hello 1")); + token = GetToken(WriteSession); + + auto e = WriteSession->GetEvent(true); + auto* acksEvent = std::get_if(&*e); + UNIT_ASSERT(acksEvent); + } + + // Wait for two requests to the federation discovery service. + // This way we ensure the federated write session has had enough time to request + // the updated state of the federation from its federation observer. + + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultWithUnavailableDatabase(1)); + + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultWithUnavailableDatabase(1)); + + { + UNIT_ASSERT(token.Defined()); + WriteSession->Write(std::move(*token), NTopic::TWriteMessage("hello 2")); + + token = GetToken(WriteSession); + + auto e = WriteSession->GetEvent(true); + auto* acksEvent = std::get_if(&*e); + UNIT_ASSERT(acksEvent); + } + + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + fdsRequest = fdsMock.WaitNextPendingRequest(); + fdsRequest.Result.SetValue(fdsMock.ComposeOkResultAvailableDatabases()); + + { + UNIT_ASSERT(token.Defined()); + WriteSession->Write(std::move(*token), NTopic::TWriteMessage("hello 3")); + + token = GetToken(WriteSession); + + auto e = WriteSession->GetEvent(true); + auto* acksEvent = std::get_if(&*e); + UNIT_ASSERT(acksEvent); + } + + setup->ShutdownGRpc(); + + WriteSession->Close(TDuration::Seconds(5)); + + UNIT_ASSERT_VALUES_EQUAL(otherSessionClosedEvents, 1); + UNIT_ASSERT_VALUES_EQUAL(successfulSessionClosedEvents, 0); + } + + Y_UNIT_TEST(WriteSessionWriteInHandlers) { + // Write messages from all event handlers. It shouldn't deadlock. + + auto setup = std::make_shared( + TEST_CASE_NAME, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 2); + setup->Start(true, true); + NYdb::TDriverConfig cfg; + cfg.SetEndpoint(TStringBuilder() << "localhost:" << setup->GetGrpcPort()); + cfg.SetDatabase("/Root"); + cfg.SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)); + NYdb::TDriver driver(cfg); + NYdb::NFederatedTopic::TFederatedTopicClient client(driver); + + auto writeSettings = NTopic::TWriteSessionSettings() + .DirectWriteToPartition(false) + .Path(setup->GetTestTopic()) + .MessageGroupId("src_id"); + + std::shared_ptr WriteSession; + + // 1. The session issues the first token: write a message inside AcksHandler. + // 2. The session issues another token: close the session, write a message inside SessionClosedHandler. + + bool gotAcksEvent = false; + std::optional token; + auto [acksHandlerPromise, sessionClosedHandlerPromise] = std::tuple{NThreading::NewPromise(), NThreading::NewPromise()}; + auto [sentFromAcksHandler, sentFromSessionClosedHandler] = std::tuple{acksHandlerPromise.GetFuture(), sessionClosedHandlerPromise.GetFuture()}; + writeSettings.EventHandlers( + NTopic::TWriteSessionSettings::TEventHandlers() + .HandlersExecutor(NTopic::CreateSyncExecutor()) + .ReadyToAcceptHandler([&token](NTopic::TWriteSessionEvent::TReadyToAcceptEvent& e) { + Cerr << "=== Inside ReadyToAcceptHandler" << Endl; + token = std::move(e.ContinuationToken); + }) + .AcksHandler([&token, &gotAcksEvent, &WriteSession, promise = std::move(acksHandlerPromise)](NTopic::TWriteSessionEvent::TAcksEvent&) mutable { + Cerr << "=== Inside AcksHandler" << Endl; + if (!gotAcksEvent) { + gotAcksEvent = true; + WriteSession->Write(std::move(*token), "From AcksHandler"); + promise.SetValue(); + } + }) + .SessionClosedHandler([&token, &WriteSession, promise = std::move(sessionClosedHandlerPromise)](NTopic::TSessionClosedEvent const&) mutable { + Cerr << "=== Inside SessionClosedHandler" << Endl; + WriteSession->Write(std::move(*token), "From SessionClosedHandler"); + promise.SetValue(); + }) + ); + + Cerr << "=== Before CreateWriteSession" << Endl; + WriteSession = client.CreateWriteSession(writeSettings); + Cerr << "=== Session created" << Endl; + + WriteSession->Write(std::move(*token), "After CreateWriteSession"); + + sentFromAcksHandler.Wait(); + Cerr << "=== AcksHandler has written a message, closing the session" << Endl; + + WriteSession->Close(); + sentFromSessionClosedHandler.Wait(); + Cerr << "=== SessionClosedHandler has 'written' a message" << Endl; + } + + void AddDatabase(std::vector>& dbInfos, int id, int weight) { + auto db = std::make_shared(); + db->set_id(ToString(id)); + db->set_name("db" + ToString(id)); + db->set_location("dc" + ToString(id)); + db->set_weight(weight); + db->set_status(Ydb::FederationDiscovery::DatabaseInfo_Status_AVAILABLE); + dbInfos.push_back(db); + } + + void EnableDatabase(std::vector>& dbInfos, int id) { + dbInfos[id - 1]->set_status(Ydb::FederationDiscovery::DatabaseInfo_Status_AVAILABLE); + } + + void DisableDatabase(std::vector>& dbInfos, int id) { + dbInfos[id - 1]->set_status(Ydb::FederationDiscovery::DatabaseInfo_Status_UNAVAILABLE); + } + + Y_UNIT_TEST(SelectDatabaseByHash) { + std::vector> dbInfos; + using Settings = TFederatedWriteSessionSettings; + + { + auto [db, status] = SelectDatabaseByHashImpl(Settings(), dbInfos); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::NOT_FOUND); + } + + AddDatabase(dbInfos, 1, 0); + + { + auto [db, status] = SelectDatabaseByHashImpl(Settings(), dbInfos); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::NOT_FOUND); + } + + AddDatabase(dbInfos, 2, 100); + + { + auto [db, status] = SelectDatabaseByHashImpl(Settings(), dbInfos); + UNIT_ASSERT(db); + UNIT_ASSERT_EQUAL(db->id(), "2"); + UNIT_ASSERT_EQUAL(status, EStatus::SUCCESS); + } + } + + Y_UNIT_TEST(SelectDatabase) { + std::vector> dbInfos; + for (int i = 1; i < 11; ++i) { + AddDatabase(dbInfos, i, 1000); + } + + using Settings = TFederatedWriteSessionSettings; + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | set | not found | - | any | NOT_FOUND | + */ + for (bool allow : {false, true}) { + auto settings = Settings().PreferredDatabase("db0").AllowFallback(allow); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::NOT_FOUND); + } + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | set | available | - | any | preferred | + */ + for (bool allow : {false, true}) { + auto settings = Settings().PreferredDatabase("db8").AllowFallback(allow); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(db); + UNIT_ASSERT_EQUAL(db->id(), "8"); + } + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | set | unavailable | - | false | UNAVAILABLE | + */ + DisableDatabase(dbInfos, 8); + auto settings = Settings().PreferredDatabase("db8").AllowFallback(false); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::UNAVAILABLE); + EnableDatabase(dbInfos, 8); + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | set | unavailable | - | true | by hash | + */ + DisableDatabase(dbInfos, 8); + auto settings = Settings().PreferredDatabase("db8").AllowFallback(true); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(db); + UNIT_ASSERT_UNEQUAL(db->id(), "8"); + UNIT_ASSERT_EQUAL(status, EStatus::SUCCESS); + EnableDatabase(dbInfos, 8); + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | unset | - | not found | false | NOT_FOUND | + */ + auto settings = Settings().AllowFallback(false); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc0"); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::NOT_FOUND); + } + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | unset | - | not found | true | by hash | + */ + auto settings = Settings().AllowFallback(true); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc0"); + UNIT_ASSERT(db); + UNIT_ASSERT_EQUAL(status, EStatus::SUCCESS); + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | unset | - | available | any | local | + */ + for (bool allow : {false, true}) { + auto settings = Settings().AllowFallback(allow); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(db); + UNIT_ASSERT_EQUAL(db->id(), "1"); + } + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | unset | - | unavailable | false | UNAVAILABLE | + */ + DisableDatabase(dbInfos, 1); + auto settings = Settings().AllowFallback(false); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(!db); + UNIT_ASSERT_EQUAL(status, EStatus::UNAVAILABLE); + EnableDatabase(dbInfos, 1); + } + + { + /* + | PreferredDb | Preferred state | Local state | AllowFallback | Return | + |-------------+-----------------+-------------+---------------+-------------| + | unset | - | unavailable | true | by hash | + */ + DisableDatabase(dbInfos, 1); + auto settings = Settings().AllowFallback(true); + auto [db, status] = SelectDatabaseImpl(settings, dbInfos, "dc1"); + UNIT_ASSERT(db); + UNIT_ASSERT_UNEQUAL(db->id(), "1"); + UNIT_ASSERT_EQUAL(status, EStatus::SUCCESS); + EnableDatabase(dbInfos, 1); + } + } + +} + +} diff --git a/src/client/federated_topic/ut/fds_mock.h b/src/client/federated_topic/ut/fds_mock.h new file mode 100644 index 0000000000..365cd0d962 --- /dev/null +++ b/src/client/federated_topic/ut/fds_mock.h @@ -0,0 +1,178 @@ +#pragma once + +#include "util/string/builder.h" +#include +#include + +#include +#include + +namespace NYdb::NFederatedTopic::NTests { + +// ctor gets ---list of response--- tmaybe{response} +// ListFederationDatabases gets 1 element under lock and respond. otherwise +// create 2 queues, for requests and responses +// getrequest() - put request and returns> +// sendresponse() +class TFederationDiscoveryServiceMock: public Ydb::FederationDiscovery::V1::FederationDiscoveryService::Service { +public: + using TRequest = Ydb::FederationDiscovery::ListFederationDatabasesRequest; + using TResponse = Ydb::FederationDiscovery::ListFederationDatabasesResponse; + + struct TGrpcResult { + TResponse Response; + grpc::Status Status; + }; + + struct TManualRequest { + const TRequest* Request; + NThreading::TPromise Result; + }; + +public: + ~TFederationDiscoveryServiceMock() { + while (auto r = GetNextPendingRequest()) { + r->Result.SetValue({}); + } + } + + std::optional GetNextPendingRequest() { + std::optional result; + with_lock(Lock) { + if (!PendingRequests.empty()) { + result = PendingRequests.front(); + PendingRequests.pop_front(); + } + } + return result; + } + + TManualRequest WaitNextPendingRequest() { + do { + auto result = GetNextPendingRequest(); + if (result.has_value()) { + return *result; + } + Sleep(TDuration::MilliSeconds(50)); + } while (true); + } + + virtual grpc::Status ListFederationDatabases(grpc::ServerContext*, + const TRequest* request, + TResponse* response) override { + Y_UNUSED(request); + + auto p = NThreading::NewPromise(); + auto f = p.GetFuture(); + + with_lock(Lock) { + PendingRequests.push_back({request, std::move(p)}); + } + + f.Wait(TDuration::Seconds(35)); + if (f.HasValue()) { + auto result = f.ExtractValueSync(); + Cerr << ">>> Ready to answer: " << (result.Status.ok() ? "ok" : "err") << Endl; + *response = std::move(result.Response); + return result.Status; + } + + return grpc::Status(grpc::StatusCode::UNKNOWN, "No response after timeout"); + } + + TGrpcResult ComposeOkResult(::Ydb::FederationDiscovery::DatabaseInfo::Status status) { + Ydb::FederationDiscovery::ListFederationDatabasesResponse okResponse; + + auto op = okResponse.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + okResponse.mutable_operation()->set_ready(true); + okResponse.mutable_operation()->set_id("12345"); + + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("fancy_datacenter"); + auto c1 = mockResult.add_federation_databases(); + c1->set_name("dc1"); + c1->set_path("/Root"); + c1->set_id("account-dc1"); + c1->set_endpoint("localhost:" + ToString(Port)); + c1->set_location("dc1"); + c1->set_status(status); + c1->set_weight(1000); + auto c2 = mockResult.add_federation_databases(); + c2->set_name("dc2"); + c2->set_path("/Root"); + c2->set_id("account-dc2"); + c2->set_endpoint("localhost:" + ToString(Port)); + c2->set_location("dc2"); + c2->set_status(status); + c2->set_weight(500); + auto c3 = mockResult.add_federation_databases(); + c3->set_name("dc3"); + c3->set_path("/Root"); + c3->set_id("account-dc3"); + c3->set_endpoint("localhost:" + ToString(Port)); + c3->set_location("dc3"); + c3->set_status(status); + c3->set_weight(500); + + op->mutable_result()->PackFrom(mockResult); + + return {okResponse, grpc::Status::OK}; + } + + TGrpcResult ComposeOkResultAvailableDatabases() { + return ComposeOkResult(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + } + + TGrpcResult ComposeOkResultUnavailableDatabases() { + return ComposeOkResult(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_UNAVAILABLE); + } + + TGrpcResult ComposeUnavailableResult() { + Ydb::FederationDiscovery::ListFederationDatabasesResponse response; + auto op = response.mutable_operation(); + op->set_status(Ydb::StatusIds::UNAVAILABLE); + response.mutable_operation()->set_ready(true); + response.mutable_operation()->set_id("12345"); + return {response, grpc::Status::OK}; + } + + TGrpcResult ComposeOkResultWithUnavailableDatabase(int unavailableDb) { + Ydb::FederationDiscovery::ListFederationDatabasesResponse okResponse; + + auto op = okResponse.mutable_operation(); + op->set_status(Ydb::StatusIds::SUCCESS); + okResponse.mutable_operation()->set_ready(true); + okResponse.mutable_operation()->set_id("12345"); + + Ydb::FederationDiscovery::ListFederationDatabasesResult mockResult; + mockResult.set_control_plane_endpoint("cp.logbroker-federation:2135"); + mockResult.set_self_location("fancy_datacenter"); + for (int i = 1; i <= 3; ++i) { + auto c1 = mockResult.add_federation_databases(); + c1->set_name(TStringBuilder() << "dc" << i); + c1->set_path("/Root"); + c1->set_id(TStringBuilder() << "account-dc" << i); + c1->set_endpoint("localhost:" + ToString(Port)); + c1->set_location(TStringBuilder() << "dc" << i); + if (i == unavailableDb) { + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_UNAVAILABLE); + } else { + c1->set_status(::Ydb::FederationDiscovery::DatabaseInfo::Status::DatabaseInfo_Status_AVAILABLE); + } + c1->set_weight(i == 0 ? 1000 : 500); + } + + op->mutable_result()->PackFrom(mockResult); + + return {okResponse, grpc::Status::OK}; + } + +public: + ui16 Port; + std::deque PendingRequests; + TAdaptiveLock Lock; +}; + +} // namespace NYdb::NFederatedTopic::NTests diff --git a/src/client/federated_topic/ut/ya.make b/src/client/federated_topic/ut/ya.make new file mode 100644 index 0000000000..8688c7519d --- /dev/null +++ b/src/client/federated_topic/ut/ya.make @@ -0,0 +1,38 @@ +UNITTEST_FOR(ydb/public/sdk/cpp/client/ydb_federated_topic) + +IF (SANITIZER_TYPE == "thread" OR WITH_VALGRIND) + TIMEOUT(1200) + SIZE(LARGE) + TAG(ya:fat) +ELSE() + TIMEOUT(600) + SIZE(MEDIUM) +ENDIF() + +FORK_SUBTESTS() + +PEERDIR( + library/cpp/testing/gmock_in_unittest + ydb/core/testlib/default + ydb/public/lib/json_value + ydb/public/lib/yson_value + ydb/public/sdk/cpp/client/ydb_driver + ydb/public/sdk/cpp/client/ydb_persqueue_public + ydb/public/sdk/cpp/client/ydb_persqueue_public/include + ydb/public/sdk/cpp/client/ydb_persqueue_public/ut/ut_utils + + ydb/public/sdk/cpp/client/ydb_topic/codecs + ydb/public/sdk/cpp/client/ydb_topic + ydb/public/sdk/cpp/client/ydb_topic/impl + + ydb/public/sdk/cpp/client/ydb_federated_topic + ydb/public/sdk/cpp/client/ydb_federated_topic/impl +) + +YQL_LAST_ABI_VERSION() + +SRCS( + basic_usage_ut.cpp +) + +END() From 8fe04eb324fddf659ab04b29d6dac820a0070f11 Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Fri, 23 Aug 2024 17:23:40 +0000 Subject: [PATCH 02/35] Moved commit "fix reordering of messages" from ydb repo --- src/client/topic/impl/write_session_impl.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index b0a37ef6b5..3fe4b79d29 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -846,6 +846,11 @@ void TWriteSessionImpl::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t co } } } + + for (auto& event : processResult.Events) { + EventsQueue->PushEvent(std::move(event)); + } + if (doRead) ReadFromProcessor(); @@ -860,9 +865,6 @@ void TWriteSessionImpl::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t co CloseImpl(std::move(errorStatus)); } } - for (auto& event : processResult.Events) { - EventsQueue->PushEvent(std::move(event)); - } if (needSetValue) { InitSeqNoPromise.SetValue(*processResult.InitSeqNo); processResult.HandleResult.DoSetSeqNo = false; // Redundant. Just in case. From 27c607bb7cb8f747bca2e23a1021fa4cb410d6c2 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy <79596613+GrigoriyPA@users.noreply.github.com> Date: Fri, 23 Aug 2024 17:28:33 +0000 Subject: [PATCH 03/35] Moved commit "Added add field for pool id into gRPC api" from ydb repo --- include/ydb-cpp-sdk/client/query/query.h | 2 ++ src/client/query/client.cpp | 1 + src/client/query/impl/exec_query.cpp | 1 + 3 files changed, 4 insertions(+) diff --git a/include/ydb-cpp-sdk/client/query/query.h b/include/ydb-cpp-sdk/client/query/query.h index bdb01e5b79..7a07eb53fc 100644 --- a/include/ydb-cpp-sdk/client/query/query.h +++ b/include/ydb-cpp-sdk/client/query/query.h @@ -74,6 +74,7 @@ struct TExecuteQuerySettings : public TRequestSettings { FLUENT_SETTING_DEFAULT(EExecMode, ExecMode, EExecMode::Execute); FLUENT_SETTING_DEFAULT(EStatsMode, StatsMode, EStatsMode::None); FLUENT_SETTING_OPTIONAL(bool, ConcurrentResultSets); + FLUENT_SETTING(std::string, PoolId); }; struct TBeginTxSettings : public TRequestSettings {}; @@ -97,6 +98,7 @@ struct TExecuteScriptSettings : public TOperationRequestSettings, public auto request = MakeOperationRequest(settings); request.set_exec_mode(::Ydb::Query::ExecMode(settings.ExecMode_)); request.set_stats_mode(::Ydb::Query::StatsMode(settings.StatsMode_)); + request.set_pool_id(settings.PoolId_); request.mutable_script_content()->set_syntax(::Ydb::Query::Syntax(settings.Syntax_)); request.mutable_script_content()->set_text(TStringType{script}); SetDuration(settings.ResultsTtl_, *request.mutable_results_ttl()); diff --git a/src/client/query/impl/exec_query.cpp b/src/client/query/impl/exec_query.cpp index 63dcf24a3c..a570cb72ea 100644 --- a/src/client/query/impl/exec_query.cpp +++ b/src/client/query/impl/exec_query.cpp @@ -212,6 +212,7 @@ TFuture> StreamExecuteQueryIm auto request = MakeRequest(); request.set_exec_mode(::Ydb::Query::ExecMode(settings.ExecMode_)); request.set_stats_mode(::Ydb::Query::StatsMode(settings.StatsMode_)); + request.set_pool_id(settings.PoolId_); request.mutable_query_content()->set_text(TStringType{query}); request.mutable_query_content()->set_syntax(::Ydb::Query::Syntax(settings.Syntax_)); if (session.has_value()) { From a9953921927b932cb3e065bc9f139ac799f0e317 Mon Sep 17 00:00:00 2001 From: Ilnaz Nizametdinov Date: Fri, 23 Aug 2024 17:35:09 +0000 Subject: [PATCH 04/35] Moved commit "Report & show replication lag" from ydb repo --- .../ydb-cpp-sdk/client/draft/ydb_replication.h | 14 +++++++++++++- src/client/draft/ydb_replication.cpp | 17 ++++++++++++++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/include/ydb-cpp-sdk/client/draft/ydb_replication.h b/include/ydb-cpp-sdk/client/draft/ydb_replication.h index 4ae8019484..51cb072890 100644 --- a/include/ydb-cpp-sdk/client/draft/ydb_replication.h +++ b/include/ydb-cpp-sdk/client/draft/ydb_replication.h @@ -5,6 +5,8 @@ #include +#include + namespace Ydb::Replication { class ConnectionParams; class DescribeReplicationResult; @@ -56,7 +58,17 @@ class TConnectionParams: private TCommonClientSettings { > Credentials_; }; -struct TRunningState {}; +struct TRunningState { +public: + TRunningState() = default; + explicit TRunningState(const std::optional& lag); + + const std::optional& GetLag() const; + +private: + std::optional Lag_; +}; + struct TDoneState {}; class TErrorState { diff --git a/src/client/draft/ydb_replication.cpp b/src/client/draft/ydb_replication.cpp index 8e7fccf96f..b569bcaefc 100644 --- a/src/client/draft/ydb_replication.cpp +++ b/src/client/draft/ydb_replication.cpp @@ -10,6 +10,7 @@ #include #include +#include #include namespace NYdb { @@ -58,6 +59,15 @@ const TOAuthCredentials& TConnectionParams::GetOAuthCredentials() const { return std::get(Credentials_); } +TRunningState::TRunningState(const std::optional& lag) + : Lag_(lag) +{ +} + +const std::optional& TRunningState::GetLag() const { + return Lag_; +} + class TErrorState::TImpl { public: NYql::TIssues Issues; @@ -77,6 +87,10 @@ const NYql::TIssues& TErrorState::GetIssues() const { return Impl_->Issues; } +TDuration DurationToDuration(const google::protobuf::Duration& value) { + return TDuration::MilliSeconds(google::protobuf::util::TimeUtil::DurationToMilliseconds(value)); +} + template NYql::TIssues IssuesFromMessage(const ::google::protobuf::RepeatedPtrField& message) { NYql::TIssues issues; @@ -100,7 +114,8 @@ TReplicationDescription::TReplicationDescription(const Ydb::Replication::Describ switch (desc.state_case()) { case Ydb::Replication::DescribeReplicationResult::kRunning: - State_ = TRunningState(); + State_ = TRunningState(desc.running().has_lag() + ? std::make_optional(DurationToDuration(desc.running().lag())) : std::nullopt); break; case Ydb::Replication::DescribeReplicationResult::kError: From edcad847eb17a8025ce6916085d464b22fcf5c1b Mon Sep 17 00:00:00 2001 From: Ilia Shakhov Date: Fri, 23 Aug 2024 17:42:46 +0000 Subject: [PATCH 05/35] Moved commit "Add start and end times to Operations API" from ydb repo --- .../client/types/operation/operation.h | 5 ++++ src/client/export/export.cpp | 6 ----- src/client/import/import.cpp | 6 ----- src/client/types/operation/operation.cpp | 26 +++++++++++++++++++ 4 files changed, 31 insertions(+), 12 deletions(-) diff --git a/include/ydb-cpp-sdk/client/types/operation/operation.h b/include/ydb-cpp-sdk/client/types/operation/operation.h index 722bf947a5..f8f01891ad 100644 --- a/include/ydb-cpp-sdk/client/types/operation/operation.h +++ b/include/ydb-cpp-sdk/client/types/operation/operation.h @@ -5,6 +5,7 @@ #include #include +#include #include namespace Ydb { @@ -31,6 +32,8 @@ class TOperation { const TOperationId& Id() const; bool Ready() const; const TStatus& Status() const; + TInstant StartTime() const; + TInstant EndTime() const; std::string ToString() const; std::string ToJsonString() const; @@ -46,4 +49,6 @@ class TOperation { using TAsyncOperation = NThreading::TFuture; +TInstant ProtoTimestampToInstant(const google::protobuf::Timestamp& timestamp); + } // namespace NYdb diff --git a/src/client/export/export.cpp b/src/client/export/export.cpp index 44d952d889..8760903181 100644 --- a/src/client/export/export.cpp +++ b/src/client/export/export.cpp @@ -24,12 +24,6 @@ using namespace Ydb::Export; /// Common namespace { -TInstant ProtoTimestampToInstant(const google::protobuf::Timestamp& timestamp) { - ui64 us = timestamp.seconds() * 1000000; - us += timestamp.nanos() / 1000; - return TInstant::MicroSeconds(us); -} - std::vector ItemsProgressFromProto(const google::protobuf::RepeatedPtrField& proto) { std::vector result(proto.size()); diff --git a/src/client/import/import.cpp b/src/client/import/import.cpp index 64f352c844..1db747edd5 100644 --- a/src/client/import/import.cpp +++ b/src/client/import/import.cpp @@ -19,12 +19,6 @@ using namespace Ydb::Import; /// Common namespace { -TInstant ProtoTimestampToInstant(const google::protobuf::Timestamp& timestamp) { - ui64 us = timestamp.seconds() * 1000000; - us += timestamp.nanos() / 1000; - return TInstant::MicroSeconds(us); -} - std::vector ItemsProgressFromProto(const google::protobuf::RepeatedPtrField& proto) { std::vector result; result.reserve(proto.size()); diff --git a/src/client/types/operation/operation.cpp b/src/client/types/operation/operation.cpp index 92847a087a..78d3224169 100644 --- a/src/client/types/operation/operation.cpp +++ b/src/client/types/operation/operation.cpp @@ -21,6 +21,8 @@ class TOperation::TImpl { : Id_(operation.id(), true /* allowEmpty */) , Status_(std::move(status)) , Ready_(operation.ready()) + , StartTime_(ProtoTimestampToInstant(operation.start_time())) + , EndTime_(ProtoTimestampToInstant(operation.end_time())) , Operation_(std::move(operation)) { } @@ -37,6 +39,14 @@ class TOperation::TImpl { return Status_; } + TInstant StartTime() const { + return StartTime_; + } + + TInstant EndTime() const { + return EndTime_; + } + const Ydb::Operations::Operation& GetProto() const { return Operation_; } @@ -45,6 +55,8 @@ class TOperation::TImpl { const TOperationId Id_; const TStatus Status_; const bool Ready_; + const TInstant StartTime_; + const TInstant EndTime_; const Ydb::Operations::Operation Operation_; }; @@ -68,6 +80,14 @@ const TStatus& TOperation::Status() const { return Impl_->Status(); } +TInstant TOperation::StartTime() const { + return Impl_->StartTime(); +} + +TInstant TOperation::EndTime() const { + return Impl_->EndTime(); +} + std::string TOperation::ToString() const { TString result; TStringOutput out(result); @@ -92,4 +112,10 @@ const Ydb::Operations::Operation& TOperation::GetProto() const { return Impl_->GetProto(); } +TInstant ProtoTimestampToInstant(const google::protobuf::Timestamp& timestamp) { + ui64 us = timestamp.seconds() * 1000000; + us += timestamp.nanos() / 1000; + return TInstant::MicroSeconds(us); +} + } // namespace NYdb From 102877cc07e76e6f2b948638d14e9f6e9ff47666 Mon Sep 17 00:00:00 2001 From: Daniil Demin Date: Fri, 23 Aug 2024 17:58:35 +0000 Subject: [PATCH 06/35] Moved commit "Extend index tables' partitioning settings to be able to restore them from a backup with the same partitioning as before" from ydb repo --- include/ydb-cpp-sdk/client/table/table.h | 84 +++++++++++++------- src/client/table/table.cpp | 97 ++++++++++++++++++------ 2 files changed, 131 insertions(+), 50 deletions(-) diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 245b9f1742..eb0c9bd9a8 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -21,6 +21,8 @@ class CreateTableRequest; class Changefeed; class ChangefeedDescription; class DescribeTableResult; +class ExplicitPartitions; +class GlobalIndexSettings; class PartitioningSettings; class DateTypeColumnModeSettings; class TtlSettings; @@ -146,17 +148,67 @@ struct TAlterTableColumn { //////////////////////////////////////////////////////////////////////////////// +//! Represents table partitioning settings +class TPartitioningSettings { +public: + TPartitioningSettings(); + explicit TPartitioningSettings(const Ydb::Table::PartitioningSettings& proto); + + const Ydb::Table::PartitioningSettings& GetProto() const; + + std::optional GetPartitioningBySize() const; + std::optional GetPartitioningByLoad() const; + uint64_t GetPartitionSizeMb() const; + uint64_t GetMinPartitionsCount() const; + uint64_t GetMaxPartitionsCount() const; + +private: + class TImpl; + std::shared_ptr Impl_; +}; + +struct TExplicitPartitions { + using TSelf = TExplicitPartitions; + + FLUENT_SETTING_VECTOR(TValue, SplitPoints); + + template + static TExplicitPartitions FromProto(const TProto& proto); + + void SerializeTo(Ydb::Table::ExplicitPartitions& proto) const; +}; + +struct TGlobalIndexSettings { + using TUniformOrExplicitPartitions = std::variant; + + TPartitioningSettings PartitioningSettings; + TUniformOrExplicitPartitions Partitions; + + template + static TGlobalIndexSettings FromProto(const TProto& proto); + + void SerializeTo(Ydb::Table::GlobalIndexSettings& proto) const; +}; + //! Represents index description class TIndexDescription { friend class NYdb::TProtoAccessor; public: TIndexDescription( - const std::string& name, EIndexType type, + const std::string& name, + EIndexType type, const std::vector& indexColumns, - const std::vector& dataColumns = std::vector()); + const std::vector& dataColumns = {}, + const TGlobalIndexSettings& settings = {} + ); - TIndexDescription(const std::string& name, const std::vector& indexColumns, const std::vector& dataColumns = std::vector()); + TIndexDescription( + const std::string& name, + const std::vector& indexColumns, + const std::vector& dataColumns = {}, + const TGlobalIndexSettings& settings = {} + ); const std::string& GetIndexName() const; EIndexType GetIndexType() const; @@ -180,6 +232,7 @@ class TIndexDescription { EIndexType IndexType_; std::vector IndexColumns_; std::vector DataColumns_; + TGlobalIndexSettings GlobalIndexSettings_; uint64_t SizeBytes = 0; }; @@ -423,25 +476,6 @@ class TColumnFamilyDescription { std::shared_ptr Impl_; }; -//! Represents table partitioning settings -class TPartitioningSettings { -public: - TPartitioningSettings(); - explicit TPartitioningSettings(const Ydb::Table::PartitioningSettings& proto); - - const Ydb::Table::PartitioningSettings& GetProto() const; - - std::optional GetPartitioningBySize() const; - std::optional GetPartitioningByLoad() const; - uint64_t GetPartitionSizeMb() const; - uint64_t GetMinPartitionsCount() const; - uint64_t GetMaxPartitionsCount() const; - -private: - class TImpl; - std::shared_ptr Impl_; -}; - //! Represents table read replicas settings class TReadReplicasSettings { public: @@ -1219,12 +1253,6 @@ struct TStoragePolicy { FLUENT_SETTING_VECTOR(TColumnFamilyPolicy, ColumnFamilies); }; -struct TExplicitPartitions { - using TSelf = TExplicitPartitions; - - FLUENT_SETTING_VECTOR(TValue, SplitPoints); -}; - struct TPartitioningPolicy { using TSelf = TPartitioningPolicy; diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index ec6d575c2e..0317229f54 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -251,6 +251,24 @@ static void SerializeTo(const TRenameIndex& rename, Ydb::Table::RenameIndexItem& proto.set_replace_destination(rename.ReplaceDestination_); } +template +TExplicitPartitions TExplicitPartitions::FromProto(const TProto& proto) { + TExplicitPartitions out; + for (const auto& splitPoint : proto.split_points()) { + TValue value(TType(splitPoint.type()), splitPoint.value()); + out.AppendSplitPoints(value); + } + return out; +} + +void TExplicitPartitions::SerializeTo(Ydb::Table::ExplicitPartitions& proto) const { + for (const auto& splitPoint : SplitPoints_) { + auto* boundary = proto.add_split_points(); + boundary->mutable_type()->CopyFrom(TProtoAccessor::GetProto(splitPoint.GetType())); + boundary->mutable_value()->CopyFrom(TProtoAccessor::GetProto(splitPoint)); + } +} + class TTableDescription::TImpl { using EUnit = TValueSinceUnixEpochModeSettings::EUnit; @@ -419,13 +437,7 @@ class TTableDescription::TImpl { break; case Ydb::Table::CreateTableRequest::kPartitionAtKeys: { - TExplicitPartitions partitionAtKeys; - for (const auto& splitPoint : request.partition_at_keys().split_points()) { - TValue value(TType(splitPoint.type()), splitPoint.value()); - partitionAtKeys.AppendSplitPoints(value); - } - - SetPartitionAtKeys(partitionAtKeys); + SetPartitionAtKeys(TExplicitPartitions::FromProto(request.partition_at_keys())); break; } @@ -922,12 +934,7 @@ void TTableDescription::SerializeTo(Ydb::Table::CreateTableRequest& request) con } if (const auto& partitionAtKeys = Impl_->GetPartitionAtKeys()) { - auto* borders = request.mutable_partition_at_keys(); - for (const auto& splitPoint : partitionAtKeys->SplitPoints_) { - auto* border = borders->add_split_points(); - border->mutable_type()->CopyFrom(TProtoAccessor::GetProto(splitPoint.GetType())); - border->mutable_value()->CopyFrom(TProtoAccessor::GetProto(splitPoint)); - } + partitionAtKeys->SerializeTo(*request.mutable_partition_at_keys()); } else if (Impl_->GetProto().shard_key_bounds_size()) { request.mutable_partition_at_keys()->mutable_split_points()->CopyFrom(Impl_->GetProto().shard_key_bounds()); } @@ -2204,16 +2211,25 @@ bool TRenameItem::ReplaceDestination() const { //////////////////////////////////////////////////////////////////////////////// -TIndexDescription::TIndexDescription(const std::string& name, EIndexType type, - const std::vector& indexColumns, const std::vector& dataColumns) - : IndexName_(name) +TIndexDescription::TIndexDescription( + const std::string& name, + EIndexType type, + const std::vector& indexColumns, + const std::vector& dataColumns, + const TGlobalIndexSettings& settings +) : IndexName_(name) , IndexType_(type) , IndexColumns_(indexColumns) , DataColumns_(dataColumns) + , GlobalIndexSettings_(settings) {} -TIndexDescription::TIndexDescription(const std::string& name, const std::vector& indexColumns, const std::vector& dataColumns) - : TIndexDescription(name, EIndexType::GlobalSync, indexColumns, dataColumns) +TIndexDescription::TIndexDescription( + const std::string& name, + const std::vector& indexColumns, + const std::vector& dataColumns, + const TGlobalIndexSettings& settings +) : TIndexDescription(name, EIndexType::GlobalSync, indexColumns, dataColumns, settings) {} TIndexDescription::TIndexDescription(const Ydb::Table::TableIndex& tableIndex) @@ -2244,11 +2260,45 @@ uint64_t TIndexDescription::GetSizeBytes() const { return SizeBytes; } +template +TGlobalIndexSettings TGlobalIndexSettings::FromProto(const TProto& proto) { + auto partitionsFromProto = [](const auto& proto) -> TUniformOrExplicitPartitions { + switch (proto.partitions_case()) { + case TProto::kUniformPartitions: + return proto.uniform_partitions(); + case TProto::kPartitionAtKeys: + return TExplicitPartitions::FromProto(proto.partition_at_keys()); + default: + return {}; + } + }; + + return { + .PartitioningSettings = TPartitioningSettings(proto.partitioning_settings()), + .Partitions = partitionsFromProto(proto) + }; +} + +void TGlobalIndexSettings::SerializeTo(Ydb::Table::GlobalIndexSettings& settings) const { + *settings.mutable_partitioning_settings() = PartitioningSettings.GetProto(); + + auto variantVisitor = [&settings](auto&& partitions) { + using T = std::decay_t; + if constexpr (std::is_same_v) { + settings.set_uniform_partitions(partitions); + } else if constexpr (std::is_same_v) { + partitions.SerializeTo(*settings.mutable_partition_at_keys()); + } + }; + std::visit(std::move(variantVisitor), Partitions); +} + template TIndexDescription TIndexDescription::FromProto(const TProto& proto) { EIndexType type; std::vector indexColumns; std::vector dataColumns; + TGlobalIndexSettings globalIndexSettings; indexColumns.assign(proto.index_columns().begin(), proto.index_columns().end()); dataColumns.assign(proto.data_columns().begin(), proto.data_columns().end()); @@ -2256,19 +2306,22 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { switch (proto.type_case()) { case TProto::kGlobalIndex: type = EIndexType::GlobalSync; + globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_index().settings()); break; case TProto::kGlobalAsyncIndex: type = EIndexType::GlobalAsync; + globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_async_index().settings()); break; case TProto::kGlobalUniqueIndex: type = EIndexType::GlobalUnique; + globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_unique_index().settings()); break; default: // fallback to global sync type = EIndexType::GlobalSync; break; } - auto result = TIndexDescription(proto.name(), type, indexColumns, dataColumns); + auto result = TIndexDescription(proto.name(), type, indexColumns, dataColumns, globalIndexSettings); if constexpr (std::is_same_v) { result.SizeBytes = proto.size_bytes(); } @@ -2286,13 +2339,13 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { switch (IndexType_) { case EIndexType::GlobalSync: - *proto.mutable_global_index() = Ydb::Table::GlobalIndex(); + GlobalIndexSettings_.SerializeTo(*proto.mutable_global_index()->mutable_settings()); break; case EIndexType::GlobalAsync: - *proto.mutable_global_async_index() = Ydb::Table::GlobalAsyncIndex(); + GlobalIndexSettings_.SerializeTo(*proto.mutable_global_async_index()->mutable_settings()); break; case EIndexType::GlobalUnique: - *proto.mutable_global_unique_index() = Ydb::Table::GlobalUniqueIndex(); + GlobalIndexSettings_.SerializeTo(*proto.mutable_global_unique_index()->mutable_settings()); break; case EIndexType::Unknown: break; From 8e429b87cda4ac5d27228d4cb570ee874d41d35e Mon Sep 17 00:00:00 2001 From: qyryq Date: Fri, 23 Aug 2024 18:17:01 +0000 Subject: [PATCH 07/35] Moved commit "Federated write session: fix deadlock" from ydb repo --- .../impl/federated_write_session.cpp | 13 +++++++++++-- .../federated_topic/impl/federated_write_session.h | 1 + 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index ca2ab80d21..ce41c0e318 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -19,7 +19,7 @@ bool DatabasesAreSame(std::shared_ptr lhs, std::shared_ptr rhs if (!lhs || !rhs) { return false; } - return lhs->path() == rhs->path() && lhs->endpoint() == rhs->endpoint(); + return lhs->name() == rhs->name() && lhs->path() == rhs->path() && lhs->endpoint() == rhs->endpoint(); } NTopic::TTopicClientSettings FromFederated(const TFederatedTopicClientSettings& settings); @@ -107,8 +107,10 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) Y_ABORT_UNLESS(Lock.IsLocked()); if (Subsession) { PendingToken.reset(); - Subsession->Close(TDuration::Zero()); + OldSubsession = std::move(Subsession); + OldSubsession->Close(TDuration::Zero()); } + auto clientSettings = SubclientSettings; clientSettings .Database(db->path()) @@ -148,6 +150,11 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) } }) .SessionClosedHandler([selfCtx = SelfContext](const NTopic::TSessionClosedEvent & ev) { + if (ev.IsSuccess()) { + // The subsession was closed by the federated write session itself while creating a new subsession. + // In this case we get SUCCESS status and don't need to propagate it further. + return; + } if (auto self = selfCtx->LockShared()) { with_lock(self->Lock) { self->CloseImpl(ev); @@ -291,8 +298,10 @@ void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration del auto cb = [selfCtx = SelfContext](bool ok) { if (ok) { if (auto self = selfCtx->LockShared()) { + std::shared_ptr old; with_lock(self->Lock) { self->UpdateFederationStateImpl(); + old = std::move(self->OldSubsession); } } } diff --git a/src/client/federated_topic/impl/federated_write_session.h b/src/client/federated_topic/impl/federated_write_session.h index af8f1b7a70..8df0a4d735 100644 --- a/src/client/federated_topic/impl/federated_write_session.h +++ b/src/client/federated_topic/impl/federated_write_session.h @@ -122,6 +122,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, TAdaptiveLock Lock; std::shared_ptr Subsession; + std::shared_ptr OldSubsession; std::shared_ptr ClientEventsQueue; From 24f16414e1a2de3050075b3d5597418897f83982 Mon Sep 17 00:00:00 2001 From: Ilnaz Nizametdinov Date: Fri, 23 Aug 2024 18:28:12 +0000 Subject: [PATCH 08/35] Moved commit "CDC Initial Scan progress" from ydb repo --- include/ydb-cpp-sdk/client/table/table.h | 18 +++++++++++- src/client/table/table.cpp | 36 ++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index eb0c9bd9a8..44440346a5 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -264,10 +264,24 @@ class TBuildIndexOperation : public TOperation { //////////////////////////////////////////////////////////////////////////////// -//! Represents index description +//! Represents changefeed description class TChangefeedDescription { friend class NYdb::TProtoAccessor; +public: + class TInitialScanProgress { + public: + explicit TInitialScanProgress(uint32_t total, uint32_t completed); + + uint32_t GetPartsTotal() const; + uint32_t GetPartsCompleted() const; + float GetProgress() const; // percentage + + private: + uint32_t PartsTotal; + uint32_t PartsCompleted; + }; + public: TChangefeedDescription(const std::string& name, EChangefeedMode mode, EChangefeedFormat format); @@ -295,6 +309,7 @@ class TChangefeedDescription { bool GetInitialScan() const; const std::unordered_map& GetAttributes() const; const std::string& GetAwsRegion() const; + const std::optional& GetInitialScanProgress() const; void SerializeTo(Ydb::Table::Changefeed& proto) const; std::string ToString() const; @@ -318,6 +333,7 @@ class TChangefeedDescription { bool InitialScan_ = false; std::unordered_map Attributes_; std::string AwsRegion_; + std::optional InitialScanProgress_; }; bool operator==(const TChangefeedDescription& lhs, const TChangefeedDescription& rhs); diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 0317229f54..4dbbad239b 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -2398,6 +2398,27 @@ TChangefeedDescription::TChangefeedDescription(const Ydb::Table::ChangefeedDescr : TChangefeedDescription(FromProto(proto)) {} +TChangefeedDescription::TInitialScanProgress::TInitialScanProgress(uint32_t total, uint32_t completed) + : PartsTotal(total) + , PartsCompleted(completed) +{} + +uint32_t TChangefeedDescription::TInitialScanProgress::GetPartsTotal() const { + return PartsTotal; +} + +uint32_t TChangefeedDescription::TInitialScanProgress::GetPartsCompleted() const { + return PartsCompleted; +} + +float TChangefeedDescription::TInitialScanProgress::GetProgress() const { + if (PartsTotal == 0) { + return 0; + } + + return 100 * float(PartsCompleted) / float(PartsTotal); +} + TChangefeedDescription& TChangefeedDescription::WithVirtualTimestamps() { VirtualTimestamps_ = true; return *this; @@ -2474,6 +2495,10 @@ const std::string& TChangefeedDescription::GetAwsRegion() const { return AwsRegion_; } +const std::optional& TChangefeedDescription::GetInitialScanProgress() const { + return InitialScanProgress_; +} + template TChangefeedDescription TChangefeedDescription::FromProto(const TProto& proto) { EChangefeedMode mode; @@ -2541,6 +2566,13 @@ TChangefeedDescription TChangefeedDescription::FromProto(const TProto& proto) { ret.State_ = EChangefeedState::Unknown; break; } + + if (proto.has_initial_scan_progress()) { + ret.InitialScanProgress_ = std::make_optional( + proto.initial_scan_progress().parts_total(), + proto.initial_scan_progress().parts_completed() + ); + } } for (const auto& [key, value] : proto.attributes()) { @@ -2628,6 +2660,10 @@ void TChangefeedDescription::Out(IOutputStream& o) const { o << ", aws_region: " << AwsRegion_; } + if (InitialScanProgress_) { + o << ", initial_scan_progress: " << InitialScanProgress_->GetProgress() << "%"; + } + o << " }"; } From e584d377d1184d5d3ee2e986151e4b0b1a18c1a3 Mon Sep 17 00:00:00 2001 From: niksaveliev Date: Fri, 23 Aug 2024 18:42:38 +0000 Subject: [PATCH 09/35] Moved commit "PQ_V1 SDK auto partitioning support" from ydb repo --- .../ydb-cpp-sdk/client/topic/control_plane.h | 2 +- .../persqueue_public/impl/persqueue.cpp | 13 +++++++-- .../persqueue_public/impl/persqueue_impl.h | 28 +++++++++++++++++++ .../persqueue_public/include/control_plane.h | 26 +++++++++++++++++ 4 files changed, 66 insertions(+), 3 deletions(-) diff --git a/include/ydb-cpp-sdk/client/topic/control_plane.h b/include/ydb-cpp-sdk/client/topic/control_plane.h index 3369f8c984..4fbc92ad4a 100644 --- a/include/ydb-cpp-sdk/client/topic/control_plane.h +++ b/include/ydb-cpp-sdk/client/topic/control_plane.h @@ -160,7 +160,7 @@ struct TAlterTopicSettings; struct TAutoPartitioningSettings { friend struct TAutoPartitioningSettingsBuilder; public: -TAutoPartitioningSettings() + TAutoPartitioningSettings() : Strategy_(EAutoPartitioningStrategy::Disabled) , StabilizationWindow_(TDuration::Seconds(0)) , DownUtilizationPercent_(0) diff --git a/src/client/persqueue_public/impl/persqueue.cpp b/src/client/persqueue_public/impl/persqueue.cpp index 366e54170b..82f5c7332f 100644 --- a/src/client/persqueue_public/impl/persqueue.cpp +++ b/src/client/persqueue_public/impl/persqueue.cpp @@ -85,11 +85,20 @@ TDescribeTopicResult::TDescribeTopicResult(TStatus status, const Ydb::PersQueue: } TDescribeTopicResult::TTopicSettings::TTopicSettings(const Ydb::PersQueue::V1::TopicSettings& settings) { - - PartitionsCount_ = settings.partitions_count(); RetentionPeriod_ = TDuration::MilliSeconds(settings.retention_period_ms()); SupportedFormat_ = static_cast(settings.supported_format()); + if (settings.has_auto_partitioning_settings()) { + PartitionsCount_ = settings.auto_partitioning_settings().min_active_partitions(); + MaxPartitionsCount_ = settings.auto_partitioning_settings().max_active_partitions(); + StabilizationWindow_ = TDuration::Seconds(settings.auto_partitioning_settings().partition_write_speed().stabilization_window().seconds()); + UpUtilizationPercent_ = settings.auto_partitioning_settings().partition_write_speed().up_utilization_percent(); + DownUtilizationPercent_ = settings.auto_partitioning_settings().partition_write_speed().down_utilization_percent(); + AutoPartitioningStrategy_ = settings.auto_partitioning_settings().strategy(); + } else { + PartitionsCount_ = settings.partitions_count(); + } + for (const auto& codec : settings.supported_codecs()) { SupportedCodecs_.push_back(static_cast(codec)); } diff --git a/src/client/persqueue_public/impl/persqueue_impl.h b/src/client/persqueue_public/impl/persqueue_impl.h index a0edce9ca0..a653e3c328 100644 --- a/src/client/persqueue_public/impl/persqueue_impl.h +++ b/src/client/persqueue_public/impl/persqueue_impl.h @@ -55,6 +55,34 @@ class TPersQueueClient::TImpl : public TClientImplCommonset_max_active_partitions(settings.PartitionsCount_); + autoPartitioningSettingsDefined = true; + } + if (settings.AutoPartitioningStrategy_.has_value()) { + props.mutable_auto_partitioning_settings()->set_strategy(*settings.AutoPartitioningStrategy_); + autoPartitioningSettingsDefined = true; + } + if (settings.DownUtilizationPercent_.has_value()) { + props.mutable_auto_partitioning_settings()->mutable_partition_write_speed()->set_down_utilization_percent(*settings.DownUtilizationPercent_); + autoPartitioningSettingsDefined = true; + } + if (settings.UpUtilizationPercent_.has_value()) { + props.mutable_auto_partitioning_settings()->mutable_partition_write_speed()->set_up_utilization_percent(*settings.UpUtilizationPercent_); + autoPartitioningSettingsDefined = true; + } + if (settings.StabilizationWindow_.has_value()) { + props.mutable_auto_partitioning_settings()->mutable_partition_write_speed()->mutable_stabilization_window()->set_seconds((*settings.StabilizationWindow_).Seconds()); + autoPartitioningSettingsDefined = true; + } + + if (!autoPartitioningSettingsDefined) { + props.set_partitions_count(settings.PartitionsCount_); + } else { + props.mutable_auto_partitioning_settings()->set_min_active_partitions(settings.PartitionsCount_); + } + props.set_retention_period_ms(settings.RetentionPeriod_.MilliSeconds()); props.set_supported_format(static_cast(settings.SupportedFormat_)); for (const auto& codec : settings.SupportedCodecs_) { diff --git a/src/client/persqueue_public/include/control_plane.h b/src/client/persqueue_public/include/control_plane.h index c8cc2ec7e3..0511fb8e50 100644 --- a/src/client/persqueue_public/include/control_plane.h +++ b/src/client/persqueue_public/include/control_plane.h @@ -117,6 +117,12 @@ struct TDescribeTopicResult : public TStatus { } GETTER(std::optional, RemoteMirrorRule); + GETTER(std::optional, MaxPartitionsCount); + GETTER(std::optional, StabilizationWindow); + GETTER(std::optional, UpUtilizationPercent); + GETTER(std::optional, DownUtilizationPercent); + GETTER(std::optional, AutoPartitioningStrategy); + #undef GETTER @@ -138,6 +144,12 @@ struct TDescribeTopicResult : public TStatus { std::optional AbcId_; std::optional AbcSlug_; std::string FederationAccount_; + + std::optional MaxPartitionsCount_; + std::optional StabilizationWindow_; + std::optional UpUtilizationPercent_; + std::optional DownUtilizationPercent_; + std::optional AutoPartitioningStrategy_; }; TDescribeTopicResult(TStatus status, const Ydb::PersQueue::V1::DescribeTopicResult& result); @@ -191,6 +203,7 @@ struct TReadRuleSettings { // Settings for topic. template struct TTopicSettings : public TOperationRequestSettings { + friend class TPersQueueClient; struct TRemoteMirrorRuleSettings { TRemoteMirrorRuleSettings() {} @@ -266,9 +279,22 @@ struct TTopicSettings : public TOperationRequestSettings { if (settings.RemoteMirrorRule()) { RemoteMirrorRule_ = TRemoteMirrorRuleSettings().SetSettings(settings.RemoteMirrorRule().value()); } + + MaxPartitionsCount_ = settings.MaxPartitionsCount(); + StabilizationWindow_ = settings.StabilizationWindow(); + UpUtilizationPercent_ = settings.UpUtilizationPercent(); + DownUtilizationPercent_ = settings.DownUtilizationPercent(); + AutoPartitioningStrategy_ = settings.AutoPartitioningStrategy(); + return static_cast(*this); } +private: + std::optional MaxPartitionsCount_; + std::optional StabilizationWindow_; + std::optional UpUtilizationPercent_; + std::optional DownUtilizationPercent_; + std::optional AutoPartitioningStrategy_; }; From 6984ae678ad4ee1d0c37666f0d824492bf3e0d32 Mon Sep 17 00:00:00 2001 From: qyryq Date: Fri, 23 Aug 2024 18:57:15 +0000 Subject: [PATCH 10/35] Moved commit "Federated Topic: Log on federation discovery failure" from ydb repo --- .../impl/federation_observer.cpp | 22 +++++++++---------- src/client/topic/impl/read_session_impl.ipp | 2 +- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/client/federated_topic/impl/federation_observer.cpp b/src/client/federated_topic/impl/federation_observer.cpp index 89a9c2d134..53ec045325 100644 --- a/src/client/federated_topic/impl/federation_observer.cpp +++ b/src/client/federated_topic/impl/federation_observer.cpp @@ -147,27 +147,27 @@ void TFederatedDbObserverImpl::OnFederationDiscovery(TStatus&& status, Ydb::Fede FederatedDbState->DbInfos.emplace_back(std::move(db)); } else { - if (!status.IsSuccess()) { + if (status.IsSuccess()) { + ScheduleFederationDiscoveryImpl(REDISCOVERY_DELAY); + } else { + LOG_LAZY(DbDriverState_->Log, TLOG_ERR, TStringBuilder() + << "OnFederationDiscovery: Got error. Status: " << status.GetStatus() + << ". Description: " << status.GetIssues().ToOneLineString()); if (!FederationDiscoveryRetryState) { FederationDiscoveryRetryState = FederationDiscoveryRetryPolicy->CreateRetryState(); } - auto retryDelay = FederationDiscoveryRetryState->GetNextRetryDelay(status.GetStatus()); - - if (retryDelay) { - ScheduleFederationDiscoveryImpl(*retryDelay); + if (auto d = FederationDiscoveryRetryState->GetNextRetryDelay(status.GetStatus())) { + ScheduleFederationDiscoveryImpl(*d); return; } - // If retryDelay is Nothing, meaning there won't be another retry, - // we replace FederatedDbState with the unsuccessful one and then set the PromiseToInitState if needed, - // and the observer becomes stale (see IsStale method). - } else { - ScheduleFederationDiscoveryImpl(REDISCOVERY_DELAY); + // If there won't be another retry, we replace FederatedDbState with the unsuccessful one + // and set the PromiseToInitState to make the observer stale (see IsStale method). } // TODO validate new state and check if differs from previous - auto newInfo = std::make_shared(std::move(result), std::move(status)); + // TODO update only if new state differs std::swap(FederatedDbState, newInfo); } diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index fb070e3b7a..c05b85b2be 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -258,7 +258,7 @@ bool TSingleClusterReadSessionImpl::Reconnect(const TPlain if (!status.Ok()) { LOG_LAZY(Log, TLOG_ERR, GetLogPrefix() << "Got error. Status: " << status.Status - << ". Description: " << IssuesSingleLineString(status.Issues)); + << ". Description: " << IssuesSingleLineString(status.Issues)); } NYdbGrpc::IQueueClientContextPtr delayContext = nullptr; From c65fe0f255401d0783d1eb4e3bb73cefd691a9a6 Mon Sep 17 00:00:00 2001 From: Ilnaz Nizametdinov Date: Fri, 23 Aug 2024 19:07:45 +0000 Subject: [PATCH 11/35] Moved commit "Replication stats (total & per item): lag, initial scan progress" from ydb repo --- .../client/draft/ydb_replication.h | 27 ++++++++++++--- include/ydb-cpp-sdk/client/table/table.h | 3 ++ src/client/draft/ydb_replication.cpp | 33 ++++++++++++++----- src/client/table/table.cpp | 11 +++++++ 4 files changed, 61 insertions(+), 13 deletions(-) diff --git a/include/ydb-cpp-sdk/client/draft/ydb_replication.h b/include/ydb-cpp-sdk/client/draft/ydb_replication.h index 51cb072890..749e00f5e4 100644 --- a/include/ydb-cpp-sdk/client/draft/ydb_replication.h +++ b/include/ydb-cpp-sdk/client/draft/ydb_replication.h @@ -10,6 +10,7 @@ namespace Ydb::Replication { class ConnectionParams; class DescribeReplicationResult; + class DescribeReplicationResult_Stats; } namespace NYdb { @@ -24,7 +25,11 @@ namespace NYdb::NReplication { class TDescribeReplicationResult; using TAsyncDescribeReplicationResult = NThreading::TFuture; -struct TDescribeReplicationSettings: public TOperationRequestSettings {}; + +struct TDescribeReplicationSettings: public TOperationRequestSettings { + using TSelf = TDescribeReplicationSettings; + FLUENT_SETTING_DEFAULT(bool, IncludeStats, false); +}; struct TStaticCredentials { std::string User; @@ -58,15 +63,28 @@ class TConnectionParams: private TCommonClientSettings { > Credentials_; }; -struct TRunningState { +class TStats { public: - TRunningState() = default; - explicit TRunningState(const std::optional& lag); + TStats() = default; + TStats(const Ydb::Replication::DescribeReplicationResult_Stats& stats); const std::optional& GetLag() const; + const std::optional& GetInitialScanProgress() const; private: std::optional Lag_; + std::optional InitialScanProgress_; +}; + +class TRunningState { +public: + TRunningState() = default; + explicit TRunningState(const TStats& stats); + + const TStats& GetStats() const; + +private: + TStats Stats_; }; struct TDoneState {}; @@ -89,6 +107,7 @@ class TReplicationDescription { uint64_t Id; std::string SrcPath; std::string DstPath; + TStats Stats; std::optional SrcChangefeedName; }; diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 44440346a5..d7619197d7 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -271,8 +271,11 @@ class TChangefeedDescription { public: class TInitialScanProgress { public: + TInitialScanProgress(); explicit TInitialScanProgress(uint32_t total, uint32_t completed); + TInitialScanProgress& operator+=(const TInitialScanProgress& other); + uint32_t GetPartsTotal() const; uint32_t GetPartsCompleted() const; float GetProgress() const; // percentage diff --git a/src/client/draft/ydb_replication.cpp b/src/client/draft/ydb_replication.cpp index b569bcaefc..7e4ce3e0ce 100644 --- a/src/client/draft/ydb_replication.cpp +++ b/src/client/draft/ydb_replication.cpp @@ -59,15 +59,33 @@ const TOAuthCredentials& TConnectionParams::GetOAuthCredentials() const { return std::get(Credentials_); } -TRunningState::TRunningState(const std::optional& lag) - : Lag_(lag) +static TDuration DurationToDuration(const google::protobuf::Duration& value) { + return TDuration::MilliSeconds(google::protobuf::util::TimeUtil::DurationToMilliseconds(value)); +} + +TStats::TStats(const Ydb::Replication::DescribeReplicationResult_Stats& stats) + : Lag_(stats.has_lag() ? std::make_optional(DurationToDuration(stats.lag())) : std::nullopt) + , InitialScanProgress_(stats.has_initial_scan_progress() ? std::make_optional(stats.initial_scan_progress()) : std::nullopt) { } -const std::optional& TRunningState::GetLag() const { +const std::optional& TStats::GetLag() const { return Lag_; } +const std::optional& TStats::GetInitialScanProgress() const { + return InitialScanProgress_; +} + +TRunningState::TRunningState(const TStats& stats) + : Stats_(stats) +{ +} + +const TStats& TRunningState::GetStats() const { + return Stats_; +} + class TErrorState::TImpl { public: NYql::TIssues Issues; @@ -87,10 +105,6 @@ const NYql::TIssues& TErrorState::GetIssues() const { return Impl_->Issues; } -TDuration DurationToDuration(const google::protobuf::Duration& value) { - return TDuration::MilliSeconds(google::protobuf::util::TimeUtil::DurationToMilliseconds(value)); -} - template NYql::TIssues IssuesFromMessage(const ::google::protobuf::RepeatedPtrField& message) { NYql::TIssues issues; @@ -107,6 +121,7 @@ TReplicationDescription::TReplicationDescription(const Ydb::Replication::Describ .Id = item.id(), .SrcPath = item.source_path(), .DstPath = item.destination_path(), + .Stats = TStats(item.stats()), .SrcChangefeedName = item.has_source_changefeed_name() ? std::make_optional(item.source_changefeed_name()) : std::nullopt, }); @@ -114,8 +129,7 @@ TReplicationDescription::TReplicationDescription(const Ydb::Replication::Describ switch (desc.state_case()) { case Ydb::Replication::DescribeReplicationResult::kRunning: - State_ = TRunningState(desc.running().has_lag() - ? std::make_optional(DurationToDuration(desc.running().lag())) : std::nullopt); + State_ = TRunningState(desc.running().stats()); break; case Ydb::Replication::DescribeReplicationResult::kError: @@ -183,6 +197,7 @@ class TReplicationClient::TImpl: public TClientImplCommon(settings); request.set_path(TStringType{path}); + request.set_include_stats(settings.IncludeStats_); auto promise = NThreading::NewPromise(); diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 4dbbad239b..bfe83e664f 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -2398,11 +2398,22 @@ TChangefeedDescription::TChangefeedDescription(const Ydb::Table::ChangefeedDescr : TChangefeedDescription(FromProto(proto)) {} +TChangefeedDescription::TInitialScanProgress::TInitialScanProgress() + : PartsTotal(0) + , PartsCompleted(0) +{} + TChangefeedDescription::TInitialScanProgress::TInitialScanProgress(uint32_t total, uint32_t completed) : PartsTotal(total) , PartsCompleted(completed) {} +TChangefeedDescription::TInitialScanProgress& TChangefeedDescription::TInitialScanProgress::operator+=(const TInitialScanProgress& other) { + PartsTotal += other.PartsTotal; + PartsCompleted += other.PartsCompleted; + return *this; +} + uint32_t TChangefeedDescription::TInitialScanProgress::GetPartsTotal() const { return PartsTotal; } From f4b2e42e417240d7f89a37385696afd674ee6de7 Mon Sep 17 00:00:00 2001 From: qyryq Date: Fri, 23 Aug 2024 19:23:04 +0000 Subject: [PATCH 12/35] Moved commit "Fix for federated topic" from ydb repo --- .../impl/federated_write_session.cpp | 19 ++++++++++++++++--- .../impl/federated_write_session.h | 2 ++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index ce41c0e318..560ac90a96 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -105,6 +105,9 @@ void TFederatedWriteSessionImpl::Start() { void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) { Y_ABORT_UNLESS(Lock.IsLocked()); + + ++SubsessionGeneration; + if (Subsession) { PendingToken.reset(); OldSubsession = std::move(Subsession); @@ -119,14 +122,20 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) auto handlers = NTopic::TWriteSessionSettings::TEventHandlers() .HandlersExecutor(Settings.EventHandlers_.HandlersExecutor_) - .ReadyToAcceptHandler([selfCtx = SelfContext](NTopic::TWriteSessionEvent::TReadyToAcceptEvent& ev) { + .ReadyToAcceptHandler([selfCtx = SelfContext, generation = SubsessionGeneration](NTopic::TWriteSessionEvent::TReadyToAcceptEvent& ev) { if (auto self = selfCtx->LockShared()) { - TDeferredWrite deferred(self->Subsession); + TDeferredWrite deferred; + with_lock(self->Lock) { + if (generation != self->SubsessionGeneration) { + return; + } + Y_ABORT_UNLESS(!self->PendingToken.has_value()); self->PendingToken = std::move(ev.ContinuationToken); self->PrepareDeferredWriteImpl(deferred); } + deferred.DoWrite(); } }) @@ -149,7 +158,7 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) self->IssueTokenIfAllowed(); } }) - .SessionClosedHandler([selfCtx = SelfContext](const NTopic::TSessionClosedEvent & ev) { + .SessionClosedHandler([selfCtx = SelfContext, generation = SubsessionGeneration](const NTopic::TSessionClosedEvent & ev) { if (ev.IsSuccess()) { // The subsession was closed by the federated write session itself while creating a new subsession. // In this case we get SUCCESS status and don't need to propagate it further. @@ -157,6 +166,9 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) } if (auto self = selfCtx->LockShared()) { with_lock(self->Lock) { + if (generation != self->SubsessionGeneration) { + return; + } self->CloseImpl(ev); } } @@ -391,6 +403,7 @@ bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferr } OriginalMessagesToGetAck.push_back(std::move(OriginalMessagesToPassDown.front())); OriginalMessagesToPassDown.pop_front(); + deferred.Writer = Subsession; deferred.Token.emplace(std::move(*PendingToken)); deferred.Message.emplace(std::move(OriginalMessagesToGetAck.back().Message)); PendingToken.reset(); diff --git a/src/client/federated_topic/impl/federated_write_session.h b/src/client/federated_topic/impl/federated_write_session.h index 8df0a4d735..f70c78e98d 100644 --- a/src/client/federated_topic/impl/federated_write_session.h +++ b/src/client/federated_topic/impl/federated_write_session.h @@ -62,6 +62,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, }; struct TDeferredWrite { + TDeferredWrite() {} explicit TDeferredWrite(std::shared_ptr writer) : Writer(std::move(writer)) { } @@ -121,6 +122,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, TAdaptiveLock Lock; + size_t SubsessionGeneration = 0; std::shared_ptr Subsession; std::shared_ptr OldSubsession; From 6ba38eb0fa7cb1ce7ddafcbc2290e7ef2be4bf9b Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Fri, 23 Aug 2024 19:51:02 +0000 Subject: [PATCH 13/35] Moved commit "autopartitioning for kinesis" from ydb repo --- .../client/datastreams/datastreams.h | 146 ++++++++++++++++++ src/client/datastreams/datastreams.cpp | 46 +++++- 2 files changed, 190 insertions(+), 2 deletions(-) diff --git a/include/ydb-cpp-sdk/client/datastreams/datastreams.h b/include/ydb-cpp-sdk/client/datastreams/datastreams.h index 15ead26b0c..7c4be777f3 100644 --- a/include/ydb-cpp-sdk/client/datastreams/datastreams.h +++ b/include/ydb-cpp-sdk/client/datastreams/datastreams.h @@ -101,13 +101,157 @@ namespace NYdb::NDataStreams::V1 { std::string ExplicitHashDecimal; }; + enum class EAutoPartitioningStrategy: uint32_t { + Unspecified = 0, + Disabled = 1, + ScaleUp = 2, + ScaleUpAndDown = 3, + }; + + struct TCreateStreamSettings; + struct TUpdateStreamSettings; + + + template + struct TPartitioningSettingsBuilder; + template + struct TAutoPartitioningSettingsBuilder; + + struct TAutoPartitioningSettings { + friend struct TAutoPartitioningSettingsBuilder; + friend struct TAutoPartitioningSettingsBuilder; + public: + TAutoPartitioningSettings() + : Strategy_(EAutoPartitioningStrategy::Disabled) + , StabilizationWindow_(TDuration::Seconds(0)) + , DownUtilizationPercent_(0) + , UpUtilizationPercent_(0) { + } + TAutoPartitioningSettings(const Ydb::DataStreams::V1::AutoPartitioningSettings& settings); + TAutoPartitioningSettings(EAutoPartitioningStrategy strategy, TDuration stabilizationWindow, uint64_t downUtilizationPercent, uint64_t upUtilizationPercent) + : Strategy_(strategy) + , StabilizationWindow_(stabilizationWindow) + , DownUtilizationPercent_(downUtilizationPercent) + , UpUtilizationPercent_(upUtilizationPercent) {} + + EAutoPartitioningStrategy GetStrategy() const { return Strategy_; }; + TDuration GetStabilizationWindow() const { return StabilizationWindow_; }; + uint32_t GetDownUtilizationPercent() const { return DownUtilizationPercent_; }; + uint32_t GetUpUtilizationPercent() const { return UpUtilizationPercent_; }; + private: + EAutoPartitioningStrategy Strategy_; + TDuration StabilizationWindow_; + uint32_t DownUtilizationPercent_; + uint32_t UpUtilizationPercent_; + }; + + + class TPartitioningSettings { + using TSelf = TPartitioningSettings; + friend struct TPartitioningSettingsBuilder; + friend struct TPartitioningSettingsBuilder; + public: + TPartitioningSettings() : MinActivePartitions_(0), MaxActivePartitions_(0), AutoPartitioningSettings_(){} + TPartitioningSettings(const Ydb::DataStreams::V1::PartitioningSettings& settings); + TPartitioningSettings(uint64_t minActivePartitions, uint64_t maxActivePartitions, TAutoPartitioningSettings autoscalingSettings = {}) + : MinActivePartitions_(minActivePartitions) + , MaxActivePartitions_(maxActivePartitions) + , AutoPartitioningSettings_(autoscalingSettings) { + } + + uint64_t GetMinActivePartitions() const { return MinActivePartitions_; }; + uint64_t GetMaxActivePartitions() const { return MaxActivePartitions_; }; + TAutoPartitioningSettings GetAutoPartitioningSettings() const { return AutoPartitioningSettings_; }; + private: + uint64_t MinActivePartitions_; + uint64_t MaxActivePartitions_; + TAutoPartitioningSettings AutoPartitioningSettings_; + }; + struct TCreateStreamSettings : public NYdb::TOperationRequestSettings { FLUENT_SETTING(uint32_t, ShardCount); FLUENT_SETTING_OPTIONAL(uint32_t, RetentionPeriodHours); FLUENT_SETTING_OPTIONAL(uint32_t, RetentionStorageMegabytes); FLUENT_SETTING(uint64_t, WriteQuotaKbPerSec); FLUENT_SETTING_OPTIONAL(EStreamMode, StreamMode); + + + FLUENT_SETTING_OPTIONAL(TPartitioningSettings, PartitioningSettings); + TPartitioningSettingsBuilder BeginConfigurePartitioningSettings(); }; + + template + struct TAutoPartitioningSettingsBuilder { + using TSelf = TAutoPartitioningSettingsBuilder; + public: + TAutoPartitioningSettingsBuilder(TPartitioningSettingsBuilder& parent, TAutoPartitioningSettings& settings): Parent_(parent), Settings_(settings) {} + + TSelf Strategy(EAutoPartitioningStrategy value) { + Settings_.Strategy_ = value; + return *this; + } + + TSelf StabilizationWindow(TDuration value) { + Settings_.StabilizationWindow_ = value; + return *this; + } + + TSelf DownUtilizationPercent(uint32_t value) { + Settings_.DownUtilizationPercent_ = value; + return *this; + } + + TSelf UpUtilizationPercent(uint32_t value) { + Settings_.UpUtilizationPercent_ = value; + return *this; + } + + TPartitioningSettingsBuilder& EndConfigureAutoPartitioningSettings() { + return Parent_; + } + + private: + TPartitioningSettingsBuilder& Parent_; + TAutoPartitioningSettings& Settings_; + }; + + template + struct TPartitioningSettingsBuilder { + using TSelf = TPartitioningSettingsBuilder; + public: + TPartitioningSettingsBuilder(TSettings& parent): Parent_(parent) {} + + TSelf MinActivePartitions(uint64_t value) { + if (!Parent_.PartitioningSettings_.has_value()) { + Parent_.PartitioningSettings_.emplace(); + } + (*Parent_.PartitioningSettings_).MinActivePartitions_ = value; + return *this; + } + + TSelf MaxActivePartitions(uint64_t value) { + if (!Parent_.PartitioningSettings_.has_value()) { + Parent_.PartitioningSettings_.emplace(); + } + (*Parent_.PartitioningSettings_).MaxActivePartitions_ = value; + return *this; + } + + TAutoPartitioningSettingsBuilder BeginConfigureAutoPartitioningSettings() { + if (!Parent_.PartitioningSettings_.has_value()) { + Parent_.PartitioningSettings_.emplace(); + } + return {*this, (*Parent_.PartitioningSettings_).AutoPartitioningSettings_}; + } + + TSettings& EndConfigurePartitioningSettings() { + return Parent_; + } + + private: + TSettings& Parent_; + }; + struct TListStreamsSettings : public NYdb::TOperationRequestSettings { FLUENT_SETTING(uint32_t, Limit); FLUENT_SETTING(std::string, ExclusiveStartStreamName); @@ -155,6 +299,8 @@ namespace NYdb::NDataStreams::V1 { FLUENT_SETTING(uint64_t, WriteQuotaKbPerSec); FLUENT_SETTING_OPTIONAL(EStreamMode, StreamMode); + FLUENT_SETTING_OPTIONAL(TPartitioningSettings, PartitioningSettings); + TPartitioningSettingsBuilder BeginConfigurePartitioningSettings(); }; struct TPutRecordSettings : public NYdb::TOperationRequestSettings {}; struct TPutRecordsSettings : public NYdb::TOperationRequestSettings {}; diff --git a/src/client/datastreams/datastreams.cpp b/src/client/datastreams/datastreams.cpp index e014026fb1..2701deb7d1 100644 --- a/src/client/datastreams/datastreams.cpp +++ b/src/client/datastreams/datastreams.cpp @@ -5,12 +5,47 @@ #undef INCLUDE_YDB_INTERNAL_H #include -//#include +#include #include namespace NYdb::NDataStreams::V1 { + TPartitioningSettingsBuilder TCreateStreamSettings::BeginConfigurePartitioningSettings() { + return { *this }; + } + + TPartitioningSettingsBuilder TUpdateStreamSettings::BeginConfigurePartitioningSettings() { + return { *this }; + } + + void SetPartitionSettings(const TPartitioningSettings& ps, ::Ydb::DataStreams::V1::PartitioningSettings* pt) { + pt->set_max_active_partitions(ps.GetMaxActivePartitions()); + pt->set_min_active_partitions(ps.GetMinActivePartitions()); + + ::Ydb::DataStreams::V1::AutoPartitioningStrategy strategy; + switch (ps.GetAutoPartitioningSettings().GetStrategy()) { + case EAutoPartitioningStrategy::Unspecified: + case EAutoPartitioningStrategy::Disabled: + strategy = ::Ydb::DataStreams::V1::AutoPartitioningStrategy::AUTO_PARTITIONING_STRATEGY_DISABLED; + break; + case EAutoPartitioningStrategy::ScaleUp: + strategy = ::Ydb::DataStreams::V1::AutoPartitioningStrategy::AUTO_PARTITIONING_STRATEGY_SCALE_UP; + break; + case EAutoPartitioningStrategy::ScaleUpAndDown: + strategy = ::Ydb::DataStreams::V1::AutoPartitioningStrategy::AUTO_PARTITIONING_STRATEGY_SCALE_UP_AND_DOWN; + break; + } + + pt->mutable_auto_partitioning_settings()->set_strategy(strategy); + pt->mutable_auto_partitioning_settings()->mutable_partition_write_speed() + ->mutable_stabilization_window()->set_seconds(ps.GetAutoPartitioningSettings().GetStabilizationWindow().Seconds()); + pt->mutable_auto_partitioning_settings()->mutable_partition_write_speed() + ->set_up_utilization_percent(ps.GetAutoPartitioningSettings().GetUpUtilizationPercent()); + pt->mutable_auto_partitioning_settings()->mutable_partition_write_speed() + ->set_down_utilization_percent(ps.GetAutoPartitioningSettings().GetDownUtilizationPercent()); + } + class TDataStreamsClient::TImpl : public TClientImplCommon { public: TImpl(std::shared_ptr &&connections, const TCommonClientSettings &settings) @@ -88,6 +123,10 @@ namespace NYdb::NDataStreams::V1 { *settings.StreamMode_ == ESM_PROVISIONED ? Ydb::DataStreams::V1::StreamMode::PROVISIONED : Ydb::DataStreams::V1::StreamMode::ON_DEMAND); } + + if (settings.PartitioningSettings_.has_value()) { + SetPartitionSettings(*settings.PartitioningSettings_, req.mutable_partitioning_settings()); + } }); } @@ -372,6 +411,10 @@ namespace NYdb::NDataStreams::V1 { *settings.StreamMode_ == ESM_PROVISIONED ? Ydb::DataStreams::V1::StreamMode::PROVISIONED : Ydb::DataStreams::V1::StreamMode::ON_DEMAND); } + + if (settings.PartitioningSettings_.has_value()) { + SetPartitionSettings(*settings.PartitioningSettings_, req.mutable_partitioning_settings()); + } }); } @@ -907,4 +950,3 @@ namespace NYdb::NDataStreams::V1 { TProtoRequestSettings settings ); } - From 2fe54ce41d0d4f1aef4b9d6c0a290b805a60983b Mon Sep 17 00:00:00 2001 From: Alek5andr-Kotov Date: Fri, 23 Aug 2024 20:06:36 +0000 Subject: [PATCH 14/35] Moved commit "Memory leak in Topic SDK" from ydb repo --- src/client/topic/impl/read_session_impl.h | 13 ++++- src/client/topic/impl/read_session_impl.ipp | 58 ++++++++++++++++----- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/src/client/topic/impl/read_session_impl.h b/src/client/topic/impl/read_session_impl.h index b7464cae4c..e4da1a0530 100644 --- a/src/client/topic/impl/read_session_impl.h +++ b/src/client/topic/impl/read_session_impl.h @@ -140,7 +140,7 @@ class TDeferredActions { } void DeferReadFromProcessor(const typename IProcessor::TPtr& processor, TServerMessage* dst, typename IProcessor::TReadCallback callback); - void DeferStartExecutorTask(const typename IAExecutor::TPtr& executor, typename IAExecutor::TFunction task); + void DeferStartExecutorTask(const typename IAExecutor::TPtr& executor, typename IAExecutor::TFunction&& task); void DeferAbortSession(TCallbackContextPtr cbContext, TASessionClosedEvent&& closeEvent); void DeferAbortSession(TCallbackContextPtr cbContext, EStatus statusCode, NYql::TIssues&& issues); void DeferAbortSession(TCallbackContextPtr cbContext, EStatus statusCode, const std::string& message); @@ -206,6 +206,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this> partitionStream); + void OnDestroyReadSession(); + bool IsReady() const { return SourceDataNotProcessed == 0; } @@ -303,6 +305,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this> PartitionStream; @@ -1096,6 +1100,11 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext& deferred); void OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount, i64 serverBytesSize = 0); @@ -1285,6 +1294,8 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext BatchInfo; TIntrusivePtr> PartitionStream; }; diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index c05b85b2be..5e7d976f68 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -218,6 +218,15 @@ void TRawPartitionStreamEventQueue::DeleteNotReadyTail(TDe swap(ready, NotReady); } +//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TDecompressionQueueItem + +template +void TSingleClusterReadSessionImpl::TDecompressionQueueItem::OnDestroyReadSession() +{ + BatchInfo->OnDestroyReadSession(); +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TSingleClusterReadSessionImpl @@ -226,6 +235,10 @@ TSingleClusterReadSessionImpl::~TSingleClusterReadSessionI for (auto&& [_, partitionStream] : PartitionStreams) { partitionStream->ClearQueue(); } + + for (auto& e : DecompressionQueue) { + e.OnDestroyReadSession(); + } } template @@ -1565,6 +1578,7 @@ void TSingleClusterReadSessionImpl::OnDecompressionInfoDes template void TSingleClusterReadSessionImpl::OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount, i64 serverBytesSize) { + TDeferredActions deferred; Y_ABORT_UNLESS(DecompressionTasksInflight > 0); @@ -2512,6 +2526,14 @@ void TDataDecompressionInfo::PlanDecompressionTasks(double } } +template +void TDataDecompressionInfo::OnDestroyReadSession() +{ + for (auto& task : Tasks) { + task.ClearParent(); + } +} + template void TDataDecompressionEvent::TakeData(TIntrusivePtr> partitionStream, std::vector::TMessage>& messages, @@ -2661,19 +2683,23 @@ TDataDecompressionInfo::TDecompressionTask::TDecompression template void TDataDecompressionInfo::TDecompressionTask::operator()() { + auto parent = Parent; + if (!parent) { + return; + } i64 minOffset = Max(); i64 maxOffset = 0; - const i64 partition_id = [this](){ + const i64 partition_id = [parent](){ if constexpr (UseMigrationProtocol) { - return Parent->ServerMessage.partition(); + return parent->ServerMessage.partition(); } else { - return Parent->ServerMessage.partition_session_id(); + return parent->ServerMessage.partition_session_id(); } }(); i64 dataProcessed = 0; size_t messagesProcessed = 0; for (const TMessageRange& messages : Messages) { - auto& batch = *Parent->ServerMessage.mutable_batches(messages.Batch); + auto& batch = *parent->ServerMessage.mutable_batches(messages.Batch); for (size_t i = messages.MessageRange.first; i < messages.MessageRange.second; ++i) { auto& data = *batch.mutable_message_data(i); @@ -2684,7 +2710,7 @@ void TDataDecompressionInfo::TDecompressionTask::operator( try { if constexpr (UseMigrationProtocol) { - if (Parent->DoDecompress + if (parent->DoDecompress && data.codec() != Ydb::PersQueue::V1::CODEC_RAW && data.codec() != Ydb::PersQueue::V1::CODEC_UNSPECIFIED ) { @@ -2694,7 +2720,7 @@ void TDataDecompressionInfo::TDecompressionTask::operator( data.set_codec(Ydb::PersQueue::V1::CODEC_RAW); } } else { - if (Parent->DoDecompress + if (parent->DoDecompress && static_cast(batch.codec()) != Ydb::Topic::CODEC_RAW && static_cast(batch.codec()) != Ydb::Topic::CODEC_UNSPECIFIED ) { @@ -2706,32 +2732,38 @@ void TDataDecompressionInfo::TDecompressionTask::operator( DecompressedSize += data.data().size(); } catch (...) { - Parent->PutDecompressionError(std::current_exception(), messages.Batch, i); + parent->PutDecompressionError(std::current_exception(), messages.Batch, i); data.clear_data(); // Free memory, because we don't count it. - if (auto session = Parent->CbContext->LockShared()) { + if (auto session = parent->CbContext->LockShared()) { session->GetLog() << TLOG_INFO << "Error decompressing data: " << CurrentExceptionMessage(); } } } } - if (auto session = Parent->CbContext->LockShared()) { + if (auto session = parent->CbContext->LockShared()) { LOG_LAZY(session->GetLog(), TLOG_DEBUG, TStringBuilder() << "Decompression task done. Partition/PartitionSessionId: " << partition_id << " (" << minOffset << "-" << maxOffset << ")"); } Y_ASSERT(dataProcessed == SourceDataSize); - Parent->OnDataDecompressed(SourceDataSize, EstimatedDecompressedSize, DecompressedSize, messagesProcessed); + parent->OnDataDecompressed(SourceDataSize, EstimatedDecompressedSize, DecompressedSize, messagesProcessed); - Parent->SourceDataNotProcessed -= dataProcessed; + parent->SourceDataNotProcessed -= dataProcessed; Ready->Ready = true; - if (auto session = Parent->CbContext->LockShared()) { + if (auto session = parent->CbContext->LockShared()) { session->GetEventsQueue()->SignalReadyEvents(PartitionStream); } } +template +void TDataDecompressionInfo::TDecompressionTask::ClearParent() +{ + Parent = nullptr; +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TUserRetrievedEventsInfoAccumulator @@ -2769,7 +2801,7 @@ void TDeferredActions::DeferReadFromProcessor(const typena } template -void TDeferredActions::DeferStartExecutorTask(const typename IAExecutor::TPtr& executor, typename IAExecutor::TFunction task) { +void TDeferredActions::DeferStartExecutorTask(const typename IAExecutor::TPtr& executor, typename IAExecutor::TFunction&& task) { ExecutorsTasks.emplace_back(executor, std::move(task)); } From 0a27948e5809e423460d35685880d9ccae83515f Mon Sep 17 00:00:00 2001 From: Ilia Shakhov Date: Fri, 23 Aug 2024 20:11:15 +0000 Subject: [PATCH 15/35] Moved commit "Add created_by to Operations API" from ydb repo --- .../client/types/operation/operation.h | 3 ++- src/client/types/operation/operation.cpp | 21 +++++++++++++------ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/include/ydb-cpp-sdk/client/types/operation/operation.h b/include/ydb-cpp-sdk/client/types/operation/operation.h index f8f01891ad..88d226d2aa 100644 --- a/include/ydb-cpp-sdk/client/types/operation/operation.h +++ b/include/ydb-cpp-sdk/client/types/operation/operation.h @@ -32,8 +32,9 @@ class TOperation { const TOperationId& Id() const; bool Ready() const; const TStatus& Status() const; - TInstant StartTime() const; + TInstant CreateTime() const; TInstant EndTime() const; + const std::string& CreatedBy() const; std::string ToString() const; std::string ToJsonString() const; diff --git a/src/client/types/operation/operation.cpp b/src/client/types/operation/operation.cpp index 78d3224169..c325f465e7 100644 --- a/src/client/types/operation/operation.cpp +++ b/src/client/types/operation/operation.cpp @@ -21,7 +21,7 @@ class TOperation::TImpl { : Id_(operation.id(), true /* allowEmpty */) , Status_(std::move(status)) , Ready_(operation.ready()) - , StartTime_(ProtoTimestampToInstant(operation.start_time())) + , CreateTime_(ProtoTimestampToInstant(operation.create_time())) , EndTime_(ProtoTimestampToInstant(operation.end_time())) , Operation_(std::move(operation)) { @@ -39,14 +39,18 @@ class TOperation::TImpl { return Status_; } - TInstant StartTime() const { - return StartTime_; + TInstant CreateTime() const { + return CreateTime_; } TInstant EndTime() const { return EndTime_; } + const std::string& CreatedBy() const { + return CreatedBy_; + } + const Ydb::Operations::Operation& GetProto() const { return Operation_; } @@ -55,8 +59,9 @@ class TOperation::TImpl { const TOperationId Id_; const TStatus Status_; const bool Ready_; - const TInstant StartTime_; + const TInstant CreateTime_; const TInstant EndTime_; + const std::string CreatedBy_; const Ydb::Operations::Operation Operation_; }; @@ -80,14 +85,18 @@ const TStatus& TOperation::Status() const { return Impl_->Status(); } -TInstant TOperation::StartTime() const { - return Impl_->StartTime(); +TInstant TOperation::CreateTime() const { + return Impl_->CreateTime(); } TInstant TOperation::EndTime() const { return Impl_->EndTime(); } +const std::string& TOperation::CreatedBy() const { + return Impl_->CreatedBy(); +} + std::string TOperation::ToString() const { TString result; TStringOutput out(result); From 5c359a385143bc0566c1ba7991d1d494877613d8 Mon Sep 17 00:00:00 2001 From: azevaykin <145343289+azevaykin@users.noreply.github.com> Date: Mon, 26 Aug 2024 15:56:24 +0000 Subject: [PATCH 16/35] Moved commit "Create vector index in SchemeShard" from ydb repo --- include/ydb-cpp-sdk/client/table/table.h | 56 ++++- include/ydb-cpp-sdk/client/table/table_enum.h | 1 + src/client/table/out.cpp | 66 ++++++ src/client/table/table.cpp | 214 ++++++++++++++++-- 4 files changed, 319 insertions(+), 18 deletions(-) diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index d7619197d7..382974aad8 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -23,6 +23,7 @@ class ChangefeedDescription; class DescribeTableResult; class ExplicitPartitions; class GlobalIndexSettings; +class VectorIndexSettings; class PartitioningSettings; class DateTypeColumnModeSettings; class TtlSettings; @@ -190,6 +191,45 @@ struct TGlobalIndexSettings { void SerializeTo(Ydb::Table::GlobalIndexSettings& proto) const; }; +struct TVectorIndexSettings { +public: + enum class EDistance { + Cosine, + Manhattan, + Euclidean, + + Unknown = std::numeric_limits::max() + }; + + enum class ESimilarity { + Cosine, + InnerProduct, + + Unknown = std::numeric_limits::max() + }; + + enum class EVectorType { + Float, + Uint8, + Int8, + Bit, + + Unknown = std::numeric_limits::max() + }; + using TMetric = std::variant; + + TMetric Metric; + EVectorType VectorType; + uint32_t VectorDimension; + + template + static TVectorIndexSettings FromProto(const TProto& proto); + + void SerializeTo(Ydb::Table::VectorIndexSettings& settings) const; + + void Out(IOutputStream &o) const; +}; + //! Represents index description class TIndexDescription { friend class NYdb::TProtoAccessor; @@ -200,20 +240,22 @@ class TIndexDescription { EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns = {}, - const TGlobalIndexSettings& settings = {} + const std::vector& globalIndexSettings = {}, + const std::optional& vectorIndexSettings = {} ); TIndexDescription( const std::string& name, const std::vector& indexColumns, const std::vector& dataColumns = {}, - const TGlobalIndexSettings& settings = {} + const std::vector& globalIndexSettings = {} ); const std::string& GetIndexName() const; EIndexType GetIndexType() const; const std::vector& GetIndexColumns() const; const std::vector& GetDataColumns() const; + const std::optional& GetVectorIndexSettings() const; uint64_t GetSizeBytes() const; void SerializeTo(Ydb::Table::TableIndex& proto) const; @@ -232,7 +274,8 @@ class TIndexDescription { EIndexType IndexType_; std::vector IndexColumns_; std::vector DataColumns_; - TGlobalIndexSettings GlobalIndexSettings_; + std::vector GlobalIndexSettings_; + std::optional VectorIndexSettings_; uint64_t SizeBytes = 0; }; @@ -606,6 +649,9 @@ class TTableDescription { // unique void AddUniqueSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); void AddUniqueSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); + // vector KMeansTree + void AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings); + void AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TVectorIndexSettings& vectorIndexSettings); // default void AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); @@ -825,6 +871,10 @@ class TTableBuilder { TTableBuilder& AddUniqueSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); TTableBuilder& AddUniqueSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); + // vector KMeansTree + TTableBuilder& AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings); + TTableBuilder& AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TVectorIndexSettings& vectorIndexSettings); + // 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 25b57b005b..1660706f57 100644 --- a/include/ydb-cpp-sdk/client/table/table_enum.h +++ b/include/ydb-cpp-sdk/client/table/table_enum.h @@ -28,6 +28,7 @@ enum class EIndexType { GlobalSync, GlobalAsync, GlobalUnique, + GlobalVectorKMeansTree, Unknown = std::numeric_limits::max() }; diff --git a/src/client/table/out.cpp b/src/client/table/out.cpp index 2f24688a79..401695350c 100644 --- a/src/client/table/out.cpp +++ b/src/client/table/out.cpp @@ -23,3 +23,69 @@ Y_DECLARE_OUT_SPEC(, NYdb::NTable::TCreateSessionResult, o, x) { Y_DECLARE_OUT_SPEC(, NYdb::NTable::TDescribeTableResult, o, x) { return x.Out(o); } + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::EDistance, stream, value) { + auto convertDistance = [] (auto value) -> auto { + switch (value) { + case NYdb::NTable::TVectorIndexSettings::EDistance::Cosine: + return "COSINE"; + case NYdb::NTable::TVectorIndexSettings::EDistance::Manhattan: + return "MANHATTAN"; + case NYdb::NTable::TVectorIndexSettings::EDistance::Euclidean: + return "EUCLIDEAN"; + case NYdb::NTable::TVectorIndexSettings::EDistance::Unknown: + return "UNKNOWN"; + } + }; + + stream << convertDistance(value); +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::ESimilarity, stream, value) { + auto convertSimilarity = [] (auto value) -> auto { + switch (value) { + case NYdb::NTable::TVectorIndexSettings::ESimilarity::Cosine: + return "COSINE"; + case NYdb::NTable::TVectorIndexSettings::ESimilarity::InnerProduct: + return "INNER_PRODUCT"; + case NYdb::NTable::TVectorIndexSettings::ESimilarity::Unknown: + return "UNKNOWN"; + } + }; + + stream << convertSimilarity(value); +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::EVectorType, stream, value) { + auto convertVectorType = [] (auto value) -> auto { + switch (value) { + case NYdb::NTable::TVectorIndexSettings::EVectorType::Float: + return "FLOAT"; + case NYdb::NTable::TVectorIndexSettings::EVectorType::Uint8: + return "UINT8"; + case NYdb::NTable::TVectorIndexSettings::EVectorType::Int8: + return "INT8"; + case NYdb::NTable::TVectorIndexSettings::EVectorType::Bit: + return "BIT"; + case NYdb::NTable::TVectorIndexSettings::EVectorType::Unknown: + return "UNKNOWN"; + } + }; + + stream << convertVectorType(value); +} + +Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings, stream, value) { + stream << "{"; + + if (const auto* distance = std::get_if(&value.Metric)) { + stream << " distance: " << *distance << ""; + } else if (const auto* similarity = std::get_if(&value.Metric)) { + stream << " similarity: " << *similarity << ""; + } + + stream << ", vector_type: " << value.VectorType << ""; + stream << ", vector_dimension: " << value.VectorDimension << ""; + + stream << " }"; +} diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index bfe83e664f..01119cba90 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -27,6 +27,7 @@ #include #include +#include #include @@ -468,6 +469,14 @@ class TTableDescription::TImpl { Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, dataColumns)); } + void AddVectorIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings) { + Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, {}, {}, vectorIndexSettings)); + } + + void AddVectorIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns, const TVectorIndexSettings& vectorIndexSettings) { + Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, dataColumns, {}, vectorIndexSettings)); + } + void AddChangefeed(const std::string& name, EChangefeedMode mode, EChangefeedFormat format) { Changefeeds_.emplace_back(name, mode, format); } @@ -765,6 +774,14 @@ void TTableDescription::AddUniqueSecondaryIndex(const std::string& indexName, co AddSecondaryIndex(indexName, EIndexType::GlobalUnique, indexColumns, dataColumns); } +void TTableDescription::AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings) { + Impl_->AddVectorIndex(indexName, EIndexType::GlobalVectorKMeansTree, indexColumns, vectorIndexSettings); +} + +void TTableDescription::AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TVectorIndexSettings& vectorIndexSettings) { + Impl_->AddVectorIndex(indexName, EIndexType::GlobalVectorKMeansTree, indexColumns, dataColumns, vectorIndexSettings); +} + void TTableDescription::AddSecondaryIndex(const std::string& indexName, const std::vector& indexColumns) { AddSyncSecondaryIndex(indexName, indexColumns); } @@ -1212,6 +1229,16 @@ TTableBuilder& TTableBuilder::AddUniqueSecondaryIndex(const std::string& indexNa return AddSecondaryIndex(indexName, EIndexType::GlobalUnique, indexColumns); } +TTableBuilder& TTableBuilder::AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns, const TVectorIndexSettings& vectorIndexSettings) { + TableDescription_.AddVectorKMeansTreeSecondaryIndex(indexName, indexColumns, dataColumns, vectorIndexSettings); + return *this; +} + +TTableBuilder& TTableBuilder::AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings) { + TableDescription_.AddVectorKMeansTreeSecondaryIndex(indexName, indexColumns, vectorIndexSettings); + return *this; +} + TTableBuilder& TTableBuilder::AddSecondaryIndex(const std::string& indexName, const std::string& indexColumn) { return AddSyncSecondaryIndex(indexName, indexColumn); } @@ -2216,20 +2243,22 @@ TIndexDescription::TIndexDescription( EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns, - const TGlobalIndexSettings& settings + const std::vector& globalIndexSettings, + const std::optional& vectorIndexSettings ) : IndexName_(name) , IndexType_(type) , IndexColumns_(indexColumns) , DataColumns_(dataColumns) - , GlobalIndexSettings_(settings) + , GlobalIndexSettings_(globalIndexSettings) + , VectorIndexSettings_(vectorIndexSettings) {} TIndexDescription::TIndexDescription( const std::string& name, const std::vector& indexColumns, const std::vector& dataColumns, - const TGlobalIndexSettings& settings -) : TIndexDescription(name, EIndexType::GlobalSync, indexColumns, dataColumns, settings) + const std::vector& globalIndexSettings +) : TIndexDescription(name, EIndexType::GlobalSync, indexColumns, dataColumns, globalIndexSettings) {} TIndexDescription::TIndexDescription(const Ydb::Table::TableIndex& tableIndex) @@ -2256,6 +2285,10 @@ const std::vector& TIndexDescription::GetDataColumns() const { return DataColumns_; } +const std::optional& TIndexDescription::GetVectorIndexSettings() const { + return VectorIndexSettings_; +} + uint64_t TIndexDescription::GetSizeBytes() const { return SizeBytes; } @@ -2293,12 +2326,128 @@ void TGlobalIndexSettings::SerializeTo(Ydb::Table::GlobalIndexSettings& settings std::visit(std::move(variantVisitor), Partitions); } +template +TVectorIndexSettings TVectorIndexSettings::FromProto(const TProto& proto) { + auto convertDistance = [] (auto distance) -> auto { + switch (distance) { + case Ydb::Table::VectorIndexSettings::DISTANCE_COSINE: + return EDistance::Cosine; + case Ydb::Table::VectorIndexSettings::DISTANCE_MANHATTAN: + return EDistance::Manhattan; + case Ydb::Table::VectorIndexSettings::DISTANCE_EUCLIDEAN: + return EDistance::Euclidean; + default: + return EDistance::Unknown; + } + }; + + auto convertSimilarity = [] (auto similarity) -> auto { + switch (similarity) { + case Ydb::Table::VectorIndexSettings::SIMILARITY_COSINE: + return ESimilarity::Cosine; + case Ydb::Table::VectorIndexSettings::SIMILARITY_INNER_PRODUCT: + return ESimilarity::InnerProduct; + default: + return ESimilarity::Unknown; + } + }; + + auto convertVectorType = [] (auto vectorType) -> auto { + switch (vectorType) { + case Ydb::Table::VectorIndexSettings::VECTOR_TYPE_FLOAT: + return EVectorType::Float; + case Ydb::Table::VectorIndexSettings::VECTOR_TYPE_UINT8: + return EVectorType::Uint8; + case Ydb::Table::VectorIndexSettings::VECTOR_TYPE_INT8: + return EVectorType::Int8; + case Ydb::Table::VectorIndexSettings::VECTOR_TYPE_BIT: + return EVectorType::Bit; + default: + return EVectorType::Unknown; + } + }; + + + auto metricFromProto = [&](const auto& proto) -> TVectorIndexSettings::TMetric { + switch (proto.metric_case()) { + case TProto::kDistance: + return convertDistance(proto.distance()); + case TProto::kSimilarity: + return convertSimilarity(proto.similarity()); + default: + return {}; + } + }; + + return { + .Metric = metricFromProto(proto), + .VectorType = convertVectorType(proto.vector_type()), + .VectorDimension = proto.vector_dimension() + }; +} + +void TVectorIndexSettings::SerializeTo(Ydb::Table::VectorIndexSettings& settings) const { + auto convertDistance = [] (auto distance) -> auto { + switch (distance) { + case EDistance::Cosine: + return Ydb::Table::VectorIndexSettings::DISTANCE_COSINE; + case EDistance::Manhattan: + return Ydb::Table::VectorIndexSettings::DISTANCE_MANHATTAN; + case EDistance::Euclidean: + return Ydb::Table::VectorIndexSettings::DISTANCE_EUCLIDEAN; + case EDistance::Unknown: + return Ydb::Table::VectorIndexSettings::DISTANCE_UNSPECIFIED; + } + }; + + auto convertSimilarity = [] (auto similarity) -> auto { + switch (similarity) { + case ESimilarity::Cosine: + return Ydb::Table::VectorIndexSettings::SIMILARITY_COSINE; + case ESimilarity::InnerProduct: + return Ydb::Table::VectorIndexSettings::SIMILARITY_INNER_PRODUCT; + case ESimilarity::Unknown: + return Ydb::Table::VectorIndexSettings::SIMILARITY_UNSPECIFIED; + } + }; + + auto convertVectorType = [] (auto vectorType) -> auto { + switch (vectorType) { + case EVectorType::Float: + return Ydb::Table::VectorIndexSettings::VECTOR_TYPE_FLOAT; + case EVectorType::Uint8: + return Ydb::Table::VectorIndexSettings::VECTOR_TYPE_UINT8; + case EVectorType::Int8: + return Ydb::Table::VectorIndexSettings::VECTOR_TYPE_INT8; + case EVectorType::Bit: + return Ydb::Table::VectorIndexSettings::VECTOR_TYPE_BIT; + case EVectorType::Unknown: + return Ydb::Table::VectorIndexSettings::VECTOR_TYPE_UNSPECIFIED; + } + }; + + + if (const auto* distance = std::get_if(&Metric)) { + settings.set_distance(convertDistance(*distance)); + } else if (const auto* similarity = std::get_if(&Metric)) { + settings.set_similarity(convertSimilarity(*similarity)); + } + + settings.set_vector_type(convertVectorType(VectorType)); + settings.set_vector_dimension(VectorDimension); +} + +void TVectorIndexSettings::Out(IOutputStream& o) const { + o << *this; +} + template TIndexDescription TIndexDescription::FromProto(const TProto& proto) { EIndexType type; std::vector indexColumns; std::vector dataColumns; - TGlobalIndexSettings globalIndexSettings; + std::vector globalIndexSettings; + std::optional vectorIndexSettings; indexColumns.assign(proto.index_columns().begin(), proto.index_columns().end()); dataColumns.assign(proto.data_columns().begin(), proto.data_columns().end()); @@ -2306,22 +2455,31 @@ TIndexDescription TIndexDescription::FromProto(const TProto& proto) { switch (proto.type_case()) { case TProto::kGlobalIndex: type = EIndexType::GlobalSync; - globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_index().settings()); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(proto.global_index().settings())); break; case TProto::kGlobalAsyncIndex: type = EIndexType::GlobalAsync; - globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_async_index().settings()); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(proto.global_async_index().settings())); break; case TProto::kGlobalUniqueIndex: type = EIndexType::GlobalUnique; - globalIndexSettings = TGlobalIndexSettings::FromProto(proto.global_unique_index().settings()); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(proto.global_unique_index().settings())); + break; + case TProto::kGlobalVectorKmeansTreeIndex: { + type = EIndexType::GlobalVectorKMeansTree; + const auto &vectorProto = proto.global_vector_kmeans_tree_index(); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(vectorProto.level_table_settings())); + globalIndexSettings.emplace_back(TGlobalIndexSettings::FromProto(vectorProto.posting_table_settings())); + vectorIndexSettings = TVectorIndexSettings::FromProto(vectorProto.vector_settings()); break; + } default: // fallback to global sync type = EIndexType::GlobalSync; + globalIndexSettings.resize(1); break; } - auto result = TIndexDescription(proto.name(), type, indexColumns, dataColumns, globalIndexSettings); + auto result = TIndexDescription(proto.name(), type, indexColumns, dataColumns, globalIndexSettings, vectorIndexSettings); if constexpr (std::is_same_v) { result.SizeBytes = proto.size_bytes(); } @@ -2338,15 +2496,38 @@ void TIndexDescription::SerializeTo(Ydb::Table::TableIndex& proto) const { *proto.mutable_data_columns() = {DataColumns_.begin(), DataColumns_.end()}; switch (IndexType_) { - case EIndexType::GlobalSync: - GlobalIndexSettings_.SerializeTo(*proto.mutable_global_index()->mutable_settings()); + case EIndexType::GlobalSync: { + auto& settings = *proto.mutable_global_index()->mutable_settings(); + if (GlobalIndexSettings_.size() == 1) + GlobalIndexSettings_[0].SerializeTo(settings); break; - case EIndexType::GlobalAsync: - GlobalIndexSettings_.SerializeTo(*proto.mutable_global_async_index()->mutable_settings()); + } + case EIndexType::GlobalAsync: { + auto& settings = *proto.mutable_global_async_index()->mutable_settings(); + if (GlobalIndexSettings_.size() == 1) + GlobalIndexSettings_[0].SerializeTo(settings); break; - case EIndexType::GlobalUnique: - GlobalIndexSettings_.SerializeTo(*proto.mutable_global_unique_index()->mutable_settings()); + } + case EIndexType::GlobalUnique: { + auto& settings = *proto.mutable_global_unique_index()->mutable_settings(); + if (GlobalIndexSettings_.size() == 1) + GlobalIndexSettings_[0].SerializeTo(settings); break; + } + case EIndexType::GlobalVectorKMeansTree: { + auto* global_vector_kmeans_tree_index = proto.mutable_global_vector_kmeans_tree_index(); + auto& level_settings = *global_vector_kmeans_tree_index->mutable_level_table_settings(); + auto& posting_settings = *global_vector_kmeans_tree_index->mutable_posting_table_settings(); + auto& vector_settings = *global_vector_kmeans_tree_index->mutable_vector_settings(); + if (GlobalIndexSettings_.size() == 2) { + GlobalIndexSettings_[0].SerializeTo(level_settings); + GlobalIndexSettings_[1].SerializeTo(posting_settings); + } + if (VectorIndexSettings_) { + VectorIndexSettings_->SerializeTo(vector_settings); + } + break; + } case EIndexType::Unknown: break; } @@ -2368,6 +2549,9 @@ void TIndexDescription::Out(IOutputStream& o) const { o << ", data_columns: [" << JoinSeq(", ", DataColumns_) << "]"; } + if (VectorIndexSettings_) { + o << ", vector_settings: " << *VectorIndexSettings_ << ""; + } o << " }"; } From a21ef5d7270299a215a1bba3ce47dd3c1de43116 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 16:30:04 +0000 Subject: [PATCH 17/35] Moved commit "ydb_topic: Fallback to "indirect" write if the server doesn't support direct write" from ydb repo --- src/client/topic/impl/write_session_impl.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 3fe4b79d29..e967d40b08 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -301,7 +301,15 @@ void TWriteSessionImpl::OnDescribePartition(const TStatus& status, const Ydb::To if (!status.IsSuccess()) { with_lock(Lock) { - handleResult = OnErrorImpl({status.GetStatus(), MakeIssueWithSubIssues("Failed to get partition location", status.GetIssues())}); + if (status.GetStatus() == EStatus::CLIENT_CALL_UNIMPLEMENTED) { + Settings.DirectWriteToPartition_ = false; + handleResult = OnErrorImpl({ + EStatus::UNAVAILABLE, + MakeIssueWithSubIssues("The server does not support direct write, fallback to in-direct write", status.GetIssues()) + }); + } else { + handleResult = OnErrorImpl({status.GetStatus(), MakeIssueWithSubIssues("Failed to get partition location", status.GetIssues())}); + } } ProcessHandleResult(handleResult); return; From 01a4fbc797436346deeab7720230f69a5443b7d7 Mon Sep 17 00:00:00 2001 From: Ilia Shakhov Date: Mon, 26 Aug 2024 16:33:19 +0000 Subject: [PATCH 18/35] Moved commit "Add permissions to ydb tools dump" from ydb repo --- include/ydb-cpp-sdk/client/scheme/scheme.h | 7 +++++++ src/client/scheme/scheme.cpp | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/include/ydb-cpp-sdk/client/scheme/scheme.h b/include/ydb-cpp-sdk/client/scheme/scheme.h index c418bf1ef3..4692edb9de 100644 --- a/include/ydb-cpp-sdk/client/scheme/scheme.h +++ b/include/ydb-cpp-sdk/client/scheme/scheme.h @@ -6,6 +6,8 @@ namespace Ydb { class VirtualTimestamp; namespace Scheme { class Entry; + class ModifyPermissionsRequest; + class Permissions; } } @@ -24,6 +26,8 @@ struct TPermissions { {} std::string Subject; std::vector PermissionNames; + + void SerializeTo(::Ydb::Scheme::Permissions& proto) const; }; enum class ESchemeEntryType : i32 { @@ -77,6 +81,9 @@ struct TSchemeEntry { TSchemeEntry(const ::Ydb::Scheme::Entry& proto); void Out(IOutputStream& out) const; + + // Fills ModifyPermissionsRequest proto from this entry + void SerializeTo(::Ydb::Scheme::ModifyPermissionsRequest& request) const; }; //////////////////////////////////////////////////////////////////////////////// diff --git a/src/client/scheme/scheme.cpp b/src/client/scheme/scheme.cpp index bffe90cc0a..53ab858213 100644 --- a/src/client/scheme/scheme.cpp +++ b/src/client/scheme/scheme.cpp @@ -17,6 +17,13 @@ namespace NScheme { using namespace NThreading; using namespace Ydb::Scheme; +void TPermissions::SerializeTo(::Ydb::Scheme::Permissions& proto) const { + proto.set_subject(Subject); + for (const auto& name : PermissionNames) { + *proto.mutable_permission_names()->Add() = name; + } +} + TVirtualTimestamp::TVirtualTimestamp(uint64_t planStep, uint64_t txId) : PlanStep(planStep) , TxId(txId) @@ -120,6 +127,13 @@ void TSchemeEntry::Out(IOutputStream& out) const { << " }"; } +void TSchemeEntry::SerializeTo(::Ydb::Scheme::ModifyPermissionsRequest& request) const { + request.mutable_actions()->Add()->set_change_owner(Owner); + for (const auto& permission : Permissions) { + permission.SerializeTo(*request.mutable_actions()->Add()->mutable_set()); + } +} + class TSchemeClient::TImpl : public TClientImplCommon { public: TImpl(std::shared_ptr&& connections, const TCommonClientSettings& settings) From f0124566acc05933129cf204b52c9dd9b8ea07ac Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy <79596613+GrigoriyPA@users.noreply.github.com> Date: Mon, 26 Aug 2024 16:35:56 +0000 Subject: [PATCH 19/35] Moved commit "YQ-3447 support ydb scheme ls for resource pools" from ydb repo --- include/ydb-cpp-sdk/client/scheme/scheme.h | 3 ++- src/client/scheme/scheme.cpp | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/include/ydb-cpp-sdk/client/scheme/scheme.h b/include/ydb-cpp-sdk/client/scheme/scheme.h index 4692edb9de..060a5fafdb 100644 --- a/include/ydb-cpp-sdk/client/scheme/scheme.h +++ b/include/ydb-cpp-sdk/client/scheme/scheme.h @@ -46,7 +46,8 @@ enum class ESchemeEntryType : i32 { Topic = 17, ExternalTable = 18, ExternalDataSource = 19, - View = 20 + View = 20, + ResourcePool = 21, }; struct TVirtualTimestamp { diff --git a/src/client/scheme/scheme.cpp b/src/client/scheme/scheme.cpp index 53ab858213..b37beebce6 100644 --- a/src/client/scheme/scheme.cpp +++ b/src/client/scheme/scheme.cpp @@ -102,6 +102,8 @@ static ESchemeEntryType ConvertProtoEntryType(::Ydb::Scheme::Entry::Type entry) return ESchemeEntryType::ExternalDataSource; case ::Ydb::Scheme::Entry::VIEW: return ESchemeEntryType::View; + case ::Ydb::Scheme::Entry::RESOURCE_POOL: + return ESchemeEntryType::ResourcePool; default: return ESchemeEntryType::Unknown; } From 7a70b5c642b863232921ea939b9801b90b61502c Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Mon, 26 Aug 2024 17:08:25 +0000 Subject: [PATCH 20/35] Moved commit "Add PAUSED strategy for autopartitioning of the topic" from ydb repo --- include/ydb-cpp-sdk/client/datastreams/datastreams.h | 1 + include/ydb-cpp-sdk/client/topic/control_plane.h | 1 + src/client/datastreams/datastreams.cpp | 3 +++ 3 files changed, 5 insertions(+) diff --git a/include/ydb-cpp-sdk/client/datastreams/datastreams.h b/include/ydb-cpp-sdk/client/datastreams/datastreams.h index 7c4be777f3..0e031374e5 100644 --- a/include/ydb-cpp-sdk/client/datastreams/datastreams.h +++ b/include/ydb-cpp-sdk/client/datastreams/datastreams.h @@ -106,6 +106,7 @@ namespace NYdb::NDataStreams::V1 { Disabled = 1, ScaleUp = 2, ScaleUpAndDown = 3, + Paused = 4, }; struct TCreateStreamSettings; diff --git a/include/ydb-cpp-sdk/client/topic/control_plane.h b/include/ydb-cpp-sdk/client/topic/control_plane.h index 4fbc92ad4a..429b1a6a20 100644 --- a/include/ydb-cpp-sdk/client/topic/control_plane.h +++ b/include/ydb-cpp-sdk/client/topic/control_plane.h @@ -33,6 +33,7 @@ enum class EAutoPartitioningStrategy: uint32_t { Disabled = 1, ScaleUp = 2, ScaleUpAndDown = 3, + Paused = 4, }; class TConsumer { diff --git a/src/client/datastreams/datastreams.cpp b/src/client/datastreams/datastreams.cpp index 2701deb7d1..a17009b9b9 100644 --- a/src/client/datastreams/datastreams.cpp +++ b/src/client/datastreams/datastreams.cpp @@ -35,6 +35,9 @@ namespace NYdb::NDataStreams::V1 { case EAutoPartitioningStrategy::ScaleUpAndDown: strategy = ::Ydb::DataStreams::V1::AutoPartitioningStrategy::AUTO_PARTITIONING_STRATEGY_SCALE_UP_AND_DOWN; break; + case EAutoPartitioningStrategy::Paused: + strategy = ::Ydb::DataStreams::V1::AutoPartitioningStrategy::AUTO_PARTITIONING_STRATEGY_PAUSED; + break; } pt->mutable_auto_partitioning_settings()->set_strategy(strategy); From 4b100516ab839f34bce8b466fee6b0cc5c5e3237 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 17:21:10 +0000 Subject: [PATCH 21/35] Moved commit "Use 1-thread pool executor for subsession event handlers" from ydb repo --- .../impl/federated_topic_impl.cpp | 12 +++- .../impl/federated_topic_impl.h | 8 +++ .../impl/federated_write_session.cpp | 55 ++++++++++--------- .../impl/federated_write_session.h | 38 ++++--------- 4 files changed, 58 insertions(+), 55 deletions(-) diff --git a/src/client/federated_topic/impl/federated_topic_impl.cpp b/src/client/federated_topic/impl/federated_topic_impl.cpp index 3589ca236c..4cf28a6212 100644 --- a/src/client/federated_topic/impl/federated_topic_impl.cpp +++ b/src/client/federated_topic/impl/federated_topic_impl.cpp @@ -35,7 +35,8 @@ TFederatedTopicClient::TImpl::CreateWriteSession(const TFederatedWriteSessionSet splitSettings.EventHandlers_.HandlersExecutor(ClientSettings.DefaultHandlersExecutor_); } } - auto session = std::make_shared(splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs); + auto session = std::make_shared( + splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs, GetSubsessionHandlersExecutor()); session->Start(); return std::move(session); } @@ -48,4 +49,13 @@ void TFederatedTopicClient::TImpl::InitObserver() { } } +auto TFederatedTopicClient::TImpl::GetSubsessionHandlersExecutor() -> NTopic::IExecutor::TPtr { + with_lock (Lock) { + if (!SubsessionHandlersExecutor) { + SubsessionHandlersExecutor = NTopic::CreateThreadPoolExecutor(1); + } + return SubsessionHandlersExecutor; + } +} + } diff --git a/src/client/federated_topic/impl/federated_topic_impl.h b/src/client/federated_topic/impl/federated_topic_impl.h index 14991b4bdf..3c52904fbb 100644 --- a/src/client/federated_topic/impl/federated_topic_impl.h +++ b/src/client/federated_topic/impl/federated_topic_impl.h @@ -72,12 +72,20 @@ class TFederatedTopicClient::TImpl { void InitObserver(); +private: + + // Use single-threaded executor to prevent deadlocks inside subsession event handlers. + NTopic::IExecutor::TPtr GetSubsessionHandlersExecutor(); + private: std::shared_ptr Connections; const TFederatedTopicClientSettings ClientSettings; std::shared_ptr Observer; std::shared_ptr>> ProvidedCodecs = std::make_shared>>(); + + NTopic::IExecutor::TPtr SubsessionHandlersExecutor; + TAdaptiveLock Lock; }; diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index 560ac90a96..6c9b03d1af 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -29,12 +29,14 @@ TFederatedWriteSessionImpl::TFederatedWriteSessionImpl( std::shared_ptr connections, const TFederatedTopicClientSettings& clientSettings, std::shared_ptr observer, - std::shared_ptr>> codecs + std::shared_ptr>> codecs, + NTopic::IExecutor::TPtr subsessionHandlersExecutor ) : Settings(settings) , Connections(std::move(connections)) , SubclientSettings(FromFederated(clientSettings)) , ProvidedCodecs(std::move(codecs)) + , SubsessionHandlersExecutor(subsessionHandlersExecutor) , Observer(std::move(observer)) , AsyncInit(Observer->WaitForFirstState()) , FederationState(nullptr) @@ -70,15 +72,16 @@ void TFederatedWriteSessionImpl::IssueTokenIfAllowed() { } } -void TFederatedWriteSessionImpl::UpdateFederationStateImpl() { +std::shared_ptr TFederatedWriteSessionImpl::UpdateFederationStateImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); // Even after the user has called the Close method, transitioning the session to the CLOSING state, // we keep updating the federation state, as the session may still have some messages to send in its queues, // and for that we need to know the current state of the federation. if (SessionState < State::CLOSED) { FederationState = Observer->GetState(); - OnFederationStateUpdateImpl(); + return OnFederationStateUpdateImpl(); } + return {}; } void TFederatedWriteSessionImpl::Start() { @@ -103,15 +106,16 @@ void TFederatedWriteSessionImpl::Start() { }); } -void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) { +std::shared_ptr TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) { Y_ABORT_UNLESS(Lock.IsLocked()); ++SubsessionGeneration; + std::shared_ptr oldSubsession; + if (Subsession) { PendingToken.reset(); - OldSubsession = std::move(Subsession); - OldSubsession->Close(TDuration::Zero()); + std::swap(oldSubsession, Subsession); } auto clientSettings = SubclientSettings; @@ -121,11 +125,9 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) auto subclient = std::make_shared(Connections, clientSettings); auto handlers = NTopic::TWriteSessionSettings::TEventHandlers() - .HandlersExecutor(Settings.EventHandlers_.HandlersExecutor_) + .HandlersExecutor(SubsessionHandlersExecutor) .ReadyToAcceptHandler([selfCtx = SelfContext, generation = SubsessionGeneration](NTopic::TWriteSessionEvent::TReadyToAcceptEvent& ev) { if (auto self = selfCtx->LockShared()) { - TDeferredWrite deferred; - with_lock(self->Lock) { if (generation != self->SubsessionGeneration) { return; @@ -133,10 +135,8 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) Y_ABORT_UNLESS(!self->PendingToken.has_value()); self->PendingToken = std::move(ev.ContinuationToken); - self->PrepareDeferredWriteImpl(deferred); + self->MaybeWriteImpl(); } - - deferred.DoWrite(); } }) .AcksHandler([selfCtx = SelfContext](NTopic::TWriteSessionEvent::TAcksEvent& ev) { @@ -181,6 +181,8 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr db) Subsession = subclient->CreateWriteSession(wsSettings); CurrentDatabase = db; + + return oldSubsession; } std::pair, EStatus> SelectDatabaseByHashImpl( @@ -265,13 +267,13 @@ std::pair, EStatus> SelectDatabaseImpl( return SelectDatabaseByHashImpl(settings, dbInfos); } -void TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() { +std::shared_ptr TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); if (!FederationState->Status.IsSuccess()) { // The observer became stale, it won't try to get federation state anymore due to retry policy, // so there's no reason to keep the write session alive. CloseImpl(FederationState->Status.GetStatus(), NYql::TIssues(FederationState->Status.GetIssues())); - return; + return {}; } Y_ABORT_UNLESS(!FederationState->DbInfos.empty()); @@ -290,19 +292,22 @@ void TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() { LOG_LAZY(Log, TLOG_ERR, GetLogPrefixImpl() << message << ". Status: " << status); CloseImpl(status, NYql::TIssues{NYql::TIssue(message)}); } - return; + return {}; } RetryState.reset(); + std::shared_ptr oldSubsession; if (!DatabasesAreSame(preferrableDb, CurrentDatabase)) { LOG_LAZY(Log, TLOG_INFO, GetLogPrefixImpl() << "Start federated write session to database '" << preferrableDb->name() << "' (previous was " << (CurrentDatabase ? CurrentDatabase->name() : "") << ")" << " FederationState: " << *FederationState); - OpenSubsessionImpl(preferrableDb); + oldSubsession = OpenSubsessionImpl(preferrableDb); } ScheduleFederationStateUpdateImpl(UPDATE_FEDERATION_STATE_DELAY); + + return oldSubsession; } void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration delay) { @@ -312,8 +317,10 @@ void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration del if (auto self = selfCtx->LockShared()) { std::shared_ptr old; with_lock(self->Lock) { - self->UpdateFederationStateImpl(); - old = std::move(self->OldSubsession); + old = self->UpdateFederationStateImpl(); + } + if (old) { + old->Close(TDuration::Zero()); } } } @@ -376,8 +383,6 @@ void TFederatedWriteSessionImpl::WriteEncoded(NTopic::TContinuationToken&& token } void TFederatedWriteSessionImpl::WriteInternal(NTopic::TContinuationToken&&, TWrappedWriteMessage&& wrapped) { - TDeferredWrite deferred(Subsession); - with_lock(Lock) { ClientHasToken = false; if (!wrapped.Message.CreateTimestamp_.has_value()) { @@ -385,15 +390,13 @@ void TFederatedWriteSessionImpl::WriteInternal(NTopic::TContinuationToken&&, TWr } BufferFreeSpace -= wrapped.Message.Data.size(); OriginalMessagesToPassDown.emplace_back(std::move(wrapped)); - PrepareDeferredWriteImpl(deferred); + MaybeWriteImpl(); } - deferred.DoWrite(); - IssueTokenIfAllowed(); } -bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferred) { +bool TFederatedWriteSessionImpl::MaybeWriteImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); if (!PendingToken.has_value()) { return false; @@ -403,9 +406,7 @@ bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferr } OriginalMessagesToGetAck.push_back(std::move(OriginalMessagesToPassDown.front())); OriginalMessagesToPassDown.pop_front(); - deferred.Writer = Subsession; - deferred.Token.emplace(std::move(*PendingToken)); - deferred.Message.emplace(std::move(OriginalMessagesToGetAck.back().Message)); + Subsession->Write(std::move(*PendingToken), std::move(OriginalMessagesToGetAck.back().Message)); PendingToken.reset(); return true; } diff --git a/src/client/federated_topic/impl/federated_write_session.h b/src/client/federated_topic/impl/federated_write_session.h index f70c78e98d..121eb63fbd 100644 --- a/src/client/federated_topic/impl/federated_write_session.h +++ b/src/client/federated_topic/impl/federated_write_session.h @@ -26,7 +26,8 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, std::shared_ptr connections, const TFederatedTopicClientSettings& clientSetttings, std::shared_ptr observer, - std::shared_ptr>> codecs); + std::shared_ptr>> codecs, + NTopic::IExecutor::TPtr subsessionHandlersExecutor); ~TFederatedWriteSessionImpl() = default; @@ -61,40 +62,22 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, } }; - struct TDeferredWrite { - TDeferredWrite() {} - explicit TDeferredWrite(std::shared_ptr writer) - : Writer(std::move(writer)) { - } - - void DoWrite() { - if (!Token.has_value() && !Message.has_value()) { - return; - } - Y_ABORT_UNLESS(Token.has_value() && Message.has_value()); - return Writer->Write(std::move(*Token), std::move(*Message)); - } - - std::shared_ptr Writer; - std::optional Token; - std::optional Message; - }; - private: void Start(); - void OpenSubsessionImpl(std::shared_ptr db); - void OnFederationStateUpdateImpl(); + std::shared_ptr OpenSubsessionImpl(std::shared_ptr db); + std::shared_ptr UpdateFederationStateImpl(); + std::shared_ptr OnFederationStateUpdateImpl(); + void ScheduleFederationStateUpdateImpl(TDuration delay); void WriteInternal(NTopic::TContinuationToken&&, TWrappedWriteMessage&& message); - bool PrepareDeferredWriteImpl(TDeferredWrite& deferred); + bool MaybeWriteImpl(); void CloseImpl(EStatus statusCode, NYql::TIssues&& issues); void CloseImpl(NTopic::TSessionClosedEvent const& ev); bool MessageQueuesAreEmptyImpl() const; - void UpdateFederationStateImpl(); void IssueTokenIfAllowed(); @@ -106,6 +89,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, std::shared_ptr Connections; const NTopic::TTopicClientSettings SubclientSettings; std::shared_ptr>> ProvidedCodecs; + NTopic::IExecutor::TPtr SubsessionHandlersExecutor; NTopic::IRetryPolicy::IRetryState::TPtr RetryState; std::shared_ptr Observer; @@ -124,7 +108,6 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer, size_t SubsessionGeneration = 0; std::shared_ptr Subsession; - std::shared_ptr OldSubsession; std::shared_ptr ClientEventsQueue; @@ -155,8 +138,9 @@ class TFederatedWriteSession : public NTopic::IWriteSession, std::shared_ptr connections, const TFederatedTopicClientSettings& clientSettings, std::shared_ptr observer, - std::shared_ptr>> codecs) - : TContextOwner(settings, std::move(connections), clientSettings, std::move(observer), codecs) {} + std::shared_ptr>> codecs, + NTopic::IExecutor::TPtr subsessionHandlersExecutor) + : TContextOwner(settings, std::move(connections), clientSettings, std::move(observer), codecs, subsessionHandlersExecutor) {} NThreading::TFuture WaitEvent() override { return TryGetImpl()->WaitEvent(); From 8b95d53a326edf342e3aafcd72e3ae6830bc9fb0 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 17:27:54 +0000 Subject: [PATCH 22/35] Moved commit "ydb_topic write session: lower log level for retryable errors" from ydb repo --- .../persqueue_public/impl/write_session_impl.cpp | 13 ++++--------- src/client/topic/impl/write_session_impl.cpp | 6 +++--- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/src/client/persqueue_public/impl/write_session_impl.cpp b/src/client/persqueue_public/impl/write_session_impl.cpp index 54ac517ca0..67d5e95143 100644 --- a/src/client/persqueue_public/impl/write_session_impl.cpp +++ b/src/client/persqueue_public/impl/write_session_impl.cpp @@ -93,10 +93,7 @@ TWriteSessionImpl::THandleResult TWriteSessionImpl::RestartImpl(const TPlainStat LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "Write session is aborting and will not restart"); return result; } - LOG_LAZY(DbDriverState->Log, TLOG_ERR, - LogPrefix() << "Got error. Status: " << status.Status - << ". Description: " << IssuesSingleLineString(status.Issues) - ); + SessionEstablished = false; if (!RetryState) { RetryState = Settings.RetryPolicy_->CreateRetryState(); @@ -106,13 +103,11 @@ TWriteSessionImpl::THandleResult TWriteSessionImpl::RestartImpl(const TPlainStat if (nextDelay) { result.StartDelay = *nextDelay; result.DoRestart = true; - LOG_LAZY(DbDriverState->Log, - TLOG_DEBUG, - LogPrefix() << "Write session will restart in " << result.StartDelay.MilliSeconds() << " ms" - ); + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Got error. " << status.ToDebugString()); + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Write session will restart in " << result.StartDelay); ResetForRetryImpl(); - } else { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Got error. " << status.ToDebugString()); LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Write session will not restart after a fatal error"); result.DoStop = true; CheckHandleResultImpl(result); diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index e967d40b08..3d6c5514a5 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -171,7 +171,6 @@ TWriteSessionImpl::THandleResult TWriteSessionImpl::RestartImpl(const TPlainStat LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefix() << "Write session is aborting and will not restart"); return result; } - LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Got error. " << status.ToDebugString()); SessionEstablished = false; // Keep DirectWriteToPartitionId value on temporary errors. @@ -192,10 +191,11 @@ TWriteSessionImpl::THandleResult TWriteSessionImpl::RestartImpl(const TPlainStat if (nextDelay) { result.StartDelay = *nextDelay; result.DoRestart = true; - LOG_LAZY(DbDriverState->Log, TLOG_WARNING, LogPrefix() << "Write session will restart in " << result.StartDelay); + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Got error. " << status.ToDebugString()); + LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefix() << "Write session will restart in " << result.StartDelay); ResetForRetryImpl(); - } else { + LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Got error. " << status.ToDebugString()); LOG_LAZY(DbDriverState->Log, TLOG_ERR, LogPrefix() << "Write session will not restart after a fatal error"); result.DoStop = true; CheckHandleResultImpl(result); From 566ac7a2dfd6653fdaf042b149d2b881fa8efddc Mon Sep 17 00:00:00 2001 From: Ilnaz Nizametdinov Date: Mon, 26 Aug 2024 17:31:25 +0000 Subject: [PATCH 23/35] Moved commit "Enable/disable ssl connections, return connection_string in API " from ydb repo --- include/ydb-cpp-sdk/client/draft/ydb_replication.h | 1 + src/client/draft/ydb_replication.cpp | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/include/ydb-cpp-sdk/client/draft/ydb_replication.h b/include/ydb-cpp-sdk/client/draft/ydb_replication.h index 749e00f5e4..7528395245 100644 --- a/include/ydb-cpp-sdk/client/draft/ydb_replication.h +++ b/include/ydb-cpp-sdk/client/draft/ydb_replication.h @@ -51,6 +51,7 @@ class TConnectionParams: private TCommonClientSettings { const std::string& GetDiscoveryEndpoint() const; const std::string& GetDatabase() const; + bool GetEnableSsl() const; ECredentials GetCredentials() const; const TStaticCredentials& GetStaticCredentials() const; diff --git a/src/client/draft/ydb_replication.cpp b/src/client/draft/ydb_replication.cpp index 7e4ce3e0ce..b2be0c7eed 100644 --- a/src/client/draft/ydb_replication.cpp +++ b/src/client/draft/ydb_replication.cpp @@ -19,6 +19,7 @@ namespace NReplication { TConnectionParams::TConnectionParams(const Ydb::Replication::ConnectionParams& params) { DiscoveryEndpoint(params.endpoint()); Database(params.database()); + SslCredentials(params.enable_ssl()); switch (params.credentials_case()) { case Ydb::Replication::ConnectionParams::kStaticCredentials: @@ -47,6 +48,10 @@ const std::string& TConnectionParams::GetDatabase() const { return *Database_; } +bool TConnectionParams::GetEnableSsl() const { + return SslCredentials_->IsEnabled; +} + TConnectionParams::ECredentials TConnectionParams::GetCredentials() const { return static_cast(Credentials_.index()); } From 56d4674454dee89311c97e21fd57ed18ec5666ea Mon Sep 17 00:00:00 2001 From: niksaveliev Date: Mon, 26 Aug 2024 17:43:11 +0000 Subject: [PATCH 24/35] Moved commit "Add few auto partitioning fields to describe and SDK" from ydb repo --- .../ydb-cpp-sdk/client/topic/control_plane.h | 24 +++++++++++++----- src/client/topic/impl/topic.cpp | 25 +++++++++++++++++++ 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/include/ydb-cpp-sdk/client/topic/control_plane.h b/include/ydb-cpp-sdk/client/topic/control_plane.h index 429b1a6a20..42e4b84052 100644 --- a/include/ydb-cpp-sdk/client/topic/control_plane.h +++ b/include/ydb-cpp-sdk/client/topic/control_plane.h @@ -145,14 +145,21 @@ class TPartitionInfo { const std::optional& GetPartitionConsumerStats() const; const std::optional& GetPartitionLocation() const; + const std::optional& GetFromBound() const; + const std::optional& GetToBound() const; + private: uint64_t PartitionId_; bool Active_; std::vector ChildPartitionIds_; std::vector ParentPartitionIds_; + std::optional PartitionStats_; std::optional PartitionConsumerStats_; std::optional PartitionLocation_; + + std::optional FromBound_; + std::optional ToBound_; }; struct TAlterPartitioningSettings; @@ -187,8 +194,8 @@ friend struct TAutoPartitioningSettingsBuilder; struct TAlterAutoPartitioningSettings { using TSelf = TAlterAutoPartitioningSettings; - public: - TAlterAutoPartitioningSettings(TAlterPartitioningSettings& parent): Parent_(parent) {} +public: + TAlterAutoPartitioningSettings(TAlterPartitioningSettings& parent): Parent_(parent) {} FLUENT_SETTING_OPTIONAL(EAutoPartitioningStrategy, Strategy); FLUENT_SETTING_OPTIONAL(TDuration, StabilizationWindow); @@ -197,8 +204,8 @@ struct TAlterAutoPartitioningSettings { TAlterPartitioningSettings& EndAlterAutoPartitioningSettings() { return Parent_; }; - private: - TAlterPartitioningSettings& Parent_; +private: + TAlterPartitioningSettings& Parent_; }; class TPartitioningSettings { @@ -207,11 +214,11 @@ class TPartitioningSettings { public: TPartitioningSettings() : MinActivePartitions_(0), MaxActivePartitions_(0), PartitionCountLimit_(0), AutoPartitioningSettings_(){} TPartitioningSettings(const Ydb::Topic::PartitioningSettings& settings); - TPartitioningSettings(ui64 minActivePartitions, ui64 maxActivePartitions, TAutoPartitioningSettings autoscalingSettings = {}) + TPartitioningSettings(uint64_t minActivePartitions, uint64_t maxActivePartitions, TAutoPartitioningSettings autoPartitioning = {}) : MinActivePartitions_(minActivePartitions) , MaxActivePartitions_(maxActivePartitions) , PartitionCountLimit_(0) - , AutoPartitioningSettings_(autoscalingSettings) + , AutoPartitioningSettings_(autoPartitioning) { } @@ -460,6 +467,11 @@ struct TConsumerSettings { return *this; } + TConsumerSettings& SetImportant(bool isImportant) { + Important_ = isImportant; + return *this; + } + TSettings& EndAddConsumer() { return Parent_; }; private: diff --git a/src/client/topic/impl/topic.cpp b/src/client/topic/impl/topic.cpp index a5b31dd5ea..677a75b96b 100644 --- a/src/client/topic/impl/topic.cpp +++ b/src/client/topic/impl/topic.cpp @@ -395,6 +395,7 @@ TPartitionInfo::TPartitionInfo(const Ydb::Topic::DescribeTopicResult::PartitionI for (const auto& partId : partitionInfo.parent_partition_ids()) { ParentPartitionIds_.push_back(partId); } + if (partitionInfo.has_partition_stats()) { PartitionStats_ = TPartitionStats{partitionInfo.partition_stats()}; } @@ -402,6 +403,14 @@ TPartitionInfo::TPartitionInfo(const Ydb::Topic::DescribeTopicResult::PartitionI if (partitionInfo.has_partition_location()) { PartitionLocation_ = TPartitionLocation{partitionInfo.partition_location()}; } + + if (partitionInfo.has_key_range() && partitionInfo.key_range().has_from_bound()) { + FromBound_ = std::string{partitionInfo.key_range().from_bound()}; + } + + if (partitionInfo.has_key_range() && partitionInfo.key_range().has_to_bound()) { + ToBound_ = std::string{partitionInfo.key_range().to_bound()}; + } } TPartitionInfo::TPartitionInfo(const Ydb::Topic::DescribeConsumerResult::PartitionInfo& partitionInfo) @@ -437,6 +446,14 @@ const std::optional& TPartitionInfo::GetPartitionLocation() return PartitionLocation_; } +const std::vector TPartitionInfo::GetChildPartitionIds() const { + return ChildPartitionIds_; +} + +const std::vector TPartitionInfo::GetParentPartitionIds() const { + return ParentPartitionIds_; +} + bool TPartitionInfo::GetActive() const { return Active_; } @@ -445,6 +462,14 @@ uint64_t TPartitionInfo::GetPartitionId() const { return PartitionId_; } +const std::optional& TPartitionInfo::GetFromBound() const { + return FromBound_; +} + +const std::optional& TPartitionInfo::GetToBound() const { + return ToBound_; +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // TTopicClient From db347eb925e4b9edcffd4f9983ebcad5bae48ae1 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 17:49:27 +0000 Subject: [PATCH 25/35] Moved commit "Call ByteSizeLong on each WriteRequest separately" from ydb repo --- src/client/topic/impl/write_session_impl.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 3d6c5514a5..6b351ce11c 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -1371,8 +1371,11 @@ void TWriteSessionImpl::SendImpl() { auto* writeRequest = clientMessage.mutable_write_request(); ui32 prevCodec = 0; + + ui64 currentSize = 0; + // Send blocks while we can without messages reordering. - while (IsReadyToSendNextImpl() && clientMessage.ByteSizeLong() < GetMaxGrpcMessageSize()) { + while (IsReadyToSendNextImpl() && currentSize < GetMaxGrpcMessageSize()) { const auto& block = PackedMessagesToSend.top(); Y_ABORT_UNLESS(block.Valid); if (writeRequest->messages_size() > 0 && prevCodec != block.CodecID) { @@ -1420,6 +1423,8 @@ void TWriteSessionImpl::SendImpl() { moveBlock.Move(block); SentPackedMessage.emplace(std::move(moveBlock)); PackedMessagesToSend.pop(); + + currentSize += writeRequest->ByteSizeLong(); } UpdateTokenIfNeededImpl(); LOG_LAZY(DbDriverState->Log, From a56ab71f633233be8dcc3425b2d4770678697369 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 17:52:11 +0000 Subject: [PATCH 26/35] Moved commit "Move unacknowledged messages back to OriginalMessagesToPassDown queue" from ydb repo --- .../federated_topic/impl/federated_write_session.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index 6c9b03d1af..ea0a2c30be 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -174,6 +174,14 @@ std::shared_ptr TFederatedWriteSessionImpl::OpenSubsessio } }); + { + // Unacknowledged messages should be resent. + for (auto& msg : OriginalMessagesToPassDown) { + OriginalMessagesToGetAck.emplace_back(std::move(msg)); + } + OriginalMessagesToPassDown = std::move(OriginalMessagesToGetAck); + } + NTopic::TWriteSessionSettings wsSettings = Settings; wsSettings // .MaxMemoryUsage(Settings.MaxMemoryUsage_) // to fix if split not by half on creation From 12b4b8e19e873164f652a6bc0f2b1ba066be44fd Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 17:53:29 +0000 Subject: [PATCH 27/35] Moved commit "Ignore message acks from previous subsessions" from ydb repo --- src/client/federated_topic/impl/federated_write_session.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/federated_topic/impl/federated_write_session.cpp index ea0a2c30be..dc33eb2aab 100644 --- a/src/client/federated_topic/impl/federated_write_session.cpp +++ b/src/client/federated_topic/impl/federated_write_session.cpp @@ -139,9 +139,13 @@ std::shared_ptr TFederatedWriteSessionImpl::OpenSubsessio } } }) - .AcksHandler([selfCtx = SelfContext](NTopic::TWriteSessionEvent::TAcksEvent& ev) { + .AcksHandler([selfCtx = SelfContext, generation = SubsessionGeneration](NTopic::TWriteSessionEvent::TAcksEvent& ev) { if (auto self = selfCtx->LockShared()) { with_lock(self->Lock) { + if (generation != self->SubsessionGeneration) { + return; + } + Y_ABORT_UNLESS(ev.Acks.size() <= self->OriginalMessagesToGetAck.size()); for (size_t i = 0; i < ev.Acks.size(); ++i) { From 6dfdca852adf4ec42dce48b3d475340daa8f220d Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 18:01:58 +0000 Subject: [PATCH 28/35] Moved commit "Use a separate lock for Processor->Write calls" from ydb repo --- src/client/topic/impl/write_session_impl.cpp | 32 ++++++++++++++------ src/client/topic/impl/write_session_impl.h | 7 ++++- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 6b351ce11c..7cbed1a968 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -978,7 +978,8 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess FirstTokenSent = true; } // Kickstart send after session reestablishment - SendImpl(); + FormGrpcMessagesImpl(); + SendGrpcMessages(); break; } case TServerMessage::kWriteResponse: { @@ -1160,12 +1161,14 @@ void TWriteSessionImpl::CompressImpl(TBlock&& block_) { void TWriteSessionImpl::OnCompressed(TBlock&& block, bool isSyncCompression) { TMemoryUsageChange memoryUsage; - if (!isSyncCompression) { - std::lock_guard guard(Lock); + if (isSyncCompression) { + // The Lock is already held somewhere up the stack. memoryUsage = OnCompressedImpl(std::move(block)); } else { + std::lock_guard guard(Lock); memoryUsage = OnCompressedImpl(std::move(block)); } + SendGrpcMessages(); if (memoryUsage.NowOk && !memoryUsage.WasOk) { EventsQueue->PushEvent(TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); } @@ -1181,7 +1184,7 @@ TMemoryUsageChange TWriteSessionImpl::OnCompressedImpl(TBlock&& block) { (*Counters->BytesInflightCompressed) += block.Data.size(); PackedMessagesToSend.emplace(std::move(block)); - SendImpl(); + FormGrpcMessagesImpl(); return memoryUsage; } @@ -1298,7 +1301,7 @@ size_t TWriteSessionImpl::WriteBatchImpl() { } CurrentBatch.Reset(); if (skipCompression) { - SendImpl(); + FormGrpcMessagesImpl(); } return size; } @@ -1362,7 +1365,16 @@ bool TWriteSessionImpl::TxIsChanged(const Ydb::Topic::StreamWriteMessage_WriteRe return GetTransactionId(*writeRequest) != GetTransactionId(OriginalMessagesToSend.front().Tx); } -void TWriteSessionImpl::SendImpl() { +void TWriteSessionImpl::SendGrpcMessages() { + with_lock(ProcessorLock) { + TClientMessage message; + while (GrpcMessagesToSend.Dequeue(&message)) { + Processor->Write(std::move(message)); + } + } +} + +void TWriteSessionImpl::FormGrpcMessagesImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); // External cycle splits ready blocks into multiple gRPC messages. Current gRPC message size hard limit is 64MiB. @@ -1433,7 +1445,7 @@ void TWriteSessionImpl::SendImpl() { << OriginalMessagesToSend.size() << " left), first sequence number is " << writeRequest->messages(0).seq_no() ); - Processor->Write(std::move(clientMessage)); + GrpcMessagesToSend.Enqueue(std::move(clientMessage)); } } @@ -1495,8 +1507,10 @@ void TWriteSessionImpl::HandleWakeUpImpl() { return; } if (auto self = cbContext->LockShared()) { - std::lock_guard guard(self->Lock); - self->HandleWakeUpImpl(); + with_lock(self->Lock) { + self->HandleWakeUpImpl(); + } + self->SendGrpcMessages(); } }; if (TInstant::Now() - LastTokenUpdate > UPDATE_TOKEN_PERIOD) { diff --git a/src/client/topic/impl/write_session_impl.h b/src/client/topic/impl/write_session_impl.h index d8d78a487e..230a1bca28 100644 --- a/src/client/topic/impl/write_session_impl.h +++ b/src/client/topic/impl/write_session_impl.h @@ -5,6 +5,7 @@ #include #include +#include namespace NYdb::NTopic { @@ -390,7 +391,8 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, uint64_t GetNextIdImpl(const std::optional& seqNo); uint64_t GetSeqNoImpl(uint64_t id); uint64_t GetIdImpl(uint64_t seqNo); - void SendImpl(); + void FormGrpcMessagesImpl(); + void SendGrpcMessages(); void AbortImpl(); void CloseImpl(EStatus statusCode, NYql::TIssues&& issues); void CloseImpl(EStatus statusCode, const std::string& message); @@ -451,6 +453,9 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, std::queue SentOriginalMessages; std::queue SentPackedMessage; + TLockFreeQueue GrpcMessagesToSend; + TAdaptiveLock ProcessorLock; + const size_t MaxBlockSize = std::numeric_limits::max(); const size_t MaxBlockMessageCount = 1; //!< Max message count that can be packed into a single block. In block version 0 is equal to 1 for compatibility bool Connected = false; From 985b61a092e8ebd5feb333bc4275dea7c3cf754c Mon Sep 17 00:00:00 2001 From: Alek5andr-Kotov Date: Mon, 26 Aug 2024 18:07:05 +0000 Subject: [PATCH 29/35] Moved commit "code EES_WRITTEN_IN_TX" from ydb repo --- .../ydb-cpp-sdk/client/topic/write_events.h | 3 ++- src/client/topic/impl/write_session_impl.cpp | 22 ++++++++++++------- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/include/ydb-cpp-sdk/client/topic/write_events.h b/include/ydb-cpp-sdk/client/topic/write_events.h index 56413200cd..46d195f688 100644 --- a/include/ydb-cpp-sdk/client/topic/write_events.h +++ b/include/ydb-cpp-sdk/client/topic/write_events.h @@ -56,7 +56,8 @@ struct TWriteSessionEvent { enum EEventState { EES_WRITTEN, //! Successfully written. EES_ALREADY_WRITTEN, //! Skipped on SeqNo deduplication. - EES_DISCARDED //! In case of destruction of writer or retry policy discarded future retries in this writer. + EES_DISCARDED, //! In case of destruction of writer or retry policy discarded future retries in this writer. + EES_WRITTEN_IN_TX, //! Successfully written in tx. }; //! Details of successfully written message. struct TWrittenMessageDetails { diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 7cbed1a968..9dbc23b4b4 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -1002,17 +1002,23 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess writeStat->PartitionQuotedTime = durationConv(stat.partition_quota_wait_time()); writeStat->TopicQuotedTime = durationConv(stat.topic_quota_wait_time()); - for (size_t messageIndex = 0, endIndex = batchWriteResponse.acks_size(); messageIndex != endIndex; ++messageIndex) { + for (const auto& ack : batchWriteResponse.acks()) { // TODO: Fill writer statistics - auto ack = batchWriteResponse.acks(messageIndex); uint64_t sequenceNumber = ack.seq_no(); - Y_ABORT_UNLESS(ack.has_written() || ack.has_skipped()); - auto msgWriteStatus = ack.has_written() - ? TWriteSessionEvent::TWriteAck::EES_WRITTEN - : (ack.skipped().reason() == Ydb::Topic::StreamWriteMessage_WriteResponse_WriteAck_Skipped_Reason::StreamWriteMessage_WriteResponse_WriteAck_Skipped_Reason_REASON_ALREADY_WRITTEN - ? TWriteSessionEvent::TWriteAck::EES_ALREADY_WRITTEN - : TWriteSessionEvent::TWriteAck::EES_DISCARDED); + Y_ABORT_UNLESS(ack.has_written() || ack.has_skipped() || ack.has_written_in_tx()); + + TWriteSessionEvent::TWriteAck::EEventState msgWriteStatus; + if (ack.has_written_in_tx()) { + msgWriteStatus = TWriteSessionEvent::TWriteAck::EES_WRITTEN_IN_TX; + } else if (ack.has_written()) { + msgWriteStatus = TWriteSessionEvent::TWriteAck::EES_WRITTEN; + } else { + msgWriteStatus = + (ack.skipped().reason() == Ydb::Topic::StreamWriteMessage_WriteResponse_WriteAck_Skipped_Reason::StreamWriteMessage_WriteResponse_WriteAck_Skipped_Reason_REASON_ALREADY_WRITTEN) + ? TWriteSessionEvent::TWriteAck::EES_ALREADY_WRITTEN + : TWriteSessionEvent::TWriteAck::EES_DISCARDED; + } uint64_t offset = ack.has_written() ? ack.written().offset() : 0; From 06422e1d166da03ac9201b418253f68fd89ab689 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Mon, 26 Aug 2024 18:09:37 +0000 Subject: [PATCH 30/35] Moved commit "Fix TNodeRegistrationResult" from ydb repo --- src/client/discovery/discovery.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/client/discovery/discovery.cpp b/src/client/discovery/discovery.cpp index c7b04032ee..2295886e5c 100644 --- a/src/client/discovery/discovery.cpp +++ b/src/client/discovery/discovery.cpp @@ -120,7 +120,7 @@ uint64_t TNodeRegistrationResult::GetScopePathId() const { } bool TNodeRegistrationResult::HasScopePathId() const { - return ScopePathId_.value(); + return ScopePathId_.has_value(); } bool TNodeRegistrationResult::HasNodeName() const { From 6b29678a9d6c96f6c8ead5909c505565baf1d522 Mon Sep 17 00:00:00 2001 From: qyryq Date: Mon, 26 Aug 2024 18:16:52 +0000 Subject: [PATCH 31/35] Moved commit "Revert separate lock for Processor->Write calls" from ydb repo --- src/client/topic/impl/write_session_impl.cpp | 32 ++++++-------------- src/client/topic/impl/write_session_impl.h | 7 +---- 2 files changed, 10 insertions(+), 29 deletions(-) diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index 9dbc23b4b4..1cc47bf517 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -978,8 +978,7 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess FirstTokenSent = true; } // Kickstart send after session reestablishment - FormGrpcMessagesImpl(); - SendGrpcMessages(); + SendImpl(); break; } case TServerMessage::kWriteResponse: { @@ -1167,14 +1166,12 @@ void TWriteSessionImpl::CompressImpl(TBlock&& block_) { void TWriteSessionImpl::OnCompressed(TBlock&& block, bool isSyncCompression) { TMemoryUsageChange memoryUsage; - if (isSyncCompression) { - // The Lock is already held somewhere up the stack. + if (!isSyncCompression) { + std::lock_guard guard(Lock); memoryUsage = OnCompressedImpl(std::move(block)); } else { - std::lock_guard guard(Lock); memoryUsage = OnCompressedImpl(std::move(block)); } - SendGrpcMessages(); if (memoryUsage.NowOk && !memoryUsage.WasOk) { EventsQueue->PushEvent(TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); } @@ -1190,7 +1187,7 @@ TMemoryUsageChange TWriteSessionImpl::OnCompressedImpl(TBlock&& block) { (*Counters->BytesInflightCompressed) += block.Data.size(); PackedMessagesToSend.emplace(std::move(block)); - FormGrpcMessagesImpl(); + SendImpl(); return memoryUsage; } @@ -1307,7 +1304,7 @@ size_t TWriteSessionImpl::WriteBatchImpl() { } CurrentBatch.Reset(); if (skipCompression) { - FormGrpcMessagesImpl(); + SendImpl(); } return size; } @@ -1371,16 +1368,7 @@ bool TWriteSessionImpl::TxIsChanged(const Ydb::Topic::StreamWriteMessage_WriteRe return GetTransactionId(*writeRequest) != GetTransactionId(OriginalMessagesToSend.front().Tx); } -void TWriteSessionImpl::SendGrpcMessages() { - with_lock(ProcessorLock) { - TClientMessage message; - while (GrpcMessagesToSend.Dequeue(&message)) { - Processor->Write(std::move(message)); - } - } -} - -void TWriteSessionImpl::FormGrpcMessagesImpl() { +void TWriteSessionImpl::SendImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); // External cycle splits ready blocks into multiple gRPC messages. Current gRPC message size hard limit is 64MiB. @@ -1451,7 +1439,7 @@ void TWriteSessionImpl::FormGrpcMessagesImpl() { << OriginalMessagesToSend.size() << " left), first sequence number is " << writeRequest->messages(0).seq_no() ); - GrpcMessagesToSend.Enqueue(std::move(clientMessage)); + Processor->Write(std::move(clientMessage)); } } @@ -1513,10 +1501,8 @@ void TWriteSessionImpl::HandleWakeUpImpl() { return; } if (auto self = cbContext->LockShared()) { - with_lock(self->Lock) { - self->HandleWakeUpImpl(); - } - self->SendGrpcMessages(); + std::lock_guard guard(self->Lock); + self->HandleWakeUpImpl(); } }; if (TInstant::Now() - LastTokenUpdate > UPDATE_TOKEN_PERIOD) { diff --git a/src/client/topic/impl/write_session_impl.h b/src/client/topic/impl/write_session_impl.h index 230a1bca28..d8d78a487e 100644 --- a/src/client/topic/impl/write_session_impl.h +++ b/src/client/topic/impl/write_session_impl.h @@ -5,7 +5,6 @@ #include #include -#include namespace NYdb::NTopic { @@ -391,8 +390,7 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, uint64_t GetNextIdImpl(const std::optional& seqNo); uint64_t GetSeqNoImpl(uint64_t id); uint64_t GetIdImpl(uint64_t seqNo); - void FormGrpcMessagesImpl(); - void SendGrpcMessages(); + void SendImpl(); void AbortImpl(); void CloseImpl(EStatus statusCode, NYql::TIssues&& issues); void CloseImpl(EStatus statusCode, const std::string& message); @@ -453,9 +451,6 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, std::queue SentOriginalMessages; std::queue SentPackedMessage; - TLockFreeQueue GrpcMessagesToSend; - TAdaptiveLock ProcessorLock; - const size_t MaxBlockSize = std::numeric_limits::max(); const size_t MaxBlockMessageCount = 1; //!< Max message count that can be packed into a single block. In block version 0 is equal to 1 for compatibility bool Connected = false; From f48930a8254c000bf77f178db03276ccbf5753fb Mon Sep 17 00:00:00 2001 From: Nikolay Shestakov Date: Mon, 26 Aug 2024 18:21:01 +0000 Subject: [PATCH 32/35] Moved commit "Add tests for pqv1" from ydb repo --- src/client/topic/ut/local_partition_ut.cpp | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/client/topic/ut/local_partition_ut.cpp b/src/client/topic/ut/local_partition_ut.cpp index 68f391e315..a5b6718016 100644 --- a/src/client/topic/ut/local_partition_ut.cpp +++ b/src/client/topic/ut/local_partition_ut.cpp @@ -677,18 +677,19 @@ namespace NYdb::NTopic::NTests { .MessageGroupId(TEST_MESSAGE_GROUP_ID) .DirectWriteToPartition(true); auto writeSession = client.CreateSimpleBlockingWriteSession(writeSettings); - TTestReadSession ReadSession("Session-0", client, 2); + auto ReadSession = NPQ::NTest::CreateTestReadSession({ .Name="Session-0", .Setup=setup, .Sdk = NPQ::NTest::SdkVersion::Topic, .ExpectedMessagesCount = 2 }); - UNIT_ASSERT(writeSession->Write(Msg("message_1.1", 2))); + + UNIT_ASSERT(writeSession->Write(NPQ::NTest::Msg("message_1.1", 2))); ui64 txId = 1006; - SplitPartition(setup, ++txId, 0, "a"); + NPQ::NTest::SplitPartition(setup, ++txId, 0, "a"); - UNIT_ASSERT(writeSession->Write(Msg("message_1.2", 3))); + UNIT_ASSERT(writeSession->Write(NPQ::NTest::Msg("message_1.2", 3))); - ReadSession.WaitAllMessages(); + ReadSession->WaitAllMessages(); - for (const auto& info : ReadSession.Impl->ReceivedMessages) { + for (const auto& info : ReadSession->GetReceivedMessages()) { if (info.Data == "message_1.1") { UNIT_ASSERT_EQUAL(0, info.PartitionId); UNIT_ASSERT_EQUAL(2, info.SeqNo); @@ -724,7 +725,7 @@ namespace NYdb::NTopic::NTests { auto const events = tracingBackend->GetEvents(); UNIT_ASSERT(expected.Matches(events)); - ReadSession.Close(); + ReadSession->Close(); } } } From ce4222cb57811b5c8c875a97902ca3bb5969fff2 Mon Sep 17 00:00:00 2001 From: stanislav_shchetinin Date: Mon, 26 Aug 2024 18:26:23 +0000 Subject: [PATCH 33/35] Moved commit "Added new versions of RetryQuery" from ydb repo --- include/ydb-cpp-sdk/client/query/client.h | 25 +++++++++++++++++---- src/client/query/client.cpp | 27 ++++++++++++++++++++--- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/include/ydb-cpp-sdk/client/query/client.h b/include/ydb-cpp-sdk/client/query/client.h index dbf611825b..a25f83531a 100644 --- a/include/ydb-cpp-sdk/client/query/client.h +++ b/include/ydb-cpp-sdk/client/query/client.h @@ -14,7 +14,11 @@ namespace NYdb { namespace NRetry::Async { template class TRetryContext; - } + } // namespace NRetry::Async + namespace NRetry::Sync { + template + class TRetryContext; + } // namespace NRetry::Sync } namespace NYdb::NQuery { @@ -55,10 +59,15 @@ class TSession; class TQueryClient { friend class TSession; friend class NRetry::Async::TRetryContext; + friend class NRetry::Async::TRetryContext; + friend class NRetry::Sync::TRetryContext; public: - using TQueryFunc = std::function; - using TQueryWithoutSessionFunc = std::function; + using TQueryResultFunc = std::function; + using TQueryFunc = std::function; + using TQuerySyncFunc = std::function; + using TQueryWithoutSessionFunc = std::function; + using TQueryWithoutSessionSyncFunc = std::function; using TSettings = TClientSettings; using TSession = TSession; using TCreateSessionSettings = TCreateSessionSettings; @@ -79,7 +88,15 @@ class TQueryClient { TAsyncExecuteQueryIterator StreamExecuteQuery(const std::string& query, const TTxControl& txControl, const TParams& params, const TExecuteQuerySettings& settings = TExecuteQuerySettings()); - TAsyncExecuteQueryResult RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); + TAsyncExecuteQueryResult RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); + + TAsyncStatus RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); + + TAsyncStatus RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); + + TStatus RetryQuery(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); + + TStatus RetryQuery(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings()); TAsyncExecuteQueryResult RetryQuery(const std::string& query, const TTxControl& txControl, TDuration timeout, bool isIndempotent); diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 03fccfb8bc..6f3e6636f4 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -21,7 +21,8 @@ namespace NYdb::NQuery { -using TRetryContextAsync = NRetry::Async::TRetryContext; +using TRetryContextResultAsync = NRetry::Async::TRetryContext; +using TRetryContextAsync = NRetry::Async::TRetryContext; NYdb::NRetry::TRetryOperationSettings GetRetrySettings(TDuration timeout, bool isIndempotent) { return NYdb::NRetry::TRetryOperationSettings() @@ -577,12 +578,32 @@ int64_t TQueryClient::GetCurrentPoolSize() const { return Impl_->GetCurrentPoolSize(); } -TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings) +TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings) { + TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); + return ctx->Execute(); +} + +TAsyncStatus TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings) { TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); return ctx->Execute(); } +TAsyncStatus TQueryClient::RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRetryOperationSettings settings) { + TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithoutSession(*this, std::move(queryFunc), settings)); + return ctx->Execute(); +} + +TStatus TQueryClient::RetryQuery(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings) { + NRetry::Sync::TRetryWithSession ctx(*this, queryFunc, settings); + return ctx.Execute(); +} + +TStatus TQueryClient::RetryQuery(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings) { + NRetry::Sync::TRetryWithoutSession ctx(*this, queryFunc, settings); + return ctx.Execute(); +} + TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, const TTxControl& txControl, TDuration timeout, bool isIndempotent) { @@ -590,7 +611,7 @@ TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, cons auto queryFunc = [&query, &txControl](TSession session, TDuration duration) -> TAsyncExecuteQueryResult { return session.ExecuteQuery(query, txControl, TExecuteQuerySettings().ClientTimeout(duration)); }; - TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); + TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings)); return ctx->Execute(); } From 18d21dd62f60073f97c2debb5e4e8385958f26f3 Mon Sep 17 00:00:00 2001 From: Daniil Demin Date: Mon, 26 Aug 2024 18:33:54 +0000 Subject: [PATCH 34/35] Moved commit "Restore indexes from backup with the original partitioning" from ydb repo --- include/ydb-cpp-sdk/client/table/table.h | 2 ++ src/client/table/table.cpp | 15 ++++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 382974aad8..82ae8c7291 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -640,6 +640,7 @@ class TTableDescription { // common void AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns); void AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns); + void AddSecondaryIndex(const TIndexDescription& indexDescription); // sync void AddSyncSecondaryIndex(const std::string& indexName, const std::vector& indexColumns); void AddSyncSecondaryIndex(const std::string& indexName, const std::vector& indexColumns, const std::vector& dataColumns); @@ -853,6 +854,7 @@ class TTableBuilder { TTableBuilder& SetPrimaryKeyColumn(const std::string& primaryKeyColumn); // common + TTableBuilder& AddSecondaryIndex(const TIndexDescription& indexDescription); TTableBuilder& AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns); TTableBuilder& AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns); TTableBuilder& AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::string& indexColumn); diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 01119cba90..67df1cf1cc 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -469,6 +469,10 @@ class TTableDescription::TImpl { Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, dataColumns)); } + void AddSecondaryIndex(const TIndexDescription& indexDescription) { + Indexes_.emplace_back(indexDescription); + } + void AddVectorIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const TVectorIndexSettings& vectorIndexSettings) { Indexes_.emplace_back(TIndexDescription(indexName, type, indexColumns, {}, {}, vectorIndexSettings)); } @@ -750,6 +754,10 @@ void TTableDescription::AddSecondaryIndex(const std::string& indexName, EIndexTy Impl_->AddSecondaryIndex(indexName, type, indexColumns, dataColumns); } +void TTableDescription::AddSecondaryIndex(const TIndexDescription& indexDescription) { + Impl_->AddSecondaryIndex(indexDescription); +} + void TTableDescription::AddSyncSecondaryIndex(const std::string& indexName, const std::vector& indexColumns) { AddSecondaryIndex(indexName, EIndexType::GlobalSync, indexColumns); } @@ -1174,6 +1182,11 @@ TTableBuilder& TTableBuilder::SetPrimaryKeyColumn(const std::string& primaryKeyC return *this; } +TTableBuilder& TTableBuilder::AddSecondaryIndex(const TIndexDescription& indexDescription) { + TableDescription_.AddSecondaryIndex(indexDescription); + return *this; +} + TTableBuilder& TTableBuilder::AddSecondaryIndex(const std::string& indexName, EIndexType type, const std::vector& indexColumns, const std::vector& dataColumns) { TableDescription_.AddSecondaryIndex(indexName, type, indexColumns, dataColumns); return *this; @@ -2377,7 +2390,7 @@ TVectorIndexSettings TVectorIndexSettings::FromProto(const TProto& proto) { default: return {}; } - }; + }; return { .Metric = metricFromProto(proto), From 87bd3a3a5a2544d4f87fe5f788d337a1e71bb4de Mon Sep 17 00:00:00 2001 From: Alek5andr-Kotov Date: Mon, 26 Aug 2024 18:35:35 +0000 Subject: [PATCH 35/35] Moved commit "The value of the WriteInflightSize in the main partition" from ydb repo --- src/client/topic/ut/topic_to_table_ut.cpp | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/client/topic/ut/topic_to_table_ut.cpp b/src/client/topic/ut/topic_to_table_ut.cpp index 2196fbc25e..3373f965dc 100644 --- a/src/client/topic/ut/topic_to_table_ut.cpp +++ b/src/client/topic/ut/topic_to_table_ut.cpp @@ -1898,6 +1898,29 @@ Y_UNIT_TEST_F(WriteToTopic_Demo_27, TFixture) } } +Y_UNIT_TEST_F(WriteToTopic_Demo_28, TFixture) +{ + // The test verifies that the `WriteInflightSize` is correctly considered for the main partition. + // Writing to the service partition does not change the `WriteInflightSize` of the main one. + CreateTopic("topic_A", TEST_CONSUMER); + + NTable::TSession tableSession = CreateTableSession(); + NTable::TTransaction tx = BeginTx(tableSession); + + TString message(16'000, 'a'); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, TString(16'000, 'a'), &tx, 0); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_1); + + CommitTx(tx, EStatus::SUCCESS); + + WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, TString(20'000, 'b'), nullptr, 0); + WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2); + + auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), nullptr, 0); + UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2); +} + } }