Merge branch 'main' into fix/ftf_lj

Author: bpintea

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/VectorScorerBenchmark.java renamed to benchmarks/src/main/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmark.java

@@ -55,24 +55,23 @@
 /**
  * Benchmark that compares various scalar quantized vector similarity function
  * implementations;: scalar, lucene's panama-ized, and Elasticsearch's native.
- * Run with ./gradlew -p benchmarks run --args 'VectorScorerBenchmark'
+ * Run with ./gradlew -p benchmarks run --args 'Int7uScorerBenchmark'
  */
-public class VectorScorerBenchmark {
+public class Int7uScorerBenchmark {
 
     static {
         LogConfigurator.configureESLogging(); // native access requires logging to be initialized
     }
 
     @Param({ "96", "768", "1024" })
-    int dims;
-    int size = 2; // there are only two vectors to compare
+    public int dims;
+    final int size = 2; // there are only two vectors to compare
 
     Directory dir;
     IndexInput in;
     VectorScorerFactory factory;
 
-    byte[] vec1;
-    byte[] vec2;
+    byte[] vec1, vec2;
     float vec1Offset;
     float vec2Offset;
     float scoreCorrectionConstant;
@@ -139,39 +138,6 @@ public void setup() throws IOException {
         nativeDotScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.DOT_PRODUCT, values, queryVec).get();
         luceneSqrScorerQuery = luceneScorer(values, VectorSimilarityFunction.EUCLIDEAN, queryVec);
         nativeSqrScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.EUCLIDEAN, values, queryVec).get();
-
-        // sanity
-        var f1 = dotProductLucene();
-        var f2 = dotProductNative();
-        var f3 = dotProductScalar();
-        if (f1 != f2) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
-        }
-        if (f1 != f3) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
-        }
-        // square distance
-        f1 = squareDistanceLucene();
-        f2 = squareDistanceNative();
-        f3 = squareDistanceScalar();
-        if (f1 != f2) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
-        }
-        if (f1 != f3) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
-        }
-
-        var q1 = dotProductLuceneQuery();
-        var q2 = dotProductNativeQuery();
-        if (q1 != q2) {
-            throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
-        }
-
-        var sqr1 = squareDistanceLuceneQuery();
-        var sqr2 = squareDistanceNativeQuery();
-        if (sqr1 != sqr2) {
-            throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
-        }
     }
 
     @TearDown

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/JDKVectorInt7uBenchmark.java

