diff options
author | azevaykin <azevaykin@yandex-team.com> | 2023-08-02 09:01:52 +0300 |
---|---|---|
committer | azevaykin <azevaykin@yandex-team.com> | 2023-08-02 09:01:52 +0300 |
commit | 0b51f759add67b66ecbc8a585b99f19b5856f56e (patch) | |
tree | 52c868041724276edb2636a35aeea53406a58321 | |
parent | fa5486d422e6eb00eb59c8c8046a9e64ef858177 (diff) | |
download | ydb-0b51f759add67b66ecbc8a585b99f19b5856f56e.tar.gz |
GetEndpoint onlyPreferred
5 files changed, 21 insertions, 6 deletions
diff --git a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.cpp b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.cpp index f9fdfd5153..6db78f6ee7 100644 --- a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.cpp +++ b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.cpp @@ -128,7 +128,7 @@ std::vector<TStringType> TEndpointElectorSafe::SetNewState(std::vector<TEndpoint return removed; } -TEndpointRecord TEndpointElectorSafe::GetEndpoint(const TEndpointKey& preferredEndpoint) const { +TEndpointRecord TEndpointElectorSafe::GetEndpoint(const TEndpointKey& preferredEndpoint, bool onlyPreferred) const { std::shared_lock guard(Mutex_); if (preferredEndpoint.GetNodeId()) { @@ -144,9 +144,13 @@ TEndpointRecord TEndpointElectorSafe::GetEndpoint(const TEndpointKey& preferredE return it->second; } } + + if(onlyPreferred) + return {}; + if (BestK_ == -1) { Y_ASSERT(Records_.empty()); - return TEndpointRecord(); + return {}; } else { // returns value in range [0, n) auto idx = RandomNumber<size_t>(BestK_ + 1); diff --git a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.h b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.h index f00dabc744..70c879fcc1 100644 --- a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.h +++ b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints.h @@ -82,7 +82,7 @@ public: void SetStatCollector(const NSdkStats::TStatCollector::TEndpointElectorStatCollector& endpointStatCollector); // Returns preferred (if presents) or best endpoint - TEndpointRecord GetEndpoint(const TEndpointKey& preferredEndpoint) const; + TEndpointRecord GetEndpoint(const TEndpointKey& preferredEndpoint, bool onlyPreferred = false) const; // Move endpoint to the end void PessimizeEndpoint(const TStringType& endpoint); diff --git a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints_ut.cpp b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints_ut.cpp index 4c7f3265bb..20f39debe3 100644 --- a/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints_ut.cpp +++ b/ydb/public/sdk/cpp/client/impl/ydb_endpoints/endpoints_ut.cpp @@ -114,6 +114,17 @@ Y_UNIT_TEST_SUITE(EndpointElector) { UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey()).Endpoint, ""); } + Y_UNIT_TEST(GetEndpoint) { + TEndpointElectorSafe elector; + elector.SetNewState(TVector<TEndpointRecord>{{"Two", 0, "", 2}, {"One", 0, "", 1}}); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("One", 0), true).Endpoint, "One"); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("Two", 0), true).Endpoint, "Two"); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("", 1), true).NodeId, 1); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("", 2), true).NodeId, 2); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("Three", 0), true).Endpoint, ""); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("", 3), true).NodeId, 0); + } + Y_UNIT_TEST(DiffOnRemove) { TEndpointElectorSafe elector; auto removed = elector.SetNewState(TVector<TEndpointRecord>{{"Two", 2}, {"One", 1}}); diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.cpp b/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.cpp index a427c3d279..e8b6ad328b 100644 --- a/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.cpp +++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.cpp @@ -122,8 +122,8 @@ std::pair<NThreading::TFuture<TEndpointUpdateResult>, bool> TEndpointPool::Updat return {future, true}; } -TEndpointRecord TEndpointPool::GetEndpoint(const TEndpointKey& preferredEndpoint) const { - return Elector_.GetEndpoint(preferredEndpoint); +TEndpointRecord TEndpointPool::GetEndpoint(const TEndpointKey& preferredEndpoint, bool onlyPreferred) const { + return Elector_.GetEndpoint(preferredEndpoint, onlyPreferred); } TDuration TEndpointPool::TimeSinceLastUpdate() const { diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.h b/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.h index a5f6954384..bddb0f5bd6 100644 --- a/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.h +++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/db_driver_state/endpoint_pool.h @@ -31,7 +31,7 @@ public: TEndpointPool(TListEndpointsResultProvider&& provider, const IInternalClient* client); ~TEndpointPool(); std::pair<NThreading::TFuture<TEndpointUpdateResult>, bool> UpdateAsync(); - TEndpointRecord GetEndpoint(const TEndpointKey& preferredEndpoint) const; + TEndpointRecord GetEndpoint(const TEndpointKey& preferredEndpoint, bool onlyPreferred = false) const; TDuration TimeSinceLastUpdate() const; void BanEndpoint(const TStringType& endpoint); int GetPessimizationRatio(); |