From 6539c2a8efc84bd398b5ecaa23292d9eb140b144 Mon Sep 17 00:00:00 2001
From: carlosdelest <carlos.delgado@elastic.co>
Date: Thu, 1 Feb 2024 23:44:49 +0100
Subject: [PATCH] Fix detection of installed components

---
 .../org/elasticsearch/node/NodeConstruction.java     | 12 ++++++++++--
 .../xpack/inference/InferencePlugin.java             |  6 ++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java
index 2d9a4d5419992..0b52a4eeeab2d 100644
--- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java
+++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java
@@ -1090,10 +1090,12 @@ record PluginServiceInstances(
         }
 
         // Register noop versions of inference services if Inference plugin is not available
-        if (pluginComponents.stream().noneMatch(p -> p instanceof InferenceServiceRegistry)) {
+        if (isPluginComponentDefined(pluginComponents, InferenceServiceRegistry.class) == false) {
+            logger.warn("Inference service is not available");
             modules.bindToInstance(InferenceServiceRegistry.class, new InferenceServiceRegistry.NoopInferenceServiceRegistry());
         }
-        if (pluginComponents.stream().noneMatch(p -> p instanceof ModelRegistry)) {
+        if (isPluginComponentDefined(pluginComponents, ModelRegistry.class) == false) {
+            logger.warn("Model registry is not available");
             modules.bindToInstance(ModelRegistry.class, new ModelRegistry.NoopModelRegistry());
         }
 
@@ -1102,6 +1104,12 @@ record PluginServiceInstances(
         postInjection(clusterModule, actionModule, clusterService, transportService, featureService);
     }
 
+    private static boolean isPluginComponentDefined(Collection<?> pluginComponents, Class<?> clazz) {
+        return pluginComponents.stream()
+            .map(p -> p instanceof PluginComponentBinding ? ((PluginComponentBinding) p).impl() : p)
+            .anyMatch(p -> clazz.isAssignableFrom(clazz));
+    }
+
     private ClusterService createClusterService(SettingsModule settingsModule, ThreadPool threadPool, TaskManager taskManager) {
         ClusterService clusterService = new ClusterService(
             settingsModule.getSettings(),
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
index 5dda253386cdf..63033539f364a 100644
--- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
+++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
@@ -143,13 +143,15 @@ public Collection<?> createComponents(PluginServices services) {
         registry.init(services.client());
         inferenceServiceRegistry.set(registry);
 
-        return List.of(new PluginComponentBinding<>(ModelRegistry.class, modelRegistry), registry);
+        return List.of(
+            new PluginComponentBinding<>(ModelRegistry.class, modelRegistry),
+            new PluginComponentBinding<>(InferenceServiceRegistry.class, registry)
+        );
     }
 
     @Override
     public void loadExtensions(ExtensionLoader loader) {
         inferenceServiceExtensions = loader.loadExtensions(InferenceServiceExtension.class);
-        loader.loadExtensions(ModelRegistry.class);
     }
 
     public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {