From e35a85ce042375804f0a5ea3b7ea8f8f4ae8714e Mon Sep 17 00:00:00 2001
From: Aleksei Kaigorodov <147330454+a1kaigorodov@users.noreply.github.com>
Date: Sat, 8 Feb 2025 15:45:35 +0700
Subject: [PATCH] Support preferable location for connections (#121)

---
 .../ydb/yoj/repository/ydb/YdbConfig.java     | 42 +++++++++++++++++++
 .../ydb/yoj/repository/ydb/YdbRepository.java | 11 +++++
 .../yoj/repository/ydb/yql/YqlOrderBy.java    |  2 -
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java
index d05bc378..132bbf6f 100644
--- a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java
+++ b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java
@@ -1,6 +1,8 @@
 package tech.ydb.yoj.repository.ydb;
 
 import com.google.common.net.HostAndPort;
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.NonNull;
 import lombok.Value;
@@ -26,6 +28,7 @@ public static YdbConfig createForTesting(String host, int port, String tablespac
                 database,
                 null,
                 HostAndPort.fromParts(host, port),
+                null,
                 SESSION_CREATE_TIMEOUT_DEFAULT,
                 SESSION_CREATE_RETRY_COUNT_DEFAULT,
                 SESSION_KEEP_ALIVE_TIME_DEFAULT,
@@ -61,6 +64,9 @@ public static YdbConfig createForTesting(String host, int port, String tablespac
     @With
     HostAndPort hostAndPort;
 
+    @With
+    BalancingConfig balancingConfig;
+
     @With
     Duration sessionCreationTimeout;
     @With
@@ -136,4 +142,40 @@ public Duration getTcpKeepaliveTimeout() {
     public boolean isUseSingleChannelTransport() {
         return Optional.ofNullable(useSingleChannelTransport).orElse(false);
     }
+
+    @Value
+    @AllArgsConstructor(access = AccessLevel.PRIVATE)
+    public static class BalancingConfig {
+        Policy policy;
+        String preferableLocation;
+
+        /**
+         * Use all available cluster nodes regardless datacenter locality
+         */
+        static BalancingConfig useAllNodes() {
+            return new BalancingConfig(Policy.USE_ALL_NODES, null);
+        }
+
+        /**
+         * Use preferable location (data center)
+         *
+         * @param preferableLocation a name of location
+         */
+        static BalancingConfig usePreferableLocation(@NonNull String preferableLocation) {
+            return new BalancingConfig(Policy.USE_PREFERABLE_LOCATION, preferableLocation);
+        }
+
+        /**
+         * Detecting of local DC by the latency measuring
+         */
+        static BalancingConfig detectLocalDc() {
+            return new BalancingConfig(Policy.DETECT_LOCAL_DC, null);
+        }
+
+        public enum Policy {
+            USE_ALL_NODES,
+            USE_PREFERABLE_LOCATION,
+            DETECT_LOCAL_DC,
+        }
+    }
 }
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
index cd9a59f9..58a0e633 100644
--- a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
+++ b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
@@ -10,6 +10,7 @@
 import org.slf4j.LoggerFactory;
 import tech.ydb.auth.AuthProvider;
 import tech.ydb.auth.NopAuthProvider;
+import tech.ydb.core.grpc.BalancingSettings;
 import tech.ydb.core.grpc.GrpcTransport;
 import tech.ydb.core.grpc.GrpcTransportBuilder;
 import tech.ydb.core.impl.SingleChannelTransport;
@@ -117,6 +118,16 @@ private static GrpcTransportBuilder makeGrpcTransportBuilder(@NonNull YdbConfig
             throw new IllegalArgumentException("one of [discoveryEndpoint, hostAndPort] must be set");
         }
 
+        if (config.getBalancingConfig() != null) {
+            transportBuilder.withBalancingSettings(
+                    switch (config.getBalancingConfig().getPolicy()) {
+                        case USE_ALL_NODES -> BalancingSettings.fromPolicy(BalancingSettings.Policy.USE_ALL_NODES);
+                        case USE_PREFERABLE_LOCATION -> BalancingSettings.fromLocation(config.getBalancingConfig().getPreferableLocation());
+                        case DETECT_LOCAL_DC -> BalancingSettings.detectLocalDs();
+                    }
+            );
+        }
+
         if (config.isUseTLS()) {
             if (config.isUseTrustStore()) {
                 transportBuilder.withSecureConnection();
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/yql/YqlOrderBy.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/yql/YqlOrderBy.java
index 9844973b..91141a24 100644
--- a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/yql/YqlOrderBy.java
+++ b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/yql/YqlOrderBy.java
@@ -17,9 +17,7 @@
 
 import static java.lang.String.format;
 import static java.util.Arrays.stream;
-import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.joining;
-import static java.util.stream.Collectors.toList;
 import static java.util.stream.Stream.concat;
 import static tech.ydb.yoj.repository.ydb.yql.YqlOrderBy.SortOrder.ASC;
 import static tech.ydb.yoj.repository.ydb.yql.YqlOrderBy.SortOrder.DESC;