Merge pull request #535 from alex268/master

alex268 · web-flow · commit 9c7051c9abd8 · 2025-09-02T13:26:41.000+01:00
Added endpoint pessimization when CreateSession returns OVERLOADED
diff --git a/core/src/main/java/tech/ydb/core/grpc/GrpcRequestSettings.java b/core/src/main/java/tech/ydb/core/grpc/GrpcRequestSettings.java
@@ -3,6 +3,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
 import io.grpc.Metadata;
@@ -15,17 +16,21 @@
 public class GrpcRequestSettings {
     private final long deadlineAfter;
     private final Integer preferredNodeID;
+    private final boolean directMode;
     private final String traceId;
     private final List<String> clientCapabilities;
     private final Consumer<Metadata> trailersHandler;
+    private final BooleanSupplier pessimizationHook;
     private final GrpcFlowControl flowControl;
 
     private GrpcRequestSettings(Builder builder) {
         this.deadlineAfter = builder.deadlineAfter;
         this.preferredNodeID = builder.preferredNodeID;
+        this.directMode = builder.directMode;
         this.traceId = builder.traceId;
         this.clientCapabilities = builder.clientCapabilities;
         this.trailersHandler = builder.trailersHandler;
+        this.pessimizationHook = builder.pessimizationHook;
         this.flowControl = builder.flowControl;
     }
 
@@ -41,6 +46,10 @@ public Integer getPreferredNodeID() {
         return preferredNodeID;
     }
 
+    public boolean isDirectMode() {
+        return directMode;
+    }
+
     public String getTraceId() {
         return traceId;
     }
@@ -53,16 +62,22 @@ public Consumer<Metadata> getTrailersHandler() {
         return trailersHandler;
     }
 
+    public BooleanSupplier getPessimizationHook() {
+        return pessimizationHook;
+    }
+
     public GrpcFlowControl getFlowControl() {
         return flowControl;
     }
 
     public static final class Builder {
         private long deadlineAfter = 0L;
         private Integer preferredNodeID = null;
+        private boolean directMode = false;
         private String traceId = null;
         private List<String> clientCapabilities = null;
         private Consumer<Metadata> trailersHandler = null;
+        private BooleanSupplier pessimizationHook = null;
         private GrpcFlowControl flowControl = GrpcFlows.SIMPLE_FLOW;
 
         /**
@@ -109,6 +124,11 @@ public Builder withFlowControl(GrpcFlowControl flowCtrl) {
             return this;
         }
 
+        public Builder withDirectMode(boolean directMode) {
+            this.directMode = directMode;
+            return this;
+        }
+
         public Builder withClientCapabilities(List<String> clientCapabilities) {
             this.clientCapabilities = clientCapabilities;
             return this;
@@ -127,6 +147,11 @@ public Builder withTrailersHandler(Consumer<Metadata> handler) {
             return this;
         }
 
+        public Builder withPessimizationHook(BooleanSupplier pessimizationHook) {
+            this.pessimizationHook = pessimizationHook;
+            return this;
+        }
+
         public GrpcRequestSettings build() {
             return new GrpcRequestSettings(this);
         }
diff --git a/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java b/core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java
@@ -3,6 +3,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BooleanSupplier;
 
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
@@ -29,6 +30,7 @@
 import tech.ydb.core.impl.call.ReadStreamCall;
 import tech.ydb.core.impl.call.ReadWriteStreamCall;
 import tech.ydb.core.impl.call.UnaryCall;
+import tech.ydb.core.impl.pool.EndpointRecord;
 import tech.ydb.core.impl.pool.GrpcChannel;
 
 /**
@@ -47,7 +49,10 @@ public abstract class BaseGrpcTransport implements GrpcTransport {
 
     protected abstract AuthCallOptions getAuthCallOptions();
     protected abstract GrpcChannel getChannel(GrpcRequestSettings settings);
-    protected abstract void updateChannelStatus(GrpcChannel channel, io.grpc.Status status);
+
+    protected void pessimizeEndpoint(EndpointRecord endpoint, String reason) {
+        // nothing to pessimize
+    }
 
     protected void shutdown() {
         // nothing to shutdown
@@ -226,10 +231,22 @@ private class ChannelStatusHandler implements GrpcStatusHandler {
 
         @Override
         public void accept(io.grpc.Status status, Metadata trailers) {
-            updateChannelStatus(channel, status);
+            // Usually CANCELLED is received when ClientCall is canceled on client side
+            if (!status.isOk() && status.getCode() != io.grpc.Status.Code.CANCELLED) {
+                pessimizeEndpoint(channel.getEndpoint(), "by grpc code " + status.getCode());
+            }
+
             if (settings.getTrailersHandler() != null && trailers != null) {
                 settings.getTrailersHandler().accept(trailers);
             }
         }
+
+        @Override
+        public void postComplete() {
+            BooleanSupplier hook = settings.getPessimizationHook();
+            if (hook != null && hook.getAsBoolean()) {
+                pessimizeEndpoint(channel.getEndpoint(), "by pessimization hook");
+            }
+        }
     }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/FixedCallOptionsTransport.java b/core/src/main/java/tech/ydb/core/impl/FixedCallOptionsTransport.java
@@ -2,7 +2,6 @@
 
 import java.util.concurrent.ScheduledExecutorService;
 
-import io.grpc.Status;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -60,15 +59,4 @@ public AuthCallOptions getAuthCallOptions() {
     protected GrpcChannel getChannel(GrpcRequestSettings settings) {
         return channel;
     }
-
-    @Override
-    protected void updateChannelStatus(GrpcChannel channel, Status status) {
-        if (!status.isOk()) {
-            logger.warn("grpc error {}[{}] on fixed channel {}",
-                    status.getCode(),
-                    status.getDescription(),
-                    channel.getEndpoint());
-        }
-    }
-
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/MultiChannelTransport.java b/core/src/main/java/tech/ydb/core/impl/MultiChannelTransport.java
@@ -7,7 +7,6 @@
 
 import com.google.common.base.Strings;
 import com.google.common.net.HostAndPort;
-import io.grpc.Status;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -83,13 +82,11 @@ protected GrpcChannel getChannel(GrpcRequestSettings settings) {
     }
 
     @Override
-    protected void updateChannelStatus(GrpcChannel channel, Status status) {
-        if (!status.isOk()) {
-            endpointPool.pessimizeEndpoint(channel.getEndpoint());
+    protected void pessimizeEndpoint(EndpointRecord endpoint, String reason) {
+        endpointPool.pessimizeEndpoint(endpoint, reason);
 
-            if (endpointPool.needToRunDiscovery()) {
-                endpointPool.setNewState(null, endpoints);
-            }
+        if (endpointPool.needToRunDiscovery()) {
+            endpointPool.setNewState(null, endpoints);
         }
     }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/SingleChannelTransport.java b/core/src/main/java/tech/ydb/core/impl/SingleChannelTransport.java
@@ -4,7 +4,6 @@
 import java.util.concurrent.ScheduledExecutorService;
 
 import com.google.common.base.Strings;
-import io.grpc.Status;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -65,14 +64,4 @@ public AuthCallOptions getAuthCallOptions() {
     protected GrpcChannel getChannel(GrpcRequestSettings settings) {
         return channel;
     }
-
-    @Override
-    protected void updateChannelStatus(GrpcChannel channel, Status status) {
-        if (!status.isOk()) {
-            logger.warn("grpc error {}[{}] on single channel {}",
-                    status.getCode(),
-                    status.getDescription(),
-                    channel.getEndpoint());
-        }
-    }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/YdbTransportImpl.java b/core/src/main/java/tech/ydb/core/impl/YdbTransportImpl.java
@@ -160,24 +160,21 @@ public AuthCallOptions getAuthCallOptions() {
 
     @Override
     protected GrpcChannel getChannel(GrpcRequestSettings settings) {
-        EndpointRecord endpoint = endpointPool.getEndpoint(settings.getPreferredNodeID());
+        EndpointRecord endpoint = endpointPool.getEndpoint(settings);
         if (endpoint == null) {
             long timeout = -1;
             if (settings.getDeadlineAfter() != 0) {
                 timeout = settings.getDeadlineAfter() - System.nanoTime();
             }
             discovery.waitReady(timeout);
-            endpoint = endpointPool.getEndpoint(settings.getPreferredNodeID());
+            endpoint = endpointPool.getEndpoint(settings);
         }
         return channelPool.getChannel(endpoint);
     }
 
     @Override
-    protected void updateChannelStatus(GrpcChannel channel, io.grpc.Status status) {
-        // Usually CANCELLED is received when ClientCall is canceled on client side
-        if (!status.isOk() && status.getCode() != io.grpc.Status.Code.CANCELLED) {
-            endpointPool.pessimizeEndpoint(channel.getEndpoint());
-        }
+    protected void pessimizeEndpoint(EndpointRecord endpoint, String reason) {
+        endpointPool.pessimizeEndpoint(endpoint, reason);
     }
 
     private class DiscoveryHandler implements YdbDiscovery.Handler {
diff --git a/core/src/main/java/tech/ydb/core/impl/call/GrpcStatusHandler.java b/core/src/main/java/tech/ydb/core/impl/call/GrpcStatusHandler.java
@@ -9,4 +9,6 @@
  */
 public interface GrpcStatusHandler {
     void accept(Status status, Metadata trailers);
+
+    void postComplete();
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/call/ReadStreamCall.java b/core/src/main/java/tech/ydb/core/impl/call/ReadStreamCall.java
@@ -144,5 +144,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {
         } else {
             statusFuture.complete(GrpcStatuses.toStatus(status));
         }
+
+        statusConsumer.postComplete();
     }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/call/ReadWriteStreamCall.java b/core/src/main/java/tech/ydb/core/impl/call/ReadWriteStreamCall.java
@@ -206,5 +206,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {
         } else {
             statusFuture.complete(GrpcStatuses.toStatus(status));
         }
+
+        statusConsumer.postComplete();
     }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/call/UnaryCall.java b/core/src/main/java/tech/ydb/core/impl/call/UnaryCall.java
@@ -105,5 +105,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {
         } else {
             future.complete(GrpcStatuses.toResult(status));
         }
+
+        statusConsumer.postComplete();
     }
 }
