[ML] Fixing inference stats race condition (#55163) (#55486)

`updateAndGet` could actually call the internal method more than once on contention.
If I read the JavaDocs, it says:
```* @param updateFunction a side-effect-free function```
So, it could be getting multiple updates on contention, thus having a race condition where stats are double counted.

To fix, I am going to use a `ReadWriteLock`. The `LongAdder` objects allows fast thread safe writes in high contention environments. These can be protected by the `ReadWriteLock::readLock`.

When stats are persisted, I need to call reset on all these adders. This is NOT thread safe if additions are taking place concurrently. So, I am going to protect with `ReadWriteLock::writeLock`.

This should prevent race conditions while allowing high (ish) throughput in the highly contention paths in inference.

I did some simple throughput tests and this change is not significantly slower and is simpler to grok (IMO).

closes  #54786

Author: benwtrent

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/InferenceStats.java

@@ -21,6 +21,8 @@
 import java.time.Instant;
 import java.util.Objects;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class InferenceStats implements ToXContentObject, Writeable {
 
@@ -204,6 +206,12 @@ public static class Accumulator {
         private final LongAdder failureCountAccumulator = new LongAdder();
         private final String modelId;
         private final String nodeId;
+        // curious reader
+        // you may be wondering why the lock set to the fair.
+        // When `currentStatsAndReset` is called, we want it guaranteed that it will eventually execute.
+        // If a ReadWriteLock is unfair, there are no such guarantees.
+        // A call for the `writelock::lock` could pause indefinitely.
+        private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock(true);
 
         public Accumulator(String modelId, String nodeId) {
             this.modelId = modelId;
@@ -226,22 +234,52 @@ public Accumulator merge(InferenceStats otherStats) {
         }
 
         public Accumulator incMissingFields() {
-            this.missingFieldsAccumulator.increment();
-            return this;
+            readWriteLock.readLock().lock();
+            try {
+                this.missingFieldsAccumulator.increment();
+                return this;
+            } finally {
+                readWriteLock.readLock().unlock();
+            }
         }
 
         public Accumulator incInference() {
-            this.inferenceAccumulator.increment();
-            return this;
+            readWriteLock.readLock().lock();
+            try {
+                this.inferenceAccumulator.increment();
+                return this;
+            } finally {
+                readWriteLock.readLock().unlock();
+            }
         }
 
         public Accumulator incFailure() {
-            this.failureCountAccumulator.increment();
-            return this;
+            readWriteLock.readLock().lock();
+            try {
+                this.failureCountAccumulator.increment();
+                return this;
+            } finally {
+                readWriteLock.readLock().unlock();
+            }
         }
 
-        public InferenceStats currentStats() {
-            return currentStats(Instant.now());
+        /**
+         * Thread safe.
+         *
+         * Returns the current stats and resets the values of all the counters.
+         * @return The current stats
+         */
+        public InferenceStats currentStatsAndReset() {
+            readWriteLock.writeLock().lock();
+            try {
+                InferenceStats stats = currentStats(Instant.now());
+                this.missingFieldsAccumulator.reset();
+                this.inferenceAccumulator.reset();
+                this.failureCountAccumulator.reset();
+                return stats;
+            } finally {
+                readWriteLock.writeLock().unlock();
+            }
         }
 
         public InferenceStats currentStats(Instant timeStamp) {

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/InferenceIngestIT.java

@@ -22,7 +22,9 @@
 import org.elasticsearch.test.ExternalTestCluster;
 import org.elasticsearch.test.SecuritySettingsSourceField;
 import org.elasticsearch.test.rest.ESRestTestCase;
+import org.elasticsearch.xpack.core.ml.MlStatsIndex;
 import org.elasticsearch.xpack.core.ml.inference.MlInferenceNamedXContentProvider;
+import org.elasticsearch.xpack.core.ml.inference.persistence.InferenceIndexConstants;
 import org.elasticsearch.xpack.core.ml.integration.MlRestTestStateCleaner;
 import org.junit.After;
 import org.junit.Before;
@@ -46,14 +48,16 @@ public class InferenceIngestIT extends ESRestTestCase {
         basicAuthHeaderValue("x_pack_rest_user", SecuritySettingsSourceField.TEST_PASSWORD_SECURE_STRING);
 
     @Before
-    public void createBothModels() throws Exception {
-        Request request = new Request("PUT", "_ml/inference/test_classification");
-        request.setJsonEntity(CLASSIFICATION_CONFIG);
-        client().performRequest(request);
-
-        request = new Request("PUT", "_ml/inference/test_regression");
-        request.setJsonEntity(REGRESSION_CONFIG);
-        client().performRequest(request);
+    public void setup() throws Exception {
+        Request loggingSettings = new Request("PUT", "_cluster/settings");
+        loggingSettings.setJsonEntity("" +
+            "{" +
+            "\"transient\" : {\n" +
+            "        \"logger.org.elasticsearch.xpack.ml.inference\" : \"TRACE\"\n" +
+            "    }" +
+            "}");
+        client().performRequest(loggingSettings);
+        client().performRequest(new Request("GET", "/_cluster/health?wait_for_status=green&timeout=30s"));
     }
 
     @Override
@@ -64,19 +68,33 @@ protected Settings restClientSettings() {
     @After
     public void cleanUpData() throws Exception {
         new MlRestTestStateCleaner(logger, adminClient()).clearMlMetadata();
+        client().performRequest(new Request("DELETE", InferenceIndexConstants.INDEX_PATTERN));
+        client().performRequest(new Request("DELETE", MlStatsIndex.indexPattern()));
+        Request loggingSettings = new Request("PUT", "_cluster/settings");
+        loggingSettings.setJsonEntity("" +
+            "{" +
+            "\"transient\" : {\n" +
+            "        \"logger.org.elasticsearch.xpack.ml.inference\" : null\n" +
+            "    }" +
+            "}");
+        client().performRequest(loggingSettings);
         ESRestTestCase.waitForPendingTasks(adminClient());
-        client().performRequest(new Request("DELETE", "_ml/inference/test_classification"));
-        client().performRequest(new Request("DELETE", "_ml/inference/test_regression"));
     }
 
     public void testPathologicalPipelineCreationAndDeletion() throws Exception {
+        String classificationModelId = "test_pathological_classification";
+        putModel(classificationModelId, CLASSIFICATION_CONFIG);
+
+        String regressionModelId = "test_pathological_regression";
+        putModel(regressionModelId, REGRESSION_CONFIG);
 
         for (int i = 0; i < 10; i++) {
-            client().performRequest(putPipeline("simple_classification_pipeline", CLASSIFICATION_PIPELINE));
+            client().performRequest(putPipeline("simple_classification_pipeline",
+                pipelineDefinition(classificationModelId, "classification")));
             client().performRequest(indexRequest("index_for_inference_test", "simple_classification_pipeline", generateSourceDoc()));
             client().performRequest(new Request("DELETE", "_ingest/pipeline/simple_classification_pipeline"));
 
-            client().performRequest(putPipeline("simple_regression_pipeline", REGRESSION_PIPELINE));
+            client().performRequest(putPipeline("simple_regression_pipeline", pipelineDefinition(regressionModelId, "regression")));
             client().performRequest(indexRequest("index_for_inference_test", "simple_regression_pipeline", generateSourceDoc()));
             client().performRequest(new Request("DELETE", "_ingest/pipeline/simple_regression_pipeline"));
         }
@@ -94,13 +112,30 @@ public void testPathologicalPipelineCreationAndDeletion() throws Exception {
                     QueryBuilders.existsQuery("ml.inference.classification.predicted_value"))));
 
         assertThat(EntityUtils.toString(searchResponse.getEntity()), containsString("\"value\":10"));
+        assertBusy(() -> {
+            try {
+                Response statsResponse = client().performRequest(new Request("GET",
+                    "_ml/inference/" + classificationModelId + "/_stats"));
+                assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":10"));
+                statsResponse = client().performRequest(new Request("GET", "_ml/inference/" + regressionModelId + "/_stats"));
+                assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":10"));
+            } catch (ResponseException ex) {
+                //this could just mean shard failures.
+                fail(ex.getMessage());
+            }
+        }, 30, TimeUnit.SECONDS);
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/54786")
     public void testPipelineIngest() throws Exception {
+        String classificationModelId = "test_classification";
+        putModel(classificationModelId, CLASSIFICATION_CONFIG);
 
-        client().performRequest(putPipeline("simple_classification_pipeline", CLASSIFICATION_PIPELINE));
-        client().performRequest(putPipeline("simple_regression_pipeline", REGRESSION_PIPELINE));
+        String regressionModelId = "test_regression";
+        putModel(regressionModelId, REGRESSION_CONFIG);
+
+        client().performRequest(putPipeline("simple_classification_pipeline",
+            pipelineDefinition(classificationModelId, "classification")));
+        client().performRequest(putPipeline("simple_regression_pipeline", pipelineDefinition(regressionModelId, "regression")));
 
         for (int i = 0; i < 10; i++) {
             client().performRequest(indexRequest("index_for_inference_test", "simple_classification_pipeline", generateSourceDoc()));
@@ -131,21 +166,30 @@ public void testPipelineIngest() throws Exception {
 
         assertBusy(() -> {
             try {
-                Response statsResponse = client().performRequest(new Request("GET", "_ml/inference/test_classification/_stats"));
+                Response statsResponse = client().performRequest(new Request("GET",
+                    "_ml/inference/" + classificationModelId + "/_stats"));
                 assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":10"));
-                statsResponse = client().performRequest(new Request("GET", "_ml/inference/test_regression/_stats"));
+                statsResponse = client().performRequest(new Request("GET", "_ml/inference/" + regressionModelId + "/_stats"));
                 assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":15"));
                 // can get both
                 statsResponse = client().performRequest(new Request("GET", "_ml/inference/_stats"));
-                assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":15"));
-                assertThat(EntityUtils.toString(statsResponse.getEntity()), containsString("\"inference_count\":10"));
+                String entityString = EntityUtils.toString(statsResponse.getEntity());
+                assertThat(entityString, containsString("\"inference_count\":15"));
+                assertThat(entityString, containsString("\"inference_count\":10"));
             } catch (ResponseException ex) {
                 //this could just mean shard failures.
+                fail(ex.getMessage());
             }
         }, 30, TimeUnit.SECONDS);
     }
 
     public void testSimulate() throws IOException {
+        String classificationModelId = "test_classification_simulate";
+        putModel(classificationModelId, CLASSIFICATION_CONFIG);
+
+        String regressionModelId = "test_regression_simulate";
+        putModel(regressionModelId, REGRESSION_CONFIG);
+
         String source = "{\n" +
             "  \"pipeline\": {\n" +
             "    \"processors\": [\n" +
@@ -157,7 +201,7 @@ public void testSimulate() throws IOException {
             "                \"top_classes_results_field\": \"result_class_prob\"," +
             "                \"num_top_feature_importance_values\": 2" +
             "          }},\n" +
-            "          \"model_id\": \"test_classification\",\n" +
+            "          \"model_id\": \"" + classificationModelId + "\",\n" +
             "          \"field_map\": {\n" +
             "            \"col1\": \"col1\",\n" +
             "            \"col2\": \"col2\",\n" +
@@ -169,7 +213,7 @@ public void testSimulate() throws IOException {
             "      {\n" +
             "        \"inference\": {\n" +
             "          \"target_field\": \"ml.regression\",\n" +
-            "          \"model_id\": \"test_regression\",\n" +
+            "          \"model_id\": \"" + regressionModelId + "\",\n" +
             "          \"inference_config\": {\"regression\":{}},\n" +
             "          \"field_map\": {\n" +
             "            \"col1\": \"col1\",\n" +
@@ -232,6 +276,8 @@ public void testSimulate() throws IOException {
     }
 
     public void testSimulateWithDefaultMappedField() throws IOException {
+        String classificationModelId = "test_classification_default_mapped_field";
+        putModel(classificationModelId, CLASSIFICATION_CONFIG);
         String source = "{\n" +
             "  \"pipeline\": {\n" +
             "    \"processors\": [\n" +
@@ -243,7 +289,7 @@ public void testSimulateWithDefaultMappedField() throws IOException {
             "                \"top_classes_results_field\": \"result_class_prob\"," +
             "                \"num_top_feature_importance_values\": 2" +
             "          }},\n" +
-            "          \"model_id\": \"test_classification\",\n" +
+            "          \"model_id\": \"" + classificationModelId + "\",\n" +
             "          \"field_map\": {}\n" +
             "        }\n" +
             "      }\n"+
@@ -607,36 +653,28 @@ protected NamedXContentRegistry xContentRegistry() {
         "  \"definition\": " + CLASSIFICATION_DEFINITION +
         "}";
 
-    private static final String CLASSIFICATION_PIPELINE = "{" +
-        "    \"processors\": [\n" +
-        "      {\n" +
-        "        \"inference\": {\n" +
-        "          \"model_id\": \"test_classification\",\n" +
-        "          \"tag\": \"classification\",\n" +
-        "          \"inference_config\": {\"classification\": {}},\n" +
-        "          \"field_map\": {\n" +
-        "            \"col1\": \"col1\",\n" +
-        "            \"col2\": \"col2\",\n" +
-        "            \"col3\": \"col3\",\n" +
-        "            \"col4\": \"col4\"\n" +
-        "          }\n" +
-        "        }\n" +
-        "      }]}\n";
-
-    private static final String REGRESSION_PIPELINE = "{" +
-        "    \"processors\": [\n" +
-        "      {\n" +
-        "        \"inference\": {\n" +
-        "          \"model_id\": \"test_regression\",\n" +
-        "          \"tag\": \"regression\",\n" +
-        "          \"inference_config\": {\"regression\": {}},\n" +
-        "          \"field_map\": {\n" +
-        "            \"col1\": \"col1\",\n" +
-        "            \"col2\": \"col2\",\n" +
-        "            \"col3\": \"col3\",\n" +
-        "            \"col4\": \"col4\"\n" +
-        "          }\n" +
-        "        }\n" +
-        "      }]}\n";
+    private static String pipelineDefinition(String modelId, String inferenceConfig) {
+        return "{" +
+            "    \"processors\": [\n" +
+            "      {\n" +
+            "        \"inference\": {\n" +
+            "          \"model_id\": \"" + modelId + "\",\n" +
+            "          \"tag\": \""+ inferenceConfig + "\",\n" +
+            "          \"inference_config\": {\"" + inferenceConfig + "\": {}},\n" +
+            "          \"field_map\": {\n" +
+            "            \"col1\": \"col1\",\n" +
+            "            \"col2\": \"col2\",\n" +
+            "            \"col3\": \"col3\",\n" +
+            "            \"col4\": \"col4\"\n" +
+            "          }\n" +
+            "        }\n" +
+            "      }]}\n";
+    }
+
+    private void putModel(String modelId, String modelConfiguration) throws IOException {
+        Request request = new Request("PUT", "_ml/inference/" + modelId);
+        request.setJsonEntity(modelConfiguration);
+        client().performRequest(request);
+    }
 
 }