Merge remote-tracking branch 'elastic/main' into merge-pages-limit

Author: dnhatn

docs/changelog/127939.yaml

@@ -0,0 +1,5 @@
+pr: 127939
+summary: Add Custom inference service
+area: Machine Learning
+type: enhancement
+issues: []

libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/EntitlementBootstrap.java

@@ -122,7 +122,7 @@ public static void bootstrap(
             suppressFailureLogPackages
         );
         exportInitializationToAgent();
-        loadAgent(findAgentJar());
+        loadAgent(findAgentJar(), EntitlementInitialization.class.getName());
     }
 
     private static Path getUserHome() {
@@ -134,11 +134,11 @@ private static Path getUserHome() {
     }
 
     @SuppressForbidden(reason = "The VirtualMachine API is the only way to attach a java agent dynamically")
-    private static void loadAgent(String agentPath) {
+    static void loadAgent(String agentPath, String entitlementInitializationClassName) {
         try {
             VirtualMachine vm = VirtualMachine.attach(Long.toString(ProcessHandle.current().pid()));
             try {
-                vm.loadAgent(agentPath, EntitlementInitialization.class.getName());
+                vm.loadAgent(agentPath, entitlementInitializationClassName);
             } finally {
                 vm.detach();
             }
@@ -154,7 +154,7 @@ private static void exportInitializationToAgent() {
         EntitlementInitialization.class.getModule().addExports(initPkg, unnamedModule);
     }
 
-    public static String findAgentJar() {
+    static String findAgentJar() {
         String propertyName = "es.entitlement.agentJar";
         String propertyValue = System.getProperty(propertyName);
         if (propertyValue != null) {

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java

@@ -36,11 +36,11 @@ public class EntitlementInitialization {
 
     private static final Module ENTITLEMENTS_MODULE = PolicyManager.class.getModule();
 
-    private static ElasticsearchEntitlementChecker manager;
+    private static ElasticsearchEntitlementChecker checker;
 
     // Note: referenced by bridge reflectively
     public static EntitlementChecker checker() {
-        return manager;
+        return checker;
     }
 
     /**
@@ -63,17 +63,33 @@ public static EntitlementChecker checker() {
      * @param inst the JVM instrumentation class instance
      */
     public static void initialize(Instrumentation inst) throws Exception {
-        manager = initChecker();
+        checker = initChecker(inst, createPolicyManager());
+    }
 
-        var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
-        if (verifyBytecode) {
-            ensureClassesSensitiveToVerificationAreInitialized();
-        }
+    private static PolicyCheckerImpl createPolicyChecker(PolicyManager policyManager) {
+        EntitlementBootstrap.BootstrapArgs bootstrapArgs = EntitlementBootstrap.bootstrapArgs();
+        return new PolicyCheckerImpl(
+            bootstrapArgs.suppressFailureLogPackages(),
+            ENTITLEMENTS_MODULE,
+            policyManager,
+            bootstrapArgs.pathLookup()
+        );
+    }
 
-        DynamicInstrumentation.initialize(
-            inst,
-            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
-            verifyBytecode
+    private static PolicyManager createPolicyManager() {
+        EntitlementBootstrap.BootstrapArgs bootstrapArgs = EntitlementBootstrap.bootstrapArgs();
+        Map<String, Policy> pluginPolicies = bootstrapArgs.pluginPolicies();
+        PathLookup pathLookup = bootstrapArgs.pathLookup();
+
+        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
+
+        return new PolicyManager(
+            HardcodedEntitlements.serverPolicy(pathLookup.pidFile(), bootstrapArgs.serverPolicyPatch()),
+            HardcodedEntitlements.agentEntitlements(),
+            pluginPolicies,
+            EntitlementBootstrap.bootstrapArgs().scopeResolver(),
+            EntitlementBootstrap.bootstrapArgs().sourcePaths(),
+            pathLookup
         );
     }
 
@@ -99,9 +115,8 @@ private static void ensureClassesSensitiveToVerificationAreInitialized() {
         }
     }
 
-    private static ElasticsearchEntitlementChecker initChecker() {
-        final PolicyChecker policyChecker = createPolicyChecker();
-
+    static ElasticsearchEntitlementChecker initChecker(Instrumentation inst, PolicyManager policyManager) throws Exception {
+        final PolicyChecker policyChecker = createPolicyChecker(policyManager);
         final Class<?> clazz = EntitlementCheckerUtils.getVersionSpecificCheckerClass(
             ElasticsearchEntitlementChecker.class,
             Runtime.version().feature()
@@ -113,34 +128,25 @@ private static ElasticsearchEntitlementChecker initChecker() {
         } catch (NoSuchMethodException e) {
             throw new AssertionError("entitlement impl is missing required constructor: [" + clazz.getName() + "]", e);
         }
+
+        ElasticsearchEntitlementChecker checker;
         try {
-            return (ElasticsearchEntitlementChecker) constructor.newInstance(policyChecker);
+            checker = (ElasticsearchEntitlementChecker) constructor.newInstance(policyChecker);
         } catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
             throw new AssertionError(e);
         }
-    }
-
-    private static PolicyCheckerImpl createPolicyChecker() {
-        EntitlementBootstrap.BootstrapArgs bootstrapArgs = EntitlementBootstrap.bootstrapArgs();
-        Map<String, Policy> pluginPolicies = bootstrapArgs.pluginPolicies();
-        PathLookup pathLookup = bootstrapArgs.pathLookup();
 
-        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
+        var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
+        if (verifyBytecode) {
+            ensureClassesSensitiveToVerificationAreInitialized();
+        }
 
-        PolicyManager policyManager = new PolicyManager(
-            HardcodedEntitlements.serverPolicy(pathLookup.pidFile(), bootstrapArgs.serverPolicyPatch()),
-            HardcodedEntitlements.agentEntitlements(),
-            pluginPolicies,
-            EntitlementBootstrap.bootstrapArgs().scopeResolver(),
-            EntitlementBootstrap.bootstrapArgs().sourcePaths(),
-            pathLookup
-        );
-        return new PolicyCheckerImpl(
-            bootstrapArgs.suppressFailureLogPackages(),
-            ENTITLEMENTS_MODULE,
-            policyManager,
-            bootstrapArgs.pathLookup()
+        DynamicInstrumentation.initialize(
+            inst,
+            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
+            verifyBytecode
         );
-    }
 
+        return checker;
+    }
 }

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -183,6 +183,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED_8_19 = def(8_841_0_36);
     public static final TransportVersion ML_INFERENCE_SAGEMAKER_CHAT_COMPLETION_8_19 = def(8_841_0_37);
     public static final TransportVersion ML_INFERENCE_VERTEXAI_CHATCOMPLETION_ADDED_8_19 = def(8_841_0_38);
+    public static final TransportVersion INFERENCE_CUSTOM_SERVICE_ADDED_8_19 = def(8_841_0_39);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -269,8 +270,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion SETTINGS_IN_DATA_STREAMS_DRY_RUN = def(9_081_0_00);
     public static final TransportVersion ML_INFERENCE_SAGEMAKER_CHAT_COMPLETION = def(9_082_0_00);
     public static final TransportVersion ML_INFERENCE_VERTEXAI_CHATCOMPLETION_ADDED = def(9_083_0_00);
-
-    public static final TransportVersion ESQL_LIMIT_ROW_SIZE = def(9_084_0_00);
+    public static final TransportVersion INFERENCE_CUSTOM_SERVICE_ADDED = def(9_084_0_00);
+    public static final TransportVersion ESQL_LIMIT_ROW_SIZE = def(9_085_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/plugins/PluginDescriptor.java

@@ -250,6 +250,31 @@ public static PluginDescriptor readFromProperties(final Path pluginDir) throws I
         return descriptor;
     }
 
+    /**
+     * Reads the internal descriptor for a classic plugin.
+     *
+     * @param stream the InputStream from which to read the plugin data
+     * @return the plugin info
+     * @throws IOException if an I/O exception occurred reading the plugin descriptor
+     */
+    public static PluginDescriptor readInternalDescriptorFromStream(InputStream stream) throws IOException {
+        final Map<String, String> propsMap;
+        {
+            final Properties props = new Properties();
+            props.load(stream);
+            propsMap = props.stringPropertyNames().stream().collect(Collectors.toMap(Function.identity(), props::getProperty));
+        }
+
+        PluginDescriptor descriptor = readerInternalDescriptor(propsMap, INTERNAL_DESCRIPTOR_FILENAME);
+        String name = descriptor.getName();
+
+        if (propsMap.isEmpty() == false) {
+            throw new IllegalArgumentException("Unknown properties for plugin [" + name + "] in plugin descriptor: " + propsMap.keySet());
+        }
+
+        return descriptor;
+    }
+
     private static PluginDescriptor readerInternalDescriptor(Map<String, String> propsMap, String filename) {
         String name = readNonEmptyString(propsMap, filename, "name");
         String desc = readString(propsMap, name, "description");

server/src/test/java/org/elasticsearch/plugins/PluginDescriptorTests.java

@@ -107,7 +107,7 @@ void assertBothDescriptors(Consumer<DescriptorWriter> assertions) {
         assertions.accept(PluginDescriptorTests::mockStableDescriptor);
     }
 
-    public void testReadInternalDescriptor() throws Exception {
+    public void testReadInternalDescriptorFromStream() throws Exception {
         PluginDescriptor info = mockInternalDescriptor();
         assertEquals("my_plugin", info.getName());
         assertEquals("fake desc", info.getDescription());

test/framework/build.gradle

@@ -16,6 +16,7 @@ dependencies {
   api project(':libs:ssl-config')
   api project(":server")
   api project(":libs:cli")
+  api project(":libs:entitlement:bridge")
   api "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}"
   api "junit:junit:${versions.junit}"
   api "org.hamcrest:hamcrest:${versions.hamcrest}"

test/framework/src/main/java/org/elasticsearch/bootstrap/TestBuildInfo.java

@@ -11,4 +11,4 @@
 
 import java.util.List;
 
-record TestBuildInfo(String component, List<TestBuildInfoLocation> locations) {}
+public record TestBuildInfo(String component, List<TestBuildInfoLocation> locations) {}

test/framework/src/main/java/org/elasticsearch/bootstrap/TestBuildInfoParser.java

@@ -23,7 +23,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
-class TestBuildInfoParser {
+public class TestBuildInfoParser {
 
     private static final String PLUGIN_TEST_BUILD_INFO_RESOURCES = "META-INF/plugin-test-build-info.json";
     private static final String SERVER_TEST_BUILD_INFO_RESOURCE = "META-INF/server-test-build-info.json";
@@ -75,7 +75,7 @@ static TestBuildInfo fromXContent(final XContentParser parser) throws IOExceptio
         return PARSER.parse(parser, null).build();
     }
 
-    static List<TestBuildInfo> parseAllPluginTestBuildInfo() throws IOException {
+    public static List<TestBuildInfo> parseAllPluginTestBuildInfo() throws IOException {
         var xContent = XContentFactory.xContent(XContentType.JSON);
         List<TestBuildInfo> pluginsTestBuildInfos = new ArrayList<>();
         var resources = TestBuildInfoParser.class.getClassLoader().getResources(PLUGIN_TEST_BUILD_INFO_RESOURCES);
@@ -88,7 +88,7 @@ static List<TestBuildInfo> parseAllPluginTestBuildInfo() throws IOException {
         return pluginsTestBuildInfos;
     }
 
-    static TestBuildInfo parseServerTestBuildInfo() throws IOException {
+    public static TestBuildInfo parseServerTestBuildInfo() throws IOException {
         var xContent = XContentFactory.xContent(XContentType.JSON);
         var resource = TestBuildInfoParser.class.getClassLoader().getResource(SERVER_TEST_BUILD_INFO_RESOURCE);
         // No test-build-info for server: this might be a non-gradle build. Proceed without TestBuildInfo

test/framework/src/main/java/org/elasticsearch/bootstrap/TestScopeResolver.java

@@ -21,7 +21,7 @@
 import java.util.Map;
 import java.util.function.Function;
 
-record TestScopeResolver(Map<String, PolicyManager.PolicyScope> scopeMap) {
+public record TestScopeResolver(Map<String, PolicyManager.PolicyScope> scopeMap) {
 
     private static final Logger logger = LogManager.getLogger(TestScopeResolver.class);
 
@@ -38,7 +38,7 @@ PolicyManager.PolicyScope getScope(Class<?> callerClass) {
         return scope;
     }
 
-    static Function<Class<?>, PolicyManager.PolicyScope> createScopeResolver(
+    public static Function<Class<?>, PolicyManager.PolicyScope> createScopeResolver(
         TestBuildInfo serverBuildInfo,
         List<TestBuildInfo> pluginsBuildInfo
     ) {