diff --git a/core/src/main/java/tech/ydb/core/impl/pool/EndpointPool.java b/core/src/main/java/tech/ydb/core/impl/pool/EndpointPool.java
@@ -15,14 +15,20 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import tech.ydb.core.Status;
+import tech.ydb.core.StatusCode;
+import tech.ydb.core.UnexpectedResultException;
 import tech.ydb.core.grpc.BalancingSettings;
+import tech.ydb.core.grpc.GrpcRequestSettings;
 
 /**
  * @author Nikolay Perfilov
  * @author Kirill Kurdyukov
  */
 public final class EndpointPool {
     private static final Logger logger = LoggerFactory.getLogger(EndpointPool.class);
+    private static final Status DIRECT_REQUEST_ERROR_CODE = Status.of(StatusCode.CLIENT_INTERNAL_ERROR);
+    private static final Status NODE_NOT_FOUND_ERROR_CODE = Status.of(StatusCode.TRANSPORT_UNAVAILABLE);
 
     // Maximum percent of endpoints pessimized by transport errors to start recheck
     private static final long DISCOVERY_PESSIMIZATION_THRESHOLD = 50;
@@ -44,15 +50,25 @@ public EndpointPool(BalancingSettings balancingSettings) {
     }
 
     @Nullable
-    public EndpointRecord getEndpoint(@Nullable Integer preferredNodeID) {
+    public EndpointRecord getEndpoint(GrpcRequestSettings settings) {
         recordsLock.readLock().lock();
         try {
-            if (preferredNodeID != null) {
-                PriorityEndpoint knownEndpoint = recordsByNodeId.get(preferredNodeID);
+            Integer nodeId = settings.getPreferredNodeID();
+            boolean directMode = settings.isDirectMode();
+
+            if (nodeId != null) {
+                PriorityEndpoint knownEndpoint = recordsByNodeId.get(nodeId);
                 if (knownEndpoint != null) {
                     return knownEndpoint.record;
                 }
+                if (directMode) {
+                    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);
@@ -124,7 +140,7 @@ public List<EndpointRecord> setNewState(String selfLocation, List<EndpointRecord
         return removed;
     }
 
-    public void pessimizeEndpoint(EndpointRecord endpoint) {
+    public void pessimizeEndpoint(EndpointRecord endpoint, String reason) {
         if (endpoint == null) {
             return;
         }
@@ -146,26 +162,26 @@ public void pessimizeEndpoint(EndpointRecord endpoint) {
             records.sort(PriorityEndpoint.COMPARATOR);
 
             long bestPriority = records.get(0).priority;
-            int newBestCount = 0;
-            int newPessimizedCount = 0;
+            int bestCount = 0;
+            int pcount = 0;
             for (PriorityEndpoint record : records) {
                 if (record.getPriority() == bestPriority) {
-                    newBestCount++;
+                    bestCount++;
                 }
                 if (record.isPessimized()) {
-                    newPessimizedCount++;
+                    pcount++;
                 }
             }
 
-            bestEndpointsCount = newBestCount;
-            needToRunDiscovery = 100 * newPessimizedCount > records.size() * DISCOVERY_PESSIMIZATION_THRESHOLD;
+            bestEndpointsCount = bestCount;
+            needToRunDiscovery = 100 * pcount > records.size() * DISCOVERY_PESSIMIZATION_THRESHOLD;
             if (needToRunDiscovery) {
                 logger.debug("launching discovery due to pessimization threshold is exceeded: {}/{} is more than {}",
-                        newPessimizedCount, records.size(), DISCOVERY_PESSIMIZATION_THRESHOLD);
+                        pcount, records.size(), DISCOVERY_PESSIMIZATION_THRESHOLD);
             }
 
-            logger.trace("Endpoint {} was pessimized. New pessimization ratio: {}/{}",
-                    endpoint, newPessimizedCount, records.size());
+            logger.warn("Endpoint {} was pessimized {}. New pessimization ratio: {}/{}",
+                    endpoint, reason, pcount, records.size());
         } finally {
             recordsLock.writeLock().unlock();
         }
diff --git a/core/src/test/java/tech/ydb/core/impl/pool/EndpointPoolTest.java b/core/src/test/java/tech/ydb/core/impl/pool/EndpointPoolTest.java
diff --git a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java
diff --git a/table/src/main/java/tech/ydb/table/impl/BaseSession.java b/table/src/main/java/tech/ydb/table/impl/BaseSession.java

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,6 @@`
`7`	`7`
`8`	`8`	`import com.google.common.base.Strings;`
`9`	`9`	`import com.google.common.net.HostAndPort;`
`10`		`-import io.grpc.Status;`
`11`	`10`	`import org.slf4j.Logger;`
`12`	`11`	`import org.slf4j.LoggerFactory;`
`13`	`12`
`@@ -83,13 +82,11 @@ protected GrpcChannel getChannel(GrpcRequestSettings settings) {`
`83`	`82`	`}`
`84`	`83`
`85`	`84`	`@Override`
`86`		`- protected void updateChannelStatus(GrpcChannel channel, Status status) {`
`87`		`- if (!status.isOk()) {`
`88`		`- endpointPool.pessimizeEndpoint(channel.getEndpoint());`
	`85`	`+ protected void pessimizeEndpoint(EndpointRecord endpoint, String reason) {`
	`86`	`+ endpointPool.pessimizeEndpoint(endpoint, reason);`
`89`	`87`
`90`		`- if (endpointPool.needToRunDiscovery()) {`
`91`		`- endpointPool.setNewState(null, endpoints);`
`92`		`- }`
	`88`	`+ if (endpointPool.needToRunDiscovery()) {`
	`89`	`+ endpointPool.setNewState(null, endpoints);`
`93`	`90`	`}`
`94`	`91`	`}`
`95`	`92`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,4 +9,6 @@`
`9`	`9`	`*/`
`10`	`10`	`public interface GrpcStatusHandler {`
`11`	`11`	`void accept(Status status, Metadata trailers);`
	`12`	`+`
	`13`	`+ void postComplete();`
`12`	`14`	`}`
Original file line number	Diff line number	Diff line change
`@@ -144,5 +144,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {`
`144`	`144`	`} else {`
`145`	`145`	`statusFuture.complete(GrpcStatuses.toStatus(status));`
`146`	`146`	`}`
	`147`	`+`
	`148`	`+ statusConsumer.postComplete();`
`147`	`149`	`}`
`148`	`150`	`}`
Original file line number	Diff line number	Diff line change
`@@ -206,5 +206,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {`
`206`	`206`	`} else {`
`207`	`207`	`statusFuture.complete(GrpcStatuses.toStatus(status));`
`208`	`208`	`}`
	`209`	`+`
	`210`	`+ statusConsumer.postComplete();`
`209`	`211`	`}`
`210`	`212`	`}`
Original file line number	Diff line number	Diff line change
`@@ -105,5 +105,7 @@ public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {`
`105`	`105`	`} else {`
`106`	`106`	`future.complete(GrpcStatuses.toResult(status));`
`107`	`107`	`}`
	`108`	`+`
	`109`	`+ statusConsumer.postComplete();`
`108`	`110`	`}`
`109`	`111`	`}`