Merge branch 'main' into fix/SearchWithRandomDisconnectsIT_testSearchWithRandomDisconnects

drempapis · web-flow · commit efd025121c37 · 2025-04-01T16:18:59.000+03:00
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -467,7 +467,7 @@ expensive messages that will usually be discarded:
 
 Logging is an important behaviour of the system and sometimes deserves its own
 unit tests, especially if there is complex logic for computing what is logged
-and when to log it. You can use a `org.elasticsearch.test.MockLogAppender` to
+and when to log it. You can use a `org.elasticsearch.test.MockLog` to
 make assertions about the logs that are being emitted.
 
 Logging is a powerful diagnostic technique, but it is not the only possibility.
diff --git a/docs/changelog/125477.yaml b/docs/changelog/125477.yaml
@@ -0,0 +1,6 @@
+pr: 125477
+summary: Prevent get datafeeds stats API returning an error when local tasks are slow to stop
+area: Machine Learning
+type: bug
+issues:
+  - 104160
diff --git a/server/src/main/java/org/elasticsearch/inference/InferenceServiceResults.java b/server/src/main/java/org/elasticsearch/inference/InferenceServiceResults.java
@@ -24,25 +24,12 @@ public interface InferenceServiceResults extends NamedWriteable, ChunkedToXConte
 
     /**
      * <p>Transform the result to match the format required for the TransportCoordinatedInferenceAction.
-     * For the inference plugin TextEmbeddingResults, the {@link #transformToLegacyFormat()} transforms the
-     * results into an intermediate format only used by the plugin's return value. It doesn't align with what the
-     * TransportCoordinatedInferenceAction expects. TransportCoordinatedInferenceAction expects an ml plugin
-     * TextEmbeddingResults.</p>
-     *
-     * <p>For other results like SparseEmbeddingResults, this method can be a pass through to the transformToLegacyFormat.</p>
+     * TransportCoordinatedInferenceAction expects an ml plugin TextEmbeddingResults or SparseEmbeddingResults.</p>
      */
     default List<? extends InferenceResults> transformToCoordinationFormat() {
         throw new UnsupportedOperationException("transformToCoordinationFormat() is not implemented");
     }
 
-    /**
-     * Transform the result to match the format required for versions prior to
-     * {@link org.elasticsearch.TransportVersions#V_8_12_0}
-     */
-    default List<? extends InferenceResults> transformToLegacyFormat() {
-        throw new UnsupportedOperationException("transformToLegacyFormat() is not implemented");
-    }
-
     /**
      * Convert the result to a map to aid with test assertions
      */
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java
@@ -580,12 +580,8 @@ public Flow.Publisher<InferenceServiceResults.Result> publisher() {
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
-                out.writeNamedWriteable(results);
-            } else {
-                out.writeNamedWriteable(results.transformToLegacyFormat().get(0));
-            }
             // streaming isn't supported via Writeable yet
+            out.writeNamedWriteable(results);
         }
 
         @Override
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/ChatCompletionResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/ChatCompletionResults.java
@@ -68,11 +68,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
         return results;
     }
 
-    @Override
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        throw new UnsupportedOperationException();
-    }
-
     public List<Result> getResults() {
         return results;
     }
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/RankedDocsResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/RankedDocsResults.java
@@ -196,11 +196,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
         throw new UnsupportedOperationException("Coordination format not supported by " + NAME);
     }
 
-    @Override
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        throw new UnsupportedOperationException("Legacy format not supported by " + NAME);
-    }
-
     @Override
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/SparseEmbeddingResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/SparseEmbeddingResults.java
@@ -100,11 +100,6 @@ public Map<String, Object> asMap() {
 
     @Override
     public List<? extends InferenceResults> transformToCoordinationFormat() {
-        return transformToLegacyFormat();
-    }
-
-    @Override
-    public List<? extends InferenceResults> transformToLegacyFormat() {
         return embeddings.stream()
             .map(
                 embedding -> new TextExpansionResults(
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingBitResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingBitResults.java
@@ -83,16 +83,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
             .toList();
     }
 
-    @Override
-    @SuppressWarnings("deprecation")
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        var legacyEmbedding = new LegacyTextEmbeddingResults(
-            embeddings.stream().map(embedding -> new LegacyTextEmbeddingResults.Embedding(embedding.toFloatArray())).toList()
-        );
-
-        return List.of(legacyEmbedding);
-    }
-
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
         map.put(TEXT_EMBEDDING_BITS, embeddings);
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java
@@ -85,16 +85,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
             .toList();
     }
 
-    @Override
-    @SuppressWarnings("deprecation")
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        var legacyEmbedding = new LegacyTextEmbeddingResults(
-            embeddings.stream().map(embedding -> new LegacyTextEmbeddingResults.Embedding(embedding.toFloatArray())).toList()
-        );
-
-        return List.of(legacyEmbedding);
-    }
-
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
         map.put(TEXT_EMBEDDING_BYTES, embeddings);
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingFloatResults.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingFloatResults.java
@@ -122,16 +122,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
         return embeddings.stream().map(embedding -> new MlTextEmbeddingResults(TEXT_EMBEDDING, embedding.asDoubleArray(), false)).toList();
     }
 
-    @Override
-    @SuppressWarnings("deprecation")
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        var legacyEmbedding = new LegacyTextEmbeddingResults(
-            embeddings.stream().map(embedding -> new LegacyTextEmbeddingResults.Embedding(embedding.values)).toList()
-        );
-
-        return List.of(legacyEmbedding);
-    }
-
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
         map.put(TEXT_EMBEDDING, embeddings);
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedRunningStateAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedRunningStateAction.java
@@ -146,12 +146,28 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
         private final Map<String, RunningState> datafeedRunningState;
 
+        private static RunningState selectMostRecentState(RunningState state1, RunningState state2) {
+
+            if (state1.searchInterval != null && state2.searchInterval != null) {
+                return state1.searchInterval.startMs() > state2.searchInterval.startMs() ? state1 : state2;
+            }
+
+            if (state1.searchInterval != null) {
+                return state1;
+            }
+            if (state2.searchInterval != null) {
+                return state2;
+            }
+
+            return state2;
+        }
+
         public static Response fromResponses(List<Response> responses) {
             return new Response(
                 responses.stream()
                     .flatMap(r -> r.datafeedRunningState.entrySet().stream())
                     .filter(entry -> entry.getValue() != null)
-                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
+                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, Response::selectMostRecentState))
             );
         }
 
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedRunningStateActionResponseTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedRunningStateActionResponseTests.java
@@ -9,6 +9,8 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xpack.core.ml.action.GetDatafeedRunningStateAction.Response;
+import org.elasticsearch.xpack.core.ml.action.GetDatafeedRunningStateAction.Response.RunningState;
+import org.elasticsearch.xpack.core.ml.datafeed.SearchInterval;
 import org.elasticsearch.xpack.core.ml.datafeed.SearchIntervalTests;
 
 import java.util.function.Function;
@@ -41,4 +43,68 @@ protected Writeable.Reader<Response> instanceReader() {
         return Response::new;
     }
 
+    /**
+     * Tests merging responses with the same datafeed ID but different running states,
+     * where both states have a searchInterval with different start times.
+     * The state with the more recent searchInterval (larger startMs value) should be selected.
+     */
+    public void testMergeWithDuplicateKeysAndDifferentSearchIntervals() {
+        SearchInterval olderInterval = new SearchInterval(1000L, 2000L);
+        SearchInterval newerInterval = new SearchInterval(3000L, 4000L);
+
+        RunningState olderState = new RunningState(true, true, olderInterval);
+        RunningState newerState = new RunningState(false, false, newerInterval);
+
+        String datafeedId = "test-datafeed";
+        Response response1 = Response.fromTaskAndState(datafeedId, olderState);
+        Response response2 = Response.fromTaskAndState(datafeedId, newerState);
+
+        Response mergedResponse = Response.fromResponses(java.util.List.of(response1, response2));
+
+        assertEquals(newerState, mergedResponse.getRunningState(datafeedId).orElse(null));
+
+        mergedResponse = Response.fromResponses(java.util.List.of(response2, response1));
+        assertEquals(newerState, mergedResponse.getRunningState(datafeedId).orElse(null));
+    }
+
+    /**
+     * Tests merging responses with the same datafeed ID but different running states,
+     * where only one state has a searchInterval.
+     * The state with the searchInterval should be selected, regardless of order.
+     */
+    public void testMergeWithDuplicateKeysWhenOnlyOneHasSearchInterval() {
+        SearchInterval interval = new SearchInterval(1000L, 2000L);
+
+        RunningState stateWithInterval = new RunningState(true, true, interval);
+        RunningState stateWithoutInterval = new RunningState(false, false, null);
+
+        String datafeedId = "test-datafeed";
+        Response response1 = Response.fromTaskAndState(datafeedId, stateWithInterval);
+        Response response2 = Response.fromTaskAndState(datafeedId, stateWithoutInterval);
+
+        Response mergedResponse = Response.fromResponses(java.util.List.of(response1, response2));
+
+        assertEquals(stateWithInterval, mergedResponse.getRunningState(datafeedId).orElse(null));
+
+        mergedResponse = Response.fromResponses(java.util.List.of(response2, response1));
+        assertEquals(stateWithInterval, mergedResponse.getRunningState(datafeedId).orElse(null));
+    }
+
+    /**
+     * Tests merging responses with the same datafeed ID but different running states,
+     * where neither state has a searchInterval.
+     * In this case, the second state in the list should be selected.
+     */
+    public void testMergeWithDuplicateKeysWhenNeitherHasSearchInterval() {
+        RunningState state1 = new RunningState(true, true, null);
+        RunningState state2 = new RunningState(false, false, null);
+
+        String datafeedId = "test-datafeed";
+        Response response1 = Response.fromTaskAndState(datafeedId, state1);
+        Response response2 = Response.fromTaskAndState(datafeedId, state2);
+
+        Response mergedResponse = Response.fromResponses(java.util.List.of(response1, response2));
+
+        assertEquals(state2, mergedResponse.getRunningState(datafeedId).orElse(null));
+    }
 }
diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListenerTests.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListenerTests.java
@@ -194,11 +194,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
             return List.of();
         }
 
-        @Override
-        public List<? extends InferenceResults> transformToLegacyFormat() {
-            return List.of();
-        }
-
         @Override
         public Map<String, Object> asMap() {
             return Map.of();
@@ -283,11 +278,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
             return List.of();
         }
 
-        @Override
-        public List<? extends InferenceResults> transformToLegacyFormat() {
-            return List.of();
-        }
-
         @Override
         public Map<String, Object> asMap() {
             return Map.of();
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceAuthorizationResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceAuthorizationResponseEntity.java
@@ -154,11 +154,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {
         throw new UnsupportedOperationException("Not implemented");
     }
 
-    @Override
-    public List<? extends InferenceResults> transformToLegacyFormat() {
-        throw new UnsupportedOperationException("not implemented");
-    }
-
     @Override
     public Map<String, Object> asMap() {
         throw new UnsupportedOperationException("Not implemented");
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/InferenceActionResponseTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/InferenceActionResponseTests.java
@@ -22,9 +22,6 @@
 import java.util.ArrayList;
 import java.util.List;
 
-import static org.elasticsearch.TransportVersions.V_8_12_0;
-import static org.elasticsearch.xpack.core.inference.action.InferenceAction.Response.transformToServiceResults;
-
 public class InferenceActionResponseTests extends AbstractBWCWireSerializationTestCase<InferenceAction.Response> {
 
     @Override
@@ -58,61 +55,6 @@ protected InferenceAction.Response mutateInstance(InferenceAction.Response insta
 
     @Override
     protected InferenceAction.Response mutateInstanceForVersion(InferenceAction.Response instance, TransportVersion version) {
-        if (version.before(V_8_12_0)) {
-            var singleResultList = instance.getResults().transformToLegacyFormat().subList(0, 1);
-            return new InferenceAction.Response(transformToServiceResults(singleResultList));
-        }
-
         return instance;
     }
-
-    public void testSerializesInferenceServiceResultsAddedVersion() throws IOException {
-        var instance = createTestInstance();
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    public void testSerializesOpenAiAddedVersion_UsingLegacyTextEmbeddingResult() throws IOException {
-        var embeddingResults = LegacyMlTextEmbeddingResultsTests.createRandomResults().transformToTextEmbeddingResults();
-        var instance = new InferenceAction.Response(embeddingResults);
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    public void testSerializesOpenAiAddedVersion_UsingSparseEmbeddingResult() throws IOException {
-        var embeddingResults = SparseEmbeddingResultsTests.createRandomResults();
-        var instance = new InferenceAction.Response(embeddingResults);
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    public void testSerializesMultipleInputsVersion_UsingLegacyTextEmbeddingResult() throws IOException {
-        var embeddingResults = TextEmbeddingFloatResultsTests.createRandomResults();
-        var instance = new InferenceAction.Response(embeddingResults);
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    public void testSerializesMultipleInputsVersion_UsingSparseEmbeddingResult() throws IOException {
-        var embeddingResults = SparseEmbeddingResultsTests.createRandomResults();
-        var instance = new InferenceAction.Response(embeddingResults);
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    // Technically we should never see a text embedding result in the transport version of this test because support
-    // for it wasn't added until openai
-    public void testSerializesSingleInputVersion_UsingLegacyTextEmbeddingResult() throws IOException {
-        var embeddingResults = TextEmbeddingFloatResultsTests.createRandomResults();
-        var instance = new InferenceAction.Response(embeddingResults);
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
-
-    public void testSerializesSingleVersion_UsingSparseEmbeddingResult() throws IOException {
-        var embeddingResults = SparseEmbeddingResultsTests.createRandomResults().transformToLegacyFormat().subList(0, 1);
-        var instance = new InferenceAction.Response(transformToServiceResults(embeddingResults));
-        var copy = copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), V_8_12_0);
-        assertOnBWCObject(copy, instance, V_8_12_0);
-    }
 }

Original file line number	Diff line number	Diff line change
`@@ -68,11 +68,6 @@ public List<? extends InferenceResults> transformToCoordinationFormat() {`
`68`	`68`	`return results;`
`69`	`69`	`}`
`70`	`70`
`71`		`- @Override`
`72`		`- public List<? extends InferenceResults> transformToLegacyFormat() {`
`73`		`- throw new UnsupportedOperationException();`
`74`		`- }`
`75`		`-`
`76`	`71`	`public List<Result> getResults() {`
`77`	`72`	`return results;`
`78`	`73`	`}`