Add unit-tests for Java 20 sources and update benchmarks

Signed-off-by: Ketan Verma <[email protected]>

benchmarks/build.gradle

@@ -64,6 +64,11 @@ compileJava.options.compilerArgs.addAll(["-processor", "org.openjdk.jmh.generato
 
 run.executable = "${BuildParams.runtimeJavaHome}/bin/java"
 
+// Add support for incubator modules on supported Java versions.
+if (BuildParams.runtimeJavaVersion >= JavaVersion.VERSION_20) {
+  run.jvmArgs += ["--add-modules=jdk.incubator.vector"]
+}
+
 // classes generated by JMH can use all sorts of forbidden APIs but we have no influence at all and cannot exclude these classes
 disableTasks('forbiddenApisMain')
 

benchmarks/src/main/java/org/opensearch/common/round/RoundableBenchmark.java

@@ -20,6 +20,7 @@
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
+import java.lang.invoke.MethodHandles;
 import java.util.Random;
 import java.util.function.Supplier;
 
@@ -83,7 +84,7 @@ public static class Options {
             "256" })
         public Integer size;
 
-        @Param({ "binary", "linear" })
+        @Param({ "linear", "binary", "btree" })
         public String type;
 
         @Param({ "uniform", "skewed_edge", "skewed_center" })
@@ -93,7 +94,7 @@ public static class Options {
         public Supplier<Roundable> supplier;
 
         @Setup
-        public void setup() {
+        public void setup() throws ClassNotFoundException, IllegalAccessException {
             Random random = new Random(size);
             long[] values = new long[size];
             for (int i = 1; i < values.length; i++) {
@@ -135,6 +136,19 @@ public void setup() {
                 case "linear":
                     supplier = () -> new BidirectionalLinearSearcher(values, size);
                     break;
+                case "btree":
+                    // Not supported below Java 20.
+                    // Using MethodHandles to reflectively look up the class (if available) in order to
+                    // avoid setting up separate Java 20 sources with mostly duplicate code.
+                    Class<?> clz = MethodHandles.lookup().findClass("org.opensearch.common.round.BtreeSearcher");
+                    supplier = () -> {
+                        try {
+                            return (Roundable) clz.getDeclaredConstructor(long[].class, int.class).newInstance(values, size);
+                        } catch (Exception e) {
+                            throw new RuntimeException(e);
+                        }
+                    };
+                    break;
                 default:
                     throw new IllegalArgumentException("invalid type: " + type);
             }

libs/common/build.gradle

@@ -52,14 +52,31 @@ if (BuildParams.runtimeJavaVersion >= JavaVersion.VERSION_20) {
         srcDirs = ['src/main/java20']
       }
     }
+
+    java20Test {
+      java {
+        srcDirs = ['src/test/java20']
+      }
+    }
   }
 
   configurations {
     java20Implementation.extendsFrom(implementation)
+    java20TestImplementation.extendsFrom(implementation)
   }
 
   dependencies {
     java20Implementation sourceSets.main.output
+
+    // Adding Java 20 sources as compile-only to make them visible during compilation,
+    // but using the multi-release JAR output to provide the runtime functionality.
+    // This helps avoid JarHell problems when an overridden class is present with the same name.
+    java20TestCompileOnly sourceSets.java20.output
+    java20TestImplementation files(jar.archiveFile)
+    java20TestImplementation sourceSets.test.output
+    java20TestImplementation(project(':test:framework')) {
+      exclude group: 'org.opensearch', module: 'opensearch-common'
+    }
   }
 
   compileJava20Java {
@@ -69,9 +86,11 @@ if (BuildParams.runtimeJavaVersion >= JavaVersion.VERSION_20) {
     options.compilerArgs -= '-Werror' // use of incubator modules is reported as a warning
   }
 
-  forbiddenApisJava20 {
-    failOnMissingClasses = false
-    ignoreSignaturesOfMissingClasses = true
+  compileJava20TestJava {
+    sourceCompatibility = JavaVersion.VERSION_20
+    targetCompatibility = JavaVersion.VERSION_20
+    options.compilerArgs += ['--add-modules', 'jdk.incubator.vector']
+    options.compilerArgs -= '-Werror' // use of incubator modules is reported as a warning
   }
 
   jar {
@@ -81,4 +100,24 @@ if (BuildParams.runtimeJavaVersion >= JavaVersion.VERSION_20) {
     }
     manifest.attributes('Multi-Release': 'true')
   }
+
+  forbiddenApisJava20 {
+    failOnMissingClasses = false
+    ignoreSignaturesOfMissingClasses = true
+  }
+
+  forbiddenApisJava20Test {
+    failOnMissingClasses = false
+    ignoreSignaturesOfMissingClasses = true
+  }
+
+  tasks.register('testJava20', Test) {
+    dependsOn jar
+    shouldRunAfter test
+    group = 'verification'
+    testClassesDirs = sourceSets.java20Test.output.classesDirs
+    classpath = sourceSets.java20Test.runtimeClasspath
+  }
+
+  check.dependsOn('testJava20')
 }

libs/common/src/main/java20/org/opensearch/common/round/BtreeSearcher.java

@@ -27,6 +27,10 @@ class BtreeSearcher implements Roundable {
     private final int shift;
     private final long[] values;
 
+    BtreeSearcher(long[] values, int size) {
+        this(values, size, LongVector.SPECIES_PREFERRED);
+    }
+
     BtreeSearcher(long[] values, int size, VectorSpecies<Long> species) {
         assert size > 0 : "at least one value must be present";
 
@@ -43,6 +47,7 @@ class BtreeSearcher implements Roundable {
 
     /**
      * Builds the B-tree memory layout.
+     * It builds the tree recursively, following an in-order traversal.
      *
      * <p>
      * Each block stores 'lanes' values at indices {@code i, i + 1, ..., i + lanes - 1} where {@code i} is the
@@ -61,7 +66,12 @@ private int build(long[] src, int i, int size, long[] dst, int j) {
         if (j < dst.length) {
             for (int k = 0; k < lanes; k++) {
                 i = build(src, i, size, dst, j + ((j + k) << shift));
-                dst[j + k] = (j + k < size + 1) ? src[i++] : Long.MAX_VALUE;
+
+                // Fills the B-tree as a complete tree, i.e., all levels are completely filled,
+                // except the last level which is filled from left to right.
+                // The trick is to fill the destination array between indices 1...size (inclusive / 1-indexed)
+                // and pad the remaining array with +infinity.
+                dst[j + k] = (j + k <= size) ? src[i++] : Long.MAX_VALUE;
             }
             i = build(src, i, size, dst, j + ((j + lanes) << shift));
         }

libs/common/src/test/java/org/opensearch/common/round/BidirectionalLinearSearcherTests.java

@@ -0,0 +1,16 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.common.round;
+
+public class BidirectionalLinearSearcherTests extends RoundableTestCase {
+    @Override
+    public Roundable newInstance(long[] values, int size) {
+        return new BidirectionalLinearSearcher(values, size);
+    }
+}

libs/common/src/test/java/org/opensearch/common/round/BinarySearcherTests.java

@@ -0,0 +1,16 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.common.round;
+
+public class BinarySearcherTests extends RoundableTestCase {
+    @Override
+    public Roundable newInstance(long[] values, int size) {
+        return new BinarySearcher(values, size);
+    }
+}

libs/common/src/test/java20/org/opensearch/common/round/BtreeSearcherTests.java

@@ -0,0 +1,20 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.common.round;
+
+import java.util.List;
+
+import jdk.incubator.vector.LongVector;
+
+public class BtreeSearcherTests extends RoundableTestCase {
+    @Override
+    public Roundable newInstance(long[] values, int size) {
+        return new BtreeSearcher(values, size, randomFrom(List.of(LongVector.SPECIES_128, LongVector.SPECIES_256, LongVector.SPECIES_512)));
+    }
+}

libs/common/src/test/java/org/opensearch/common/round/RoundableTests.java → test/framework/src/main/java/org/opensearch/common/round/RoundableTestCase.java

@@ -10,16 +10,19 @@
 
 import org.opensearch.test.OpenSearchTestCase;
 
-public class RoundableTests extends OpenSearchTestCase {
+public abstract class RoundableTestCase extends OpenSearchTestCase {
 
-    public void testFloor() {
-        int size = randomIntBetween(1, 256);
-        long[] values = new long[size];
-        for (int i = 1; i < values.length; i++) {
+    public abstract Roundable newInstance(long[] values, int size);
+
+    public final void testFloor() {
+        int size = randomIntBetween(1, 256); // number of values present in the array
+        int capacity = size + randomIntBetween(0, 20); // capacity of the array can be larger
+        long[] values = new long[capacity];
+        for (int i = 1; i < size; i++) {
             values[i] = values[i - 1] + (randomNonNegativeLong() % 200) + 1;
         }
 
-        Roundable[] impls = { new BinarySearcher(values, size), new BidirectionalLinearSearcher(values, size) };
+        Roundable roundable = newInstance(values, size);
 
         for (int i = 0; i < 100000; i++) {
             // Index of the expected round-down point.
@@ -35,23 +38,16 @@ public void testFloor() {
             // round-down point, which will still floor to the same value.
             long key = expected + (randomNonNegativeLong() % delta);
 
-            for (Roundable roundable : impls) {
-                assertEquals(expected, roundable.floor(key));
-            }
+            assertEquals(expected, roundable.floor(key));
         }
     }
 
-    public void testAssertions() {
+    public final void testAssertions() {
         AssertionError exception;
 
-        exception = assertThrows(AssertionError.class, () -> new BinarySearcher(new long[0], 0));
-        assertEquals("at least one value must be present", exception.getMessage());
-        exception = assertThrows(AssertionError.class, () -> new BidirectionalLinearSearcher(new long[0], 0));
+        exception = assertThrows(AssertionError.class, () -> newInstance(new long[0], 0));
         assertEquals("at least one value must be present", exception.getMessage());
-
-        exception = assertThrows(AssertionError.class, () -> new BinarySearcher(new long[] { 100 }, 1).floor(50));
-        assertEquals("key must be greater than or equal to 100", exception.getMessage());
-        exception = assertThrows(AssertionError.class, () -> new BidirectionalLinearSearcher(new long[] { 100 }, 1).floor(50));
+        exception = assertThrows(AssertionError.class, () -> newInstance(new long[] { 100 }, 1).floor(50));
         assertEquals("key must be greater than or equal to 100", exception.getMessage());
     }
 }