Merge remote-tracking branch 'origin/main' into set_thread_context_fo…

…r_streaming

.buildkite/packer_cache.sh

@@ -29,6 +29,6 @@ for branch in "${branches[@]}"; do
   fi
 
   export JAVA_HOME="$HOME/.java/$ES_BUILD_JAVA"
-  "checkout/${branch}/gradlew" --project-dir "$CHECKOUT_DIR" --parallel -s resolveAllDependencies -Dorg.gradle.warning.mode=none -DisCI
+  "checkout/${branch}/gradlew" --project-dir "$CHECKOUT_DIR" --parallel -s resolveAllDependencies -Dorg.gradle.warning.mode=none -DisCI --max-workers=4
   rm -rf "checkout/${branch}"
 done

build-tools-internal/src/main/groovy/elasticsearch.ide.gradle

@@ -122,6 +122,36 @@ if (providers.systemProperty('idea.active').getOrNull() == 'true') {
       .findAll { it != null }
   }
 
+  // force IntelliJ to generate *.iml files for each imported module
+  tasks.register("enableExternalConfiguration") {
+    group = 'ide'
+    description = 'Enable per-module *.iml files'
+
+    doLast {
+      modifyXml('.idea/misc.xml') {xml ->
+        def externalStorageConfig = xml.component.find { it.'@name' == 'ExternalStorageConfigurationManager' }
+        if (externalStorageConfig) {
+          xml.remove(externalStorageConfig)
+        }
+      }
+    }
+  }
+
+  // modifies the idea module config to enable preview features on 'elasticsearch-native' module
+  tasks.register("enablePreviewFeatures") {
+    group = 'ide'
+    description = 'Enables preview features on native library module'
+    dependsOn tasks.named("enableExternalConfiguration")
+
+    doLast {
+      ['main', 'test'].each { sourceSet ->
+        modifyXml(".idea/modules/libs/native/elasticsearch.libs.elasticsearch-native.${sourceSet}.iml") { xml ->
+          xml.component.find { it.'@name' == 'NewModuleRootManager' }?.'@LANGUAGE_LEVEL' = 'JDK_21_PREVIEW'
+        }
+      }
+    }
+  }
+
   tasks.register('buildDependencyArtifacts') {
     group = 'ide'
     description = 'Builds artifacts needed as dependency for IDE modules'
@@ -149,7 +179,10 @@ if (providers.systemProperty('idea.active').getOrNull() == 'true') {
           testRunner = 'choose_per_test'
         }
         taskTriggers {
-          afterSync tasks.named('configureIdeCheckstyle'), tasks.named('configureIdeaGradleJvm'), tasks.named('buildDependencyArtifacts')
+          afterSync tasks.named('configureIdeCheckstyle'),
+            tasks.named('configureIdeaGradleJvm'),
+            tasks.named('buildDependencyArtifacts'),
+            tasks.named('enablePreviewFeatures')
         }
         encodings {
           encoding = 'UTF-8'

build-tools-internal/src/main/resources/checkstyle_suppressions.xml

@@ -35,6 +35,8 @@
   <!-- Intentionally have multi line string for a bulk request, otherwise this needs to fallback to string concatenation  -->
   <suppress files="modules[/\\]data-streams[/\\]src[/\\]javaRestTest[/\\]java[/\\]org[/\\]elasticsearch[/\\]datastreams[/\\]TsdbDataStreamRestIT.java" checks="LineLength" />
   <suppress files="qa[/\\]rolling-upgrade[/\\]src[/\\]javaRestTest[/\\]java[/\\]org[/\\]elasticsearch[/\\]upgrades[/\\]TsdbIT.java" checks="LineLength" />
+  <suppress files="qa[/\\]rolling-upgrade[/\\]src[/\\]javaRestTest[/\\]java[/\\]org[/\\]elasticsearch[/\\]upgrades[/\\]TsdbIndexingRollingUpgradeIT.java" checks="LineLength" />
+  <suppress files="qa[/\\]rolling-upgrade[/\\]src[/\\]javaRestTest[/\\]java[/\\]org[/\\]elasticsearch[/\\]upgrades[/\\]LogsdbIndexingRollingUpgradeIT.java" checks="LineLength" />
 
   <!-- Gradle requires inputs to be seriablizable -->
   <suppress files="build-tools-internal[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]gradle[/\\]internal[/\\]precommit[/\\]TestingConventionRule.java" checks="RegexpSinglelineJava" />

distribution/tools/entitlement-agent/impl/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterImpl.java

@@ -15,8 +15,10 @@
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.FieldVisitor;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.RecordComponentVisitor;
 import org.objectweb.asm.Type;
 
 import java.io.IOException;
@@ -73,7 +75,13 @@ public byte[] instrumentClass(String className, byte[] classfileBuffer) {
     }
 
     class EntitlementClassVisitor extends ClassVisitor {
-        final String className;
+
+        private static final String ENTITLEMENT_ANNOTATION = "EntitlementInstrumented";
+
+        private final String className;
+
+        private boolean isAnnotationPresent;
+        private boolean annotationNeeded = true;
 
         EntitlementClassVisitor(int api, ClassVisitor classVisitor, String className) {
             super(api, classVisitor);
@@ -85,25 +93,85 @@ public void visit(int version, int access, String name, String signature, String
             super.visit(version, access, name + classNameSuffix, signature, superName, interfaces);
         }
 
+        @Override
+        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+            if (visible && descriptor.equals(ENTITLEMENT_ANNOTATION)) {
+                isAnnotationPresent = true;
+                annotationNeeded = false;
+            }
+            return cv.visitAnnotation(descriptor, visible);
+        }
+
+        @Override
+        public void visitNestMember(String nestMember) {
+            addClassAnnotationIfNeeded();
+            super.visitNestMember(nestMember);
+        }
+
+        @Override
+        public void visitPermittedSubclass(String permittedSubclass) {
+            addClassAnnotationIfNeeded();
+            super.visitPermittedSubclass(permittedSubclass);
+        }
+
+        @Override
+        public void visitInnerClass(String name, String outerName, String innerName, int access) {
+            addClassAnnotationIfNeeded();
+            super.visitInnerClass(name, outerName, innerName, access);
+        }
+
+        @Override
+        public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
+            addClassAnnotationIfNeeded();
+            return super.visitField(access, name, descriptor, signature, value);
+        }
+
+        @Override
+        public RecordComponentVisitor visitRecordComponent(String name, String descriptor, String signature) {
+            addClassAnnotationIfNeeded();
+            return super.visitRecordComponent(name, descriptor, signature);
+        }
+
         @Override
         public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+            addClassAnnotationIfNeeded();
             var mv = super.visitMethod(access, name, descriptor, signature, exceptions);
-            boolean isStatic = (access & ACC_STATIC) != 0;
-            var key = new MethodKey(
-                className,
-                name,
-                Stream.of(Type.getArgumentTypes(descriptor)).map(Type::getInternalName).toList(),
-                isStatic
-            );
-            var instrumentationMethod = instrumentationMethods.get(key);
-            if (instrumentationMethod != null) {
-                // LOGGER.debug("Will instrument method {}", key);
-                return new EntitlementMethodVisitor(Opcodes.ASM9, mv, isStatic, descriptor, instrumentationMethod);
-            } else {
-                // LOGGER.trace("Will not instrument method {}", key);
+            if (isAnnotationPresent == false) {
+                boolean isStatic = (access & ACC_STATIC) != 0;
+                var key = new MethodKey(
+                    className,
+                    name,
+                    Stream.of(Type.getArgumentTypes(descriptor)).map(Type::getInternalName).toList(),
+                    isStatic
+                );
+                var instrumentationMethod = instrumentationMethods.get(key);
+                if (instrumentationMethod != null) {
+                    // LOGGER.debug("Will instrument method {}", key);
+                    return new EntitlementMethodVisitor(Opcodes.ASM9, mv, isStatic, descriptor, instrumentationMethod);
+                } else {
+                    // LOGGER.trace("Will not instrument method {}", key);
+                }
             }
             return mv;
         }
+
+        /**
+         * A class annotation can be added via visitAnnotation; we need to call visitAnnotation after all other visitAnnotation
+         * calls (in case one of them detects our annotation is already present), but before any other subsequent visit* method is called
+         * (up to visitMethod -- if no visitMethod is called, there is nothing to instrument).
+         * This includes visitNestMember, visitPermittedSubclass, visitInnerClass, visitField, visitRecordComponent and, of course,
+         * visitMethod (see {@link ClassVisitor} javadoc).
+         */
+        private void addClassAnnotationIfNeeded() {
+            if (annotationNeeded) {
+                // logger.debug("Adding {} annotation", ENTITLEMENT_ANNOTATION);
+                AnnotationVisitor av = cv.visitAnnotation(ENTITLEMENT_ANNOTATION, true);
+                if (av != null) {
+                    av.visitEnd();
+                }
+                annotationNeeded = false;
+            }
+        }
     }
 
     static class EntitlementMethodVisitor extends MethodVisitor {

distribution/tools/entitlement-agent/impl/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterTests.java

@@ -9,20 +9,24 @@
 
 package org.elasticsearch.entitlement.instrumentation.impl;
 
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.entitlement.api.EntitlementChecks;
 import org.elasticsearch.entitlement.api.EntitlementProvider;
 import org.elasticsearch.entitlement.instrumentation.InstrumentationService;
-import org.elasticsearch.entitlement.instrumentation.MethodKey;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.test.ESTestCase;
 import org.junit.Before;
+import org.objectweb.asm.Type;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
-import java.util.Map;
+import java.util.Arrays;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.entitlement.instrumentation.impl.ASMUtils.bytecode2text;
+import static org.elasticsearch.entitlement.instrumentation.impl.InstrumenterImpl.getClassFileInfo;
+import static org.hamcrest.Matchers.is;
 
 /**
  * This tests {@link InstrumenterImpl} in isolation, without a java agent.
@@ -60,6 +64,10 @@ public static class ClassToInstrument implements Testable {
         public static void systemExit(int status) {
             assertEquals(123, status);
         }
+
+        public static void anotherSystemExit(int status) {
+            assertEquals(123, status);
+        }
     }
 
     static final class TestException extends RuntimeException {}
@@ -76,8 +84,11 @@ public static class TestEntitlementManager implements EntitlementChecks {
          */
         volatile boolean isActive;
 
+        int checkSystemExitCallCount = 0;
+
         @Override
         public void checkSystemExit(Class<?> callerClass, int status) {
+            checkSystemExitCallCount++;
             assertSame(InstrumenterTests.class, callerClass);
             assertEquals(123, status);
             throwIfActive();
@@ -90,18 +101,11 @@ private void throwIfActive() {
         }
     }
 
-    public void test() throws Exception {
-        // This test doesn't replace ClassToInstrument in-place but instead loads a separate
-        // class ClassToInstrument_NEW that contains the instrumentation. Because of this,
-        // we need to configure the Transformer to use a MethodKey and instrumentationMethod
-        // with slightly different signatures (using the common interface Testable) which
-        // is not what would happen when it's run by the agent.
-
-        MethodKey k1 = instrumentationService.methodKeyForTarget(ClassToInstrument.class.getMethod("systemExit", int.class));
-        Method v1 = EntitlementChecks.class.getMethod("checkSystemExit", Class.class, int.class);
-        var instrumenter = new InstrumenterImpl("_NEW", Map.of(k1, v1));
+    public void testClassIsInstrumented() throws Exception {
+        var classToInstrument = ClassToInstrument.class;
+        var instrumenter = createInstrumenter(classToInstrument, "systemExit");
 
-        byte[] newBytecode = instrumenter.instrumentClassFile(ClassToInstrument.class).bytecodes();
+        byte[] newBytecode = instrumenter.instrumentClassFile(classToInstrument).bytecodes();
 
         if (logger.isTraceEnabled()) {
             logger.trace("Bytecode after instrumentation:\n{}", bytecode2text(newBytecode));
@@ -112,22 +116,96 @@ public void test() throws Exception {
             newBytecode
         );
 
+        getTestChecks().isActive = false;
+
         // Before checking is active, nothing should throw
-        callStaticSystemExit(newClass, 123);
+        callStaticMethod(newClass, "systemExit", 123);
 
         getTestChecks().isActive = true;
 
         // After checking is activated, everything should throw
-        assertThrows(TestException.class, () -> callStaticSystemExit(newClass, 123));
+        assertThrows(TestException.class, () -> callStaticMethod(newClass, "systemExit", 123));
+    }
+
+    public void testClassIsNotInstrumentedTwice() throws Exception {
+        var classToInstrument = ClassToInstrument.class;
+        var instrumenter = createInstrumenter(classToInstrument, "systemExit");
+
+        InstrumenterImpl.ClassFileInfo initial = getClassFileInfo(classToInstrument);
+        var internalClassName = Type.getInternalName(classToInstrument);
+
+        byte[] instrumentedBytecode = instrumenter.instrumentClass(internalClassName, initial.bytecodes());
+        byte[] instrumentedTwiceBytecode = instrumenter.instrumentClass(internalClassName, instrumentedBytecode);
+
+        logger.trace(() -> Strings.format("Bytecode after 1st instrumentation:\n%s", bytecode2text(instrumentedBytecode)));
+        logger.trace(() -> Strings.format("Bytecode after 2nd instrumentation:\n%s", bytecode2text(instrumentedTwiceBytecode)));
+
+        Class<?> newClass = new TestLoader(Testable.class.getClassLoader()).defineClassFromBytes(
+            ClassToInstrument.class.getName() + "_NEW_NEW",
+            instrumentedTwiceBytecode
+        );
+
+        getTestChecks().isActive = true;
+        getTestChecks().checkSystemExitCallCount = 0;
+
+        assertThrows(TestException.class, () -> callStaticMethod(newClass, "systemExit", 123));
+        assertThat(getTestChecks().checkSystemExitCallCount, is(1));
+    }
+
+    public void testClassAllMethodsAreInstrumentedFirstPass() throws Exception {
+        var classToInstrument = ClassToInstrument.class;
+        var instrumenter = createInstrumenter(classToInstrument, "systemExit", "anotherSystemExit");
+
+        InstrumenterImpl.ClassFileInfo initial = getClassFileInfo(classToInstrument);
+        var internalClassName = Type.getInternalName(classToInstrument);
+
+        byte[] instrumentedBytecode = instrumenter.instrumentClass(internalClassName, initial.bytecodes());
+        byte[] instrumentedTwiceBytecode = instrumenter.instrumentClass(internalClassName, instrumentedBytecode);
+
+        logger.trace(() -> Strings.format("Bytecode after 1st instrumentation:\n%s", bytecode2text(instrumentedBytecode)));
+        logger.trace(() -> Strings.format("Bytecode after 2nd instrumentation:\n%s", bytecode2text(instrumentedTwiceBytecode)));
+
+        Class<?> newClass = new TestLoader(Testable.class.getClassLoader()).defineClassFromBytes(
+            ClassToInstrument.class.getName() + "_NEW_NEW",
+            instrumentedTwiceBytecode
+        );
+
+        getTestChecks().isActive = true;
+        getTestChecks().checkSystemExitCallCount = 0;
+
+        assertThrows(TestException.class, () -> callStaticMethod(newClass, "systemExit", 123));
+        assertThat(getTestChecks().checkSystemExitCallCount, is(1));
+
+        assertThrows(TestException.class, () -> callStaticMethod(newClass, "anotherSystemExit", 123));
+        assertThat(getTestChecks().checkSystemExitCallCount, is(2));
+    }
+
+    /** This test doesn't replace ClassToInstrument in-place but instead loads a separate
+     * class ClassToInstrument_NEW that contains the instrumentation. Because of this,
+     * we need to configure the Transformer to use a MethodKey and instrumentationMethod
+     * with slightly different signatures (using the common interface Testable) which
+     * is not what would happen when it's run by the agent.
+     */
+    private InstrumenterImpl createInstrumenter(Class<?> classToInstrument, String... methodNames) throws NoSuchMethodException {
+        Method v1 = EntitlementChecks.class.getMethod("checkSystemExit", Class.class, int.class);
+        var methods = Arrays.stream(methodNames).map(name -> {
+            try {
+                return instrumentationService.methodKeyForTarget(classToInstrument.getMethod(name, int.class));
+            } catch (NoSuchMethodException e) {
+                throw new RuntimeException(e);
+            }
+        }).collect(Collectors.toUnmodifiableMap(name -> name, name -> v1));
+
+        return new InstrumenterImpl("_NEW", methods);
     }
 
     /**
      * Calling a static method of a dynamically loaded class is significantly more cumbersome
      * than calling a virtual method.
      */
-    private static void callStaticSystemExit(Class<?> c, int status) throws NoSuchMethodException, IllegalAccessException {
+    private static void callStaticMethod(Class<?> c, String methodName, int status) throws NoSuchMethodException, IllegalAccessException {
         try {
-            c.getMethod("systemExit", int.class).invoke(null, status);
+            c.getMethod(methodName, int.class).invoke(null, status);
         } catch (InvocationTargetException e) {
             Throwable cause = e.getCause();
             if (cause instanceof TestException n) {

distribution/tools/plugin-cli/src/main/java/org/elasticsearch/plugins/cli/SyncPluginsAction.java

@@ -25,6 +25,7 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
@@ -60,7 +61,7 @@ public SyncPluginsAction(Terminal terminal, Environment env) {
      * @throws UserException if a plugins config file is found.
      */
     public static void ensureNoConfigFile(Environment env) throws UserException {
-        final Path pluginsConfig = env.configFile().resolve("elasticsearch-plugins.yml");
+        final Path pluginsConfig = env.configFile().resolve(ELASTICSEARCH_PLUGINS_YML);
         if (Files.exists(pluginsConfig)) {
             throw new UserException(
                 ExitCodes.USAGE,
@@ -207,9 +208,8 @@ private List<InstallablePlugin> getPluginsToUpgrade(
         Optional<PluginsConfig> cachedPluginsConfig,
         List<PluginDescriptor> existingPlugins
     ) {
-        final Map<String, String> cachedPluginIdToLocation = cachedPluginsConfig.map(
-            config -> config.getPlugins().stream().collect(Collectors.toMap(InstallablePlugin::getId, InstallablePlugin::getLocation))
-        ).orElse(Map.of());
+        final Map<String, String> cachedPluginIdToLocation = new HashMap<>();
+        cachedPluginsConfig.ifPresent(config -> config.getPlugins().forEach(p -> cachedPluginIdToLocation.put(p.getId(), p.getLocation())));
 
         return pluginsToMaybeUpgrade.stream().filter(eachPlugin -> {
             final String eachPluginId = eachPlugin.getId();