Restrict the range of JVMs that need security policy flags set

build.gradle

@@ -145,7 +145,7 @@ githubRelease.doFirst {
 }
 
 tasks.withType(Test).configureEach { task ->
-    if (JavaVersion.current().compareTo(JavaVersion.VERSION_17) > 0) {
+    if (JavaVersion.current().compareTo(JavaVersion.VERSION_17) > 0 && JavaVersion.current().compareTo(JavaVersion.VERSION_21) < 0) {
         def policyFile = file("${projectDir}/src/main/plugin-metadata/plugin-security.policy").absolutePath
         task.jvmArgs += [
             "--add-opens", "java.base/java.lang=ALL-UNNAMED",

src/main/java/com/o19s/es/ltr/query/RankerQuery.java

@@ -203,7 +203,7 @@ public boolean isCacheable(LeafReaderContext ctx) {
         // XXX: this is not thread safe and may run into extremely weird issues
         // if the searcher uses the parallel collector
         // Hopefully elastic never runs
-        MutableSupplier<LtrRanker.FeatureVector> vectorSupplier = new Suppliers.FeatureVectorSupplier();
+        MutableSupplier<LtrRanker.FeatureVector> vectorSupplier = new Suppliers.MutableSupplier<>();
         FVLtrRankerWrapper ltrRankerWrapper = new FVLtrRankerWrapper(ranker, vectorSupplier);
         LtrRewriteContext context = new LtrRewriteContext(ranker, vectorSupplier);
         for (Query q : queries) {

src/main/java/com/o19s/es/ltr/utils/Suppliers.java

@@ -16,10 +16,8 @@
 
 package com.o19s.es.ltr.utils;
 
-import com.o19s.es.ltr.ranker.LtrRanker;
-import org.opensearch.core.Assertions;
-
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 public final class Suppliers {
@@ -66,33 +64,15 @@ public E get() {
      * A mutable supplier
      */
     public static class MutableSupplier<T> implements Supplier<T> {
-        T obj;
+        private final AtomicReference<T> ref = new AtomicReference<>();
 
         @Override
         public T get() {
-            return obj;
+            return ref.get();
         }
 
         public void set(T obj) {
-            this.obj = obj;
-        }
-    }
-
-    /**
-     * Simple wrapper to make sure we run on the same thread
-     */
-    public static class FeatureVectorSupplier extends MutableSupplier<LtrRanker.FeatureVector> {
-        private final long threadId = Assertions.ENABLED ? Thread.currentThread().getId() : 0;
-
-        public LtrRanker.FeatureVector get() {
-            assert threadId == Thread.currentThread().getId();
-            return super.get();
-        }
-
-        @Override
-        public void set(LtrRanker.FeatureVector obj) {
-            assert threadId == Thread.currentThread().getId();
-            super.set(obj);
+            this.ref.set(obj);
         }
     }
 }