From e2c66e4791a6fdb7217d087c0d33c960edaebdb6 Mon Sep 17 00:00:00 2001
From: Jonathan Buttner <jonathan.buttner@elastic.co>
Date: Wed, 24 Jan 2024 14:26:48 -0500
Subject: [PATCH] Closing services

---
 .../inference/InferenceServiceRegistry.java   | 20 ++++++-------------
 .../xpack/inference/InferencePlugin.java      |  8 ++------
 2 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java b/server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java
index a0ed7bbd82b24..d5973807d9d78 100644
--- a/server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java
+++ b/server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java
@@ -9,9 +9,9 @@
 package org.elasticsearch.inference;
 
 import org.elasticsearch.client.internal.Client;
-import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
@@ -20,7 +20,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-public class InferenceServiceRegistry extends AbstractLifecycleComponent {
+public class InferenceServiceRegistry implements Closeable {
 
     private final Map<String, InferenceService> services;
     private final List<NamedWriteableRegistry.Entry> namedWriteables = new ArrayList<>();
@@ -53,17 +53,9 @@ public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
     }
 
     @Override
-    protected void doStart() {
-
-    }
-
-    @Override
-    protected void doStop() {
-
-    }
-
-    @Override
-    protected void doClose() throws IOException {
-
+    public void close() throws IOException {
+        for (var service : services.values()) {
+            service.close();
+        }
     }
 }
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 4b3da18008eac..cad5b802d11f7 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
@@ -71,8 +71,6 @@ public class InferencePlugin extends Plugin implements ActionPlugin, ExtensibleP
     public static final String NAME = "inference";
     public static final String UTILITY_THREAD_POOL_NAME = "inference_utility";
     private final Settings settings;
-    // We'll keep a reference to the http manager just in case the inference services don't get closed individually
-    private final SetOnce<HttpClientManager> httpManager = new SetOnce<>();
     private final SetOnce<HttpRequestSenderFactory> httpFactory = new SetOnce<>();
     private final SetOnce<ServiceComponents> serviceComponents = new SetOnce<>();
 
@@ -119,11 +117,9 @@ public Collection<?> createComponents(PluginServices services) {
         var truncator = new Truncator(settings, services.clusterService());
         serviceComponents.set(new ServiceComponents(services.threadPool(), throttlerManager, settings, truncator));
 
-        httpManager.set(HttpClientManager.create(settings, services.threadPool(), services.clusterService(), throttlerManager));
-
         var httpRequestSenderFactory = new HttpRequestSenderFactory(
             services.threadPool(),
-            httpManager.get(),
+            HttpClientManager.create(settings, services.threadPool(), services.clusterService(), throttlerManager),
             services.clusterService(),
             settings
         );
@@ -234,6 +230,6 @@ public void close() {
         var serviceComponentsRef = serviceComponents.get();
         var throttlerToClose = serviceComponentsRef != null ? serviceComponentsRef.throttlerManager() : null;
 
-        IOUtils.closeWhileHandlingException(httpManager.get(), throttlerToClose);
+        IOUtils.closeWhileHandlingException(inferenceServiceRegistry.get(), throttlerToClose);
     }
 }