Adding delay for model registry call

Author: jonathan-buttner

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/InferenceRevokeDefaultEndpointsIT.java

@@ -280,7 +280,7 @@ private ElasticInferenceService createElasticInferenceService() {
         return new ElasticInferenceService(
             senderFactory,
             createWithEmptySettings(threadPool),
-            new ElasticInferenceServiceComponents(gatewayUrl),
+            ElasticInferenceServiceComponents.withNoRevokeDelay(gatewayUrl),
             modelRegistry,
             new ElasticInferenceServiceAuthorizationHandler(gatewayUrl, threadPool)
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -276,7 +276,7 @@ public Collection<?> createComponents(PluginServices services) {
         ElasticInferenceServiceSettings inferenceServiceSettings = new ElasticInferenceServiceSettings(settings);
         String elasticInferenceUrl = inferenceServiceSettings.getElasticInferenceServiceUrl();
 
-        var elasticInferenceServiceComponentsInstance = new ElasticInferenceServiceComponents(elasticInferenceUrl);
+        var elasticInferenceServiceComponentsInstance = ElasticInferenceServiceComponents.withDefaultRevokeDelay(elasticInferenceUrl);
         elasticInferenceServiceComponents.set(elasticInferenceServiceComponentsInstance);
 
         var authorizationHandler = new ElasticInferenceServiceAuthorizationHandler(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

@@ -280,9 +280,19 @@ private void handleRevokedDefaultConfigs(Set<String> authorizedDefaultModelIds)
             authorizationCompletedLatch.countDown();
         });
 
-        getServiceComponents().threadPool()
-            .executor(UTILITY_THREAD_POOL_NAME)
-            .execute(() -> modelRegistry.removeDefaultConfigs(unauthorizedDefaultInferenceEndpointIds, deleteInferenceEndpointsListener));
+        Runnable removeFromRegistry = () -> {
+            logger.debug("Synchronizing default inference endpoints");
+            modelRegistry.removeDefaultConfigs(unauthorizedDefaultInferenceEndpointIds, deleteInferenceEndpointsListener);
+        };
+
+        var delay = elasticInferenceServiceComponents.revokeAuthorizationDelay();
+        if (delay == null) {
+            getServiceComponents().threadPool().executor(UTILITY_THREAD_POOL_NAME).execute(removeFromRegistry);
+        } else {
+            getServiceComponents().threadPool()
+                .schedule(removeFromRegistry, delay, getServiceComponents().threadPool().executor(UTILITY_THREAD_POOL_NAME));
+        }
+
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceComponents.java

@@ -8,5 +8,23 @@
 package org.elasticsearch.xpack.inference.services.elastic;
 
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.TimeValue;
 
-public record ElasticInferenceServiceComponents(@Nullable String elasticInferenceServiceUrl) {}
+/**
+ * @param elasticInferenceServiceUrl the upstream Elastic Inference Server's URL
+ * @param revokeAuthorizationDelay Amount of time to wait before attempting to revoke authorization to certain model ids.
+ *                                 null indicates that there should be no delay
+ */
+public record ElasticInferenceServiceComponents(@Nullable String elasticInferenceServiceUrl, @Nullable TimeValue revokeAuthorizationDelay) {
+    private static final TimeValue DEFAULT_REVOKE_AUTHORIZATION_DELAY = TimeValue.timeValueMinutes(10);
+
+    public static final ElasticInferenceServiceComponents EMPTY_INSTANCE = new ElasticInferenceServiceComponents(null, null);
+
+    public static ElasticInferenceServiceComponents withNoRevokeDelay(String elasticInferenceServiceUrl) {
+        return new ElasticInferenceServiceComponents(elasticInferenceServiceUrl, null);
+    }
+
+    public static ElasticInferenceServiceComponents withDefaultRevokeDelay(String elasticInferenceServiceUrl) {
+        return new ElasticInferenceServiceComponents(elasticInferenceServiceUrl, DEFAULT_REVOKE_AUTHORIZATION_DELAY);
+    }
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSparseEmbeddingsModelTests.java

@@ -26,7 +26,7 @@ public static ElasticInferenceServiceSparseEmbeddingsModel createModel(String ur
             new ElasticInferenceServiceSparseEmbeddingsServiceSettings(modelId, maxInputTokens, null),
             EmptyTaskSettings.INSTANCE,
             EmptySecretSettings.INSTANCE,
-            new ElasticInferenceServiceComponents(url)
+            ElasticInferenceServiceComponents.withNoRevokeDelay(url)
         );
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java

@@ -1046,7 +1046,7 @@ private void testUnifiedStreamError(int responseCode, String responseJson, Strin
                 new ElasticInferenceServiceCompletionServiceSettings("model_id", new RateLimitSettings(100)),
                 EmptyTaskSettings.INSTANCE,
                 EmptySecretSettings.INSTANCE,
-                new ElasticInferenceServiceComponents(eisGatewayUrl)
+                ElasticInferenceServiceComponents.withNoRevokeDelay(eisGatewayUrl)
             );
             PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
             service.unifiedCompletionInfer(
@@ -1099,7 +1099,7 @@ private ElasticInferenceService createServiceWithMockSender(ElasticInferenceServ
         return new ElasticInferenceService(
             mock(HttpRequestSender.Factory.class),
             createWithEmptySettings(threadPool),
-            new ElasticInferenceServiceComponents(null),
+            ElasticInferenceServiceComponents.EMPTY_INSTANCE,
             mockModelRegistry(),
             mockAuthHandler
         );
@@ -1128,7 +1128,7 @@ private ElasticInferenceService createService(
         return new ElasticInferenceService(
             senderFactory,
             createWithEmptySettings(threadPool),
-            new ElasticInferenceServiceComponents(gatewayUrl),
+            ElasticInferenceServiceComponents.withNoRevokeDelay(gatewayUrl),
             mockModelRegistry(),
             mockAuthHandler
         );
@@ -1138,7 +1138,7 @@ private ElasticInferenceService createServiceWithAuthHandler(HttpRequestSender.F
         return new ElasticInferenceService(
             senderFactory,
             createWithEmptySettings(threadPool),
-            new ElasticInferenceServiceComponents(eisGatewayUrl),
+            ElasticInferenceServiceComponents.withNoRevokeDelay(eisGatewayUrl),
             mockModelRegistry(),
             new ElasticInferenceServiceAuthorizationHandler(eisGatewayUrl, threadPool)
         );

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/completion/ElasticInferenceServiceCompletionModelTests.java

@@ -29,7 +29,7 @@ public void testOverridingModelId() {
             new ElasticInferenceServiceCompletionServiceSettings("model_id", new RateLimitSettings(100)),
             EmptyTaskSettings.INSTANCE,
             EmptySecretSettings.INSTANCE,
-            new ElasticInferenceServiceComponents("url")
+            ElasticInferenceServiceComponents.withNoRevokeDelay("url")
         );
 
         var request = new UnifiedCompletionRequest(