ignite-28305 backpressure to limit in flight partition operations per node

Author: Egor Kuts

modules/client-handler/src/integrationTest/java/org/apache/ignite/client/handler/TestServer.java

@@ -47,6 +47,7 @@
 import org.apache.ignite.internal.sql.engine.QueryProcessor;
 import org.apache.ignite.internal.table.IgniteTablesInternal;
 import org.apache.ignite.internal.tx.TxManager;
+import org.apache.ignite.internal.util.PartitionOperationInFlightLimiter;
 import org.apache.ignite.network.NetworkAddress;
 import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.TestInfo;
@@ -144,6 +145,7 @@ ClientHandlerModule start(TestInfo testInfo) {
                 EventLog.NOOP,
                 new TestLowWatermark(),
                 Runnable::run,
+                new PartitionOperationInFlightLimiter(0),
                 () -> true
         );
 

modules/client-handler/src/main/java/org/apache/ignite/client/handler/ClientHandlerModule.java

@@ -74,6 +74,7 @@
 import org.apache.ignite.internal.table.IgniteTablesInternal;
 import org.apache.ignite.internal.tx.TxManager;
 import org.apache.ignite.internal.util.IgniteSpinBusyLock;
+import org.apache.ignite.internal.util.PartitionOperationInFlightLimiter;
 import org.apache.ignite.lang.IgniteException;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -163,6 +164,8 @@ public class ClientHandlerModule implements IgniteComponent, PlatformComputeTran
 
     private final Executor partitionOperationsExecutor;
 
+    private final PartitionOperationInFlightLimiter partitionOperationInFlightLimiter;
+
     private final ConcurrentHashMap<String, CompletableFuture<PlatformComputeConnection>> computeExecutors = new ConcurrentHashMap<>();
 
     @TestOnly
@@ -186,6 +189,7 @@ public class ClientHandlerModule implements IgniteComponent, PlatformComputeTran
      * @param eventLog Event log.
      * @param lowWatermark Low watermark.
      * @param partitionOperationsExecutor Executor for a partition operation.
+     * @param partitionOperationInFlightLimiter In-flight limiter for partition operations.
      * @param ddlBatchingSuggestionEnabled Boolean supplier indicates whether the suggestion related DDL batching is enabled.
      */
     public ClientHandlerModule(
@@ -207,6 +211,7 @@ public ClientHandlerModule(
             EventLog eventLog,
             LowWatermark lowWatermark,
             Executor partitionOperationsExecutor,
+            PartitionOperationInFlightLimiter partitionOperationInFlightLimiter,
             Supplier<Boolean> ddlBatchingSuggestionEnabled
     ) {
         assert igniteTables != null;
@@ -252,6 +257,7 @@ public ClientHandlerModule(
         this.clientConnectorConfiguration = clientConnectorConfiguration;
         this.ddlBatchingSuggestionEnabled = ddlBatchingSuggestionEnabled;
         this.partitionOperationsExecutor = partitionOperationsExecutor;
+        this.partitionOperationInFlightLimiter = partitionOperationInFlightLimiter;
     }
 
     /** {@inheritDoc} */
@@ -471,6 +477,7 @@ private ClientInboundMessageHandler createInboundMessageHandler(
                 connectionId,
                 primaryReplicaTracker,
                 partitionOperationsExecutor,
+                partitionOperationInFlightLimiter,
                 SUPPORTED_FEATURES,
                 Map.of(),
                 computeExecutors::remove,

modules/client-handler/src/main/java/org/apache/ignite/client/handler/ClientInboundMessageHandler.java

@@ -153,6 +153,7 @@
 import org.apache.ignite.internal.jdbc.proto.JdbcQueryCursorHandler;
 import org.apache.ignite.internal.lang.IgniteExceptionMapperUtil;
 import org.apache.ignite.internal.lang.IgniteInternalCheckedException;
+import org.apache.ignite.internal.lang.ReplicaOverloadedException;
 import org.apache.ignite.internal.logger.IgniteLogger;
 import org.apache.ignite.internal.logger.Loggers;
 import org.apache.ignite.internal.network.ClusterService;
@@ -180,6 +181,7 @@
 import org.apache.ignite.internal.tx.TransactionKilledException;
 import org.apache.ignite.internal.tx.TxManager;
 import org.apache.ignite.internal.util.ExceptionUtils;
+import org.apache.ignite.internal.util.PartitionOperationInFlightLimiter;
 import org.apache.ignite.lang.CancelHandle;
 import org.apache.ignite.lang.ErrorGroups.Compute;
 import org.apache.ignite.lang.ErrorGroups.Sql;
@@ -278,6 +280,8 @@ public class ClientInboundMessageHandler
 
     private final Executor partitionOperationsExecutor;
 
+    private final PartitionOperationInFlightLimiter partitionOperationInFlightLimiter;
+
     private final BitSet features;
 
     private final Map<HandshakeExtension, Object> extensions;
@@ -309,6 +313,7 @@ public class ClientInboundMessageHandler
      * @param connectionId Connection ID.
      * @param primaryReplicaTracker Primary replica tracker.
      * @param partitionOperationsExecutor Partition operations executor.
+     * @param partitionOperationInFlightLimiter In-flight limiter for partition operations.
      * @param features Features.
      * @param extensions Extensions.
      * @param eventLog Event log.
@@ -330,6 +335,7 @@ public ClientInboundMessageHandler(
             long connectionId,
             ClientPrimaryReplicaTracker primaryReplicaTracker,
             Executor partitionOperationsExecutor,
+            PartitionOperationInFlightLimiter partitionOperationInFlightLimiter,
             BitSet features,
             Map<HandshakeExtension, Object> extensions,
             Function<String, CompletableFuture<PlatformComputeConnection>> computeConnectionFunc,
@@ -373,6 +379,7 @@ public ClientInboundMessageHandler(
         this.eventLog = eventLog;
         this.primaryReplicaTracker = primaryReplicaTracker;
         this.partitionOperationsExecutor = partitionOperationsExecutor;
+        this.partitionOperationInFlightLimiter = partitionOperationInFlightLimiter;
         this.handshakeEventLoopSwitcher = handshakeEventLoopSwitcher;
 
         jdbcQueryCursorHandler = new JdbcQueryCursorHandlerImpl(resources);
@@ -882,19 +889,28 @@ private void processOperation(ChannelHandlerContext ctx, ClientMessageUnpacker i
             if (ClientOp.isPartitionOperation(opCode)) {
                 long requestId0 = requestId;
                 int opCode0 = opCode;
+                if (!partitionOperationInFlightLimiter.tryAcquire()) {
+                    in.close();
 
-                partitionOperationsExecutor.execute(() -> {
-                    try {
-                        processOperationInternal(ctx, in, requestId0, opCode0, guard);
-                    } catch (Throwable t) {
-                        in.close();
-
-                        writeError(requestId0, opCode0, t, ctx, false, guard);
+                    writeError(requestId0, opCode0, new ReplicaOverloadedException(), ctx, false, guard);
 
-                        metrics.requestsFailedIncrement();
-                        metrics.requestsActiveDecrement();
-                    }
-                });
+                    metrics.requestsFailedIncrement();
+                } else {
+                    partitionOperationsExecutor.execute(() -> {
+                        try {
+                            processOperationInternal(ctx, in, requestId0, opCode0, guard);
+                        } catch (Throwable t) {
+                            in.close();
+
+                            writeError(requestId0, opCode0, t, ctx, false, guard);
+
+                            metrics.requestsFailedIncrement();
+                            metrics.requestsActiveDecrement();
+                        } finally {
+                            partitionOperationInFlightLimiter.release();
+                        }
+                    });
+                }
             } else {
                 processOperationInternal(ctx, in, requestId, opCode, guard);
             }

modules/client/src/test/java/org/apache/ignite/client/TestClientHandlerModule.java

@@ -65,6 +65,7 @@
 import org.apache.ignite.internal.schema.AlwaysSyncedSchemaSyncService;
 import org.apache.ignite.internal.security.authentication.AuthenticationManager;
 import org.apache.ignite.internal.table.IgniteTablesInternal;
+import org.apache.ignite.internal.util.PartitionOperationInFlightLimiter;
 import org.apache.ignite.lang.IgniteException;
 import org.jetbrains.annotations.Nullable;
 
@@ -272,6 +273,7 @@ protected void initChannel(Channel ch) {
                                                 new TestLowWatermark()
                                         ),
                                         Runnable::run,
+                                        new PartitionOperationInFlightLimiter(0),
                                         features,
                                         randomExtensions(),
                                         unused -> null,

modules/client/src/test/java/org/apache/ignite/client/TestServer.java

@@ -79,6 +79,7 @@
 import org.apache.ignite.internal.security.authentication.AuthenticationManagerImpl;
 import org.apache.ignite.internal.security.configuration.SecurityConfiguration;
 import org.apache.ignite.internal.table.IgniteTablesInternal;
+import org.apache.ignite.internal.util.PartitionOperationInFlightLimiter;
 import org.apache.ignite.network.NetworkAddress;
 import org.jetbrains.annotations.Nullable;
 import org.mockito.Mockito;
@@ -290,6 +291,7 @@ public void log(String type, Supplier<Event> eventProvider) {
                         EventLog.NOOP,
                         new TestLowWatermark(),
                         Runnable::run,
+                        new PartitionOperationInFlightLimiter(0),
                         () -> true
                 );
 

modules/core/src/main/java/org/apache/ignite/internal/lang/ReplicaOverloadedException.java

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.lang;
+
+import static org.apache.ignite.lang.ErrorGroups.Replicator.GROUP_OVERLOADED_ERR;
+
+/**
+ * Thrown when the node has reached the maximum number of in-flight partition operations
+ * ({@code replication.maxInFlightPartitionOperationsPerCore}) and cannot accept new requests.
+ */
+public class ReplicaOverloadedException extends IgniteInternalException {
+    private static final long serialVersionUID = -6023736883539658779L;
+
+    /** Constructor. */
+    public ReplicaOverloadedException() {
+        super(GROUP_OVERLOADED_ERR, "Node is overloaded: max in-flight partition operations limit reached.");
+    }
+}

modules/core/src/main/java/org/apache/ignite/internal/util/PartitionOperationInFlightLimiter.java

@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.util;
+
+import java.util.concurrent.Semaphore;
+import java.util.function.IntSupplier;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Limits the number of in-flight partition operations (queued or executing) across the replica manager and thin-client connector.
+ *
+ * <p>When the limit is zero or less, all operations are permitted unconditionally.
+ * When positive, {@link #tryAcquire()} returns {@code false} once the limit is reached and the caller should reject the request.
+ * A permit must be released via {@link #release()} upon operation completes.
+ */
+public class PartitionOperationInFlightLimiter {
+    private volatile Semaphore semaphore;
+
+    private final @Nullable IntSupplier limitSupplier;
+
+    private volatile boolean initialized;
+
+    /**
+     * Constructor.
+     *
+     * @param maxInFlightPartitionOperationsPerCore Max number of in-flight partition operations per CPU core, or <= 0 to disable the limit.
+     *     The total limit is {@code maxInFlightPartitionOperationsPerCore * availableProcessors}.
+     */
+    public PartitionOperationInFlightLimiter(int maxInFlightPartitionOperationsPerCore) {
+        int limit = maxInFlightPartitionOperationsPerCore <= 0 ? 0
+                : maxInFlightPartitionOperationsPerCore * Runtime.getRuntime().availableProcessors();
+        this.semaphore = limit <= 0 ? null : new Semaphore(limit);
+        this.limitSupplier = null;
+        this.initialized = true;
+    }
+
+    /**
+     * Constructor.
+     *
+     * @param maxInFlightPartitionOperationsPerCoreSupplier Supplier of the max number of in-flight partition operations per CPU core,
+     *     or 0 to disable. The total limit is {@code supplied value * availableProcessors}.
+     */
+    public PartitionOperationInFlightLimiter(@Nullable IntSupplier maxInFlightPartitionOperationsPerCoreSupplier) {
+        this.limitSupplier = maxInFlightPartitionOperationsPerCoreSupplier;
+        this.initialized = false;
+    }
+
+    /**
+     * Attempts to acquire a permit.
+     *
+     * @return {@code true} if a permit was acquired or the limit is disabled; {@code false} if the limit is reached.
+     */
+    public boolean tryAcquire() {
+        Semaphore s = resolvedSemaphore();
+        return s == null || s.tryAcquire();
+    }
+
+    /**
+     * Releases a previously acquired permit.
+     * Must only be called after a successful {@link #tryAcquire()} when the limit is enabled.
+     */
+    public void release() {
+        Semaphore s = resolvedSemaphore();
+
+        if (s != null) {
+            s.release();
+        }
+    }
+
+    private @Nullable Semaphore resolvedSemaphore() {
+        if (initialized) {
+            return semaphore;
+        }
+        synchronized (this) {
+            if (initialized) {
+                return semaphore;
+            }
+            if (limitSupplier != null) {
+                int perCore = limitSupplier.getAsInt();
+                int limit = perCore <= 0 ? 0 : perCore * Runtime.getRuntime().availableProcessors();
+
+                if (limit > 0) {
+                    this.semaphore = new Semaphore(limit);
+                }
+            }
+            this.initialized = true;
+        }
+        return semaphore;
+    }
+}