Added option to prefer ready channels to GrpcRequestSettings

Author: alex268

core/src/main/java/tech/ydb/core/grpc/GrpcRequestSettings.java

@@ -17,6 +17,7 @@ public class GrpcRequestSettings {
     private final long deadlineAfter;
     private final Integer preferredNodeID;
     private final boolean directMode;
+    private final boolean prefferReadyChannel;
     private final boolean deadlineDisabled;
     private final String traceId;
     private final List<String> clientCapabilities;
@@ -28,6 +29,7 @@ private GrpcRequestSettings(Builder builder) {
         this.deadlineAfter = builder.deadlineAfter;
         this.preferredNodeID = builder.preferredNodeID;
         this.directMode = builder.directMode;
+        this.prefferReadyChannel = builder.preferReadyChannel;
         this.deadlineDisabled = builder.deadlineDisabled;
         this.traceId = builder.traceId;
         this.clientCapabilities = builder.clientCapabilities;
@@ -56,6 +58,10 @@ public boolean isDirectMode() {
         return directMode;
     }
 
+    public boolean isPreferReadyChannel() {
+        return prefferReadyChannel;
+    }
+
     public String getTraceId() {
         return traceId;
     }
@@ -78,6 +84,7 @@ public GrpcFlowControl getFlowControl() {
 
     public static final class Builder {
         private long deadlineAfter = 0L;
+        private boolean preferReadyChannel = false;
         private boolean deadlineDisabled = false;
         private Integer preferredNodeID = null;
         private boolean directMode = false;
@@ -138,6 +145,11 @@ public Builder withDirectMode(boolean directMode) {
             return this;
         }
 
+        public Builder withPreferReadyChannel(boolean preferReady) {
+            this.preferReadyChannel = preferReady;
+            return this;
+        }
+
         public Builder withClientCapabilities(List<String> clientCapabilities) {
             this.clientCapabilities = clientCapabilities;
             return this;

core/src/main/java/tech/ydb/core/impl/MultiChannelTransport.java

@@ -1,5 +1,6 @@
 package tech.ydb.core.impl;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.ScheduledExecutorService;
@@ -77,7 +78,7 @@ public AuthCallOptions getAuthCallOptions() {
 
     @Override
     protected GrpcChannel getChannel(GrpcRequestSettings settings) {
-        EndpointRecord endpoint = endpointPool.getEndpoint(null);
+        EndpointRecord endpoint = endpointPool.getEndpoint(Collections.emptySet(), settings);
         return channelPool.getChannel(endpoint);
     }
 

core/src/main/java/tech/ydb/core/impl/YdbTransportImpl.java

@@ -165,14 +165,14 @@ public AuthCallOptions getAuthCallOptions() {
 
     @Override
     protected GrpcChannel getChannel(GrpcRequestSettings settings) {
-        EndpointRecord endpoint = endpointPool.getEndpoint(settings);
+        EndpointRecord endpoint = endpointPool.getEndpoint(channelPool.getReadyEndpoints(), settings);
         if (endpoint == null) {
             long timeout = -1;
             if (settings.getDeadlineAfter() != 0) {
                 timeout = settings.getDeadlineAfter() - System.nanoTime();
             }
             discovery.waitReady(timeout);
-            endpoint = endpointPool.getEndpoint(settings);
+            endpoint = endpointPool.getEndpoint(Collections.emptySet(), settings);
         }
         return channelPool.getChannel(endpoint);
     }

core/src/main/java/tech/ydb/core/impl/pool/EndpointPool.java

@@ -5,9 +5,11 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
 
 import javax.annotation.Nullable;
 
@@ -50,12 +52,13 @@ public EndpointPool(BalancingSettings balancingSettings) {
     }
 
     @Nullable
-    public EndpointRecord getEndpoint(GrpcRequestSettings settings) {
+    public EndpointRecord getEndpoint(Set<String> readyEnpoints, GrpcRequestSettings settings) {
+        Integer nodeId = settings.getPreferredNodeID();
+        boolean directMode = settings.isDirectMode();
+        boolean prefferReady = settings.isPreferReadyChannel();
+
         recordsLock.readLock().lock();
         try {
-            Integer nodeId = settings.getPreferredNodeID();
-            boolean directMode = settings.isDirectMode();
-
             if (nodeId != null) {
                 PriorityEndpoint knownEndpoint = recordsByNodeId.get(nodeId);
                 if (knownEndpoint != null) {
@@ -65,17 +68,32 @@ public EndpointRecord getEndpoint(GrpcRequestSettings settings) {
                     throw new UnexpectedResultException("Node " + nodeId + " not found", NODE_NOT_FOUND_ERROR_CODE);
                 }
             }
+
             if (directMode) {
                 throw new UnexpectedResultException("Cannot use direct mode without NodeId", DIRECT_REQUEST_ERROR_CODE);
             }
 
-            if (bestEndpointsCount > 0) {
-                // returns value in range [0, n)
-                int idx = ThreadLocalRandom.current().nextInt(bestEndpointsCount);
-                return records.get(idx).record;
-            } else {
+            if (bestEndpointsCount <= 0) {
+                // pool is not ready
                 return null;
             }
+
+            if (prefferReady && !readyEnpoints.isEmpty()) {
+                List<PriorityEndpoint> ready = readyEnpoints.stream()
+                        .map(recordsByEndpoint::get)
+                        .filter(pr -> pr != null && !pr.isPessimized())
+                        .collect(Collectors.toList());
+
+                if (!ready.isEmpty()) {
+                    // returns value in range [0, n)
+                    int idx = ThreadLocalRandom.current().nextInt(ready.size());
+                    return ready.get(idx).record;
+                }
+            }
+
+            // returns value in range [0, n)
+            int idx = ThreadLocalRandom.current().nextInt(bestEndpointsCount);
+            return records.get(idx).record;
         } finally {
             recordsLock.readLock().unlock();
         }

core/src/main/java/tech/ydb/core/impl/pool/GrpcChannelPool.java

@@ -4,6 +4,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledExecutorService;
@@ -29,6 +30,10 @@ public GrpcChannelPool(ManagedChannelFactory channelFactory, ScheduledExecutorSe
         this.executor = executor;
     }
 
+    public Set<String> getReadyEndpoints() {
+        return channels.keySet();
+    }
+
     public GrpcChannel getChannel(EndpointRecord endpoint) {
         // Workaround for https://bugs.openjdk.java.net/browse/JDK-8161372 to prevent unnecessary locks in Java 8
         // Was fixed in Java 9+