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();