@@ -52,7 +52,7 @@ public class JDKVectorInt7uBenchmark {
 
     Arena arena;
 
-    @Param({ "1", "128", "207", "256", "300", "512", "702", "1024" })
+    @Param({ "1", "128", "207", "256", "300", "512", "702", "1024", "1536", "2048" })
     public int size;
 
     @Setup(Level.Iteration)

benchmarks/src/test/java/org/elasticsearch/benchmark/vector/Int7uScorerBenchmarkTests.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.benchmark.vector;
+
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.apache.lucene.util.Constants;
+import org.elasticsearch.test.ESTestCase;
+import org.junit.BeforeClass;
+import org.openjdk.jmh.annotations.Param;
+
+import java.util.Arrays;
+
+public class Int7uScorerBenchmarkTests extends ESTestCase {
+
+    final double delta = 1e-3;
+    final int dims;
+
+    public Int7uScorerBenchmarkTests(int dims) {
+        this.dims = dims;
+    }
+
+    @BeforeClass
+    public static void skipWindows() {
+        assumeFalse("doesn't work on windows yet", Constants.WINDOWS);
+    }
+
+    public void testDotProduct() throws Exception {
+        for (int i = 0; i < 100; i++) {
+            var bench = new Int7uScorerBenchmark();
+            bench.dims = dims;
+            bench.setup();
+            try {
+                float expected = bench.dotProductScalar();
+                assertEquals(expected, bench.dotProductLucene(), delta);
+                assertEquals(expected, bench.dotProductNative(), delta);
+
+                expected = bench.dotProductLuceneQuery();
+                assertEquals(expected, bench.dotProductNativeQuery(), delta);
+            } finally {
+                bench.teardown();
+            }
+        }
+    }
+
+    public void testSquareDistance() throws Exception {
+        for (int i = 0; i < 100; i++) {
+            var bench = new Int7uScorerBenchmark();
+            bench.dims = dims;
+            bench.setup();
+            try {
+                float expected = bench.squareDistanceScalar();
+                assertEquals(expected, bench.squareDistanceLucene(), delta);
+                assertEquals(expected, bench.squareDistanceNative(), delta);
+
+                expected = bench.squareDistanceLuceneQuery();
+                assertEquals(expected, bench.squareDistanceNativeQuery(), delta);
+            } finally {
+                bench.teardown();
+            }
+        }
+    }
+
+    @ParametersFactory
+    public static Iterable<Object[]> parametersFactory() {
+        try {
+            var params = Int7uScorerBenchmark.class.getField("dims").getAnnotationsByType(Param.class)[0].value();
+            return () -> Arrays.stream(params).map(Integer::parseInt).map(i -> new Object[] { i }).iterator();
+        } catch (NoSuchFieldException e) {
+            throw new AssertionError(e);
+        }
+    }
+}

build-tools/src/main/java/org/elasticsearch/gradle/testclusters/MockApmServer.java

@@ -9,18 +9,30 @@
 
 package org.elasticsearch.gradle.testclusters;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpHandler;
 import com.sun.net.httpserver.HttpServer;
 
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.lang3.stream.Streams;
 import org.gradle.api.logging.Logger;
 import org.gradle.api.logging.Logging;
+import org.slf4j.LoggerFactory;
 
+import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+import javax.annotation.concurrent.NotThreadSafe;
 
 /**
  * This is a server which just accepts lines of JSON code and if the JSON
@@ -32,102 +44,127 @@
  * <p>
  * The HTTP server used is the JDK embedded com.sun.net.httpserver
  */
+@NotThreadSafe
 public class MockApmServer {
     private static final Logger logger = Logging.getLogger(MockApmServer.class);
-    private int port;
+    private static final org.slf4j.Logger log = LoggerFactory.getLogger(MockApmServer.class);
+    private final Pattern metricFilter;
+    private final Pattern transactionFilter;
+    private final Pattern transactionExcludesFilter;
 
-    public MockApmServer(int port) {
-        this.port = port;
-    }
+    private HttpServer instance;
 
-    /**
-     * Simple main that starts a mock APM server and prints the port it is
-     * running on. This is not needed
-     * for testing, it is just a convenient template for trying things out
-     * if you want play around.
-     */
-    public static void main(String[] args) throws IOException, InterruptedException {
-        MockApmServer server = new MockApmServer(9999);
-        server.start();
+    public MockApmServer(String metricFilter, String transactionFilter, String transactionExcludesFilter) {
+        this.metricFilter = createWildcardPattern(metricFilter);
+        this.transactionFilter = createWildcardPattern(transactionFilter);
+        this.transactionExcludesFilter = createWildcardPattern(transactionExcludesFilter);
     }
 
-    private static volatile HttpServer instance;
+    private Pattern createWildcardPattern(String filter) {
+        if (filter == null || filter.isEmpty()) {
+            return null;
+        }
+        var pattern = Arrays.stream(filter.split(",\\s*"))
+            .map(Pattern::quote)
+            .map(s -> s.replace("*", "\\E.*\\Q"))
+            .collect(Collectors.joining(")|(", "(", ")"));
+        return Pattern.compile(pattern);
+    }
 
     /**
      * Start the Mock APM server. Just returns empty JSON structures for every incoming message
      *
-     * @return - the port the Mock APM server started on
      * @throws IOException
      */
-    public synchronized int start() throws IOException {
+    public void start() throws IOException {
         if (instance != null) {
-            String hostname = instance.getAddress().getHostName();
-            int port = instance.getAddress().getPort();
-            logger.lifecycle("MockApmServer is already running. Reusing on address:port " + hostname + ":" + port);
-            return port;
+            throw new IllegalStateException("MockApmServer already started");
         }
-        InetSocketAddress addr = new InetSocketAddress("0.0.0.0", port);
+        InetSocketAddress addr = new InetSocketAddress("0.0.0.0", 0);
         HttpServer server = HttpServer.create(addr, 10);
-        server.createContext("/exit", new ExitHandler());
         server.createContext("/", new RootHandler());
-
         server.start();
         instance = server;
         logger.lifecycle("MockApmServer started on port " + server.getAddress().getPort());
-        return server.getAddress().getPort();
     }
 
     public int getPort() {
-        return port;
+        if (instance == null) {
+            throw new IllegalStateException("MockApmServer not started");
+        }
+        return instance.getAddress().getPort();
     }
 
     /**
      * Stop the server gracefully if possible
      */
-    public synchronized void stop() {
-        logger.lifecycle("stopping apm server");
-        instance.stop(1);
-        instance = null;
+    public void stop() {
+        if (instance != null) {
+            logger.lifecycle("stopping apm server");
+            instance.stop(1);
+            instance = null;
+        }
     }
 
     class RootHandler implements HttpHandler {
         public void handle(HttpExchange t) {
             try {
                 InputStream body = t.getRequestBody();
-                ByteArrayOutputStream bytes = new ByteArrayOutputStream();
-                byte[] buffer = new byte[8 * 1024];
-                int lengthRead;
-                while ((lengthRead = body.read(buffer)) > 0) {
-                    bytes.write(buffer, 0, lengthRead);
+                if (metricFilter == null && transactionFilter == null) {
+                    logRequestBody(body);
+                } else {
+                    logFiltered(body);
                 }
-                logger.lifecycle(("MockApmServer reading JSON objects: " + bytes.toString()));
 
                 String response = "{}";
                 t.sendResponseHeaders(200, response.length());
-                OutputStream os = t.getResponseBody();
-                os.write(response.getBytes());
-                os.close();
+                try (OutputStream os = t.getResponseBody()) {
+                    os.write(response.getBytes());
+                }
             } catch (Exception e) {
                 e.printStackTrace();
             }
         }
-    }
 
-    static class ExitHandler implements HttpHandler {
-        private static final int STOP_TIME = 3;
+        private void logRequestBody(InputStream body) throws IOException {
+            ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+            IOUtils.copy(body, bytes);
+            logger.lifecycle(("MockApmServer reading JSON objects: " + bytes.toString()));
+        }
 
-        public void handle(HttpExchange t) {
-            try {
-                InputStream body = t.getRequestBody();
-                String response = "{}";
-                t.sendResponseHeaders(200, response.length());
-                OutputStream os = t.getResponseBody();
-                os.write(response.getBytes());
-                os.close();
-                instance.stop(STOP_TIME);
-                instance = null;
-            } catch (Exception e) {
-                e.printStackTrace();
+        private void logFiltered(InputStream body) throws IOException {
+            ObjectMapper mapper = new ObjectMapper();
+            try (BufferedReader reader = new BufferedReader(new InputStreamReader(body))) {
+                String line;
+                String tier = null;
+                String node = null;
+
+                while ((line = reader.readLine()) != null) {
+                    var jsonNode = mapper.readTree(line);
+
+                    if (jsonNode.has("metadata")) {
+                        node = jsonNode.path("metadata").path("service").path("node").path("configured_name").asText(null);
+                        tier = jsonNode.path("metadata").path("labels").path("node_tier").asText(null);
+                    } else if (transactionFilter != null && jsonNode.has("transaction")) {
+                        var transaction = jsonNode.get("transaction");
+                        var name = transaction.get("name").asText();
+                        if (transactionFilter.matcher(name).matches()
+                            && (transactionExcludesFilter == null || transactionExcludesFilter.matcher(name).matches() == false)) {
+                            logger.lifecycle("Transaction [{}/{}]: {}", node, tier, transaction);
+                        }
+                    } else if (metricFilter != null && jsonNode.has("metricset")) {
+                        var metricset = jsonNode.get("metricset");
+                        var samples = (ObjectNode) metricset.get("samples");
+                        for (var name : Streams.of(samples.fieldNames()).toList()) {
+                            if (metricFilter.matcher(name).matches() == false) {
+                                samples.remove(name);
+                            }
+                        }
+                        if (samples.isEmpty() == false) {
+                            logger.lifecycle("Metricset [{}/{}]", node, tier, metricset);
+                        }
+                    }
+                }
             }
         }
     }