Fix indexing pressure replica rejections logic (#60150)

Currently the logic to rejection replica rejections is evaluate before
adding the additional bytes of the current operation. This means that
the first replica operation which should be rejected will be allowed to
proceed. This commit fixes this logic and adds unit level test to ensure
indexing pressure behavior is correct.

server/src/main/java/org/elasticsearch/index/IndexingPressure.java

@@ -112,7 +112,7 @@ public Releasable markPrimaryOperationStarted(long bytes, boolean forceExecution
     }
 
     public Releasable markReplicaOperationStarted(long bytes, boolean forceExecution) {
-        long replicaWriteBytes = this.currentReplicaBytes.getAndAdd(bytes);
+        long replicaWriteBytes = this.currentReplicaBytes.addAndGet(bytes);
         if (forceExecution == false && replicaWriteBytes > replicaLimits) {
             long replicaBytesWithoutOperation = replicaWriteBytes - bytes;
             this.currentReplicaBytes.getAndAdd(-bytes);

server/src/main/java/org/elasticsearch/index/stats/IndexingPressureStats.java

@@ -94,6 +94,50 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(replicaRejections);
     }
 
+    public long getTotalCombinedCoordinatingAndPrimaryBytes() {
+        return totalCombinedCoordinatingAndPrimaryBytes;
+    }
+
+    public long getTotalCoordinatingBytes() {
+        return totalCoordinatingBytes;
+    }
+
+    public long getTotalPrimaryBytes() {
+        return totalPrimaryBytes;
+    }
+
+    public long getTotalReplicaBytes() {
+        return totalReplicaBytes;
+    }
+
+    public long getCurrentCombinedCoordinatingAndPrimaryBytes() {
+        return currentCombinedCoordinatingAndPrimaryBytes;
+    }
+
+    public long getCurrentCoordinatingBytes() {
+        return currentCoordinatingBytes;
+    }
+
+    public long getCurrentPrimaryBytes() {
+        return currentPrimaryBytes;
+    }
+
+    public long getCurrentReplicaBytes() {
+        return currentReplicaBytes;
+    }
+
+    public long getCoordinatingRejections() {
+        return coordinatingRejections;
+    }
+
+    public long getPrimaryRejections() {
+        return primaryRejections;
+    }
+
+    public long getReplicaRejections() {
+        return replicaRejections;
+    }
+
     private static final String COMBINED = "combined_coordinating_and_primary";
     private static final String COMBINED_IN_BYTES = "combined_coordinating_and_primary_in_bytes";
     private static final String COORDINATING = "coordinating";

server/src/test/java/org/elasticsearch/index/IndexingPressureTests.java

@@ -0,0 +1,136 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.elasticsearch.index;
+
+import org.elasticsearch.common.lease.Releasable;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
+import org.elasticsearch.index.stats.IndexingPressureStats;
+import org.elasticsearch.test.ESTestCase;
+
+public class IndexingPressureTests extends ESTestCase {
+
+    private final Settings settings = Settings.builder().put(IndexingPressure.MAX_INDEXING_BYTES.getKey(), "10KB").build();
+
+    public void testMemoryBytesMarkedAndReleased() {
+        IndexingPressure indexingPressure = new IndexingPressure(settings);
+        try (Releasable coordinating = indexingPressure.markCoordinatingOperationStarted(10);
+             Releasable coordinating2 = indexingPressure.markCoordinatingOperationStarted(50);
+             Releasable primary = indexingPressure.markPrimaryOperationStarted(15, true);
+             Releasable primary2 = indexingPressure.markPrimaryOperationStarted(5, false);
+             Releasable replica = indexingPressure.markReplicaOperationStarted(25, true);
+             Releasable replica2 = indexingPressure.markReplicaOperationStarted(10, false)) {
+            IndexingPressureStats stats = indexingPressure.stats();
+            assertEquals(60, stats.getCurrentCoordinatingBytes());
+            assertEquals(20, stats.getCurrentPrimaryBytes());
+            assertEquals(80, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+            assertEquals(35, stats.getCurrentReplicaBytes());
+        }
+        IndexingPressureStats stats = indexingPressure.stats();
+        assertEquals(0, stats.getCurrentCoordinatingBytes());
+        assertEquals(0, stats.getCurrentPrimaryBytes());
+        assertEquals(0, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+        assertEquals(0, stats.getCurrentReplicaBytes());
+        assertEquals(60, stats.getTotalCoordinatingBytes());
+        assertEquals(20, stats.getTotalPrimaryBytes());
+        assertEquals(80, stats.getTotalCombinedCoordinatingAndPrimaryBytes());
+        assertEquals(35, stats.getTotalReplicaBytes());
+    }
+
+    public void testAvoidDoubleAccounting() {
+        IndexingPressure indexingPressure = new IndexingPressure(settings);
+        try (Releasable coordinating = indexingPressure.markCoordinatingOperationStarted(10);
+             Releasable primary = indexingPressure.markPrimaryOperationLocalToCoordinatingNodeStarted(15)) {
+            IndexingPressureStats stats = indexingPressure.stats();
+            assertEquals(10, stats.getCurrentCoordinatingBytes());
+            assertEquals(15, stats.getCurrentPrimaryBytes());
+            assertEquals(10, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+        }
+        IndexingPressureStats stats = indexingPressure.stats();
+        assertEquals(0, stats.getCurrentCoordinatingBytes());
+        assertEquals(0, stats.getCurrentPrimaryBytes());
+        assertEquals(0, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+        assertEquals(10, stats.getTotalCoordinatingBytes());
+        assertEquals(15, stats.getTotalPrimaryBytes());
+        assertEquals(10, stats.getTotalCombinedCoordinatingAndPrimaryBytes());
+    }
+
+    public void testCoordinatingPrimaryRejections() {
+        IndexingPressure indexingPressure = new IndexingPressure(settings);
+        try (Releasable coordinating = indexingPressure.markCoordinatingOperationStarted(1024 * 3);
+             Releasable primary = indexingPressure.markPrimaryOperationStarted(1024 * 3, false);
+             Releasable replica = indexingPressure.markReplicaOperationStarted(1024 * 3, false)) {
+            if (randomBoolean()) {
+                expectThrows(EsRejectedExecutionException.class, () -> indexingPressure.markCoordinatingOperationStarted(1024 * 2));
+                IndexingPressureStats stats = indexingPressure.stats();
+                assertEquals(1, stats.getCoordinatingRejections());
+                assertEquals(1024 * 6, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+            } else {
+                expectThrows(EsRejectedExecutionException.class, () -> indexingPressure.markPrimaryOperationStarted(1024 * 2, false));
+                IndexingPressureStats stats = indexingPressure.stats();
+                assertEquals(1, stats.getPrimaryRejections());
+                assertEquals(1024 * 6, stats.getCurrentCombinedCoordinatingAndPrimaryBytes());
+            }
+            long preForceRejections = indexingPressure.stats().getPrimaryRejections();
+            // Primary can be forced
+            Releasable forced = indexingPressure.markPrimaryOperationStarted(1024 * 2, true);
+            assertEquals(preForceRejections, indexingPressure.stats().getPrimaryRejections());
+            assertEquals(1024 * 8, indexingPressure.stats().getCurrentCombinedCoordinatingAndPrimaryBytes());
+            forced.close();
+
+            // Local to coordinating node primary actions not rejected
+            IndexingPressureStats preLocalStats = indexingPressure.stats();
+            Releasable local = indexingPressure.markPrimaryOperationLocalToCoordinatingNodeStarted(1024 * 2);
+            assertEquals(preLocalStats.getPrimaryRejections(), indexingPressure.stats().getPrimaryRejections());
+            assertEquals(1024 * 6, indexingPressure.stats().getCurrentCombinedCoordinatingAndPrimaryBytes());
+            assertEquals(preLocalStats.getCurrentPrimaryBytes() + 1024 * 2, indexingPressure.stats().getCurrentPrimaryBytes());
+            local.close();
+        }
+
+        assertEquals(1024 * 8, indexingPressure.stats().getTotalCombinedCoordinatingAndPrimaryBytes());
+    }
+
+    public void testReplicaRejections() {
+        IndexingPressure indexingPressure = new IndexingPressure(settings);
+        try (Releasable coordinating = indexingPressure.markCoordinatingOperationStarted(1024 * 3);
+             Releasable primary = indexingPressure.markPrimaryOperationStarted(1024 * 3, false);
+             Releasable replica = indexingPressure.markReplicaOperationStarted(1024 * 3, false)) {
+            // Replica will not be rejected until replica bytes > 15KB
+            Releasable replica2 = indexingPressure.markReplicaOperationStarted(1024 * 11, false);
+            assertEquals(1024 * 14, indexingPressure.stats().getCurrentReplicaBytes());
+            // Replica will be rejected once we cross 15KB
+            expectThrows(EsRejectedExecutionException.class, () -> indexingPressure.markReplicaOperationStarted(1024 * 2, false));
+            IndexingPressureStats stats = indexingPressure.stats();
+            assertEquals(1, stats.getReplicaRejections());
+            assertEquals(1024 * 14, stats.getCurrentReplicaBytes());
+
+            // Replica can be forced
+            Releasable forced = indexingPressure.markPrimaryOperationStarted(1024 * 2, true);
+            assertEquals(1, indexingPressure.stats().getReplicaRejections());
+            assertEquals(1024 * 14, indexingPressure.stats().getCurrentReplicaBytes());
+            forced.close();
+
+            replica2.close();
+            forced.close();
+        }
+
+        assertEquals(1024 * 14, indexingPressure.stats().getTotalReplicaBytes());
+    }
+}