Addressing feedback

Author: jonathan-buttner

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

@@ -17,7 +17,9 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.persistent.AllocatedPersistentTask;
+import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.threadpool.Scheduler;
 import org.elasticsearch.xpack.core.ClientHelper;
 import org.elasticsearch.xpack.core.inference.action.StoreInferenceEndpointsAction;
@@ -133,6 +135,17 @@ public void waitForAuthorizationToComplete(TimeValue waitTime) {
         }
     }
 
+    // Overriding so tests in the same package can access
+    @Override
+    protected void init(
+        PersistentTasksService persistentTasksService,
+        TaskManager taskManager,
+        String persistentTaskId,
+        long allocationId
+    ) {
+        super.init(persistentTasksService, taskManager, persistentTaskId, allocationId);
+    }
+
     @Override
     protected void onCancelled() {
         shutdown();
@@ -142,11 +155,18 @@ protected void onCancelled() {
     // default for testing
     void shutdown() {
         shutdown.set(true);
-        if (lastAuthTask.get() != null) {
-            lastAuthTask.get().cancel();
+
+        var authTask = lastAuthTask.get();
+        if (authTask != null) {
+            authTask.cancel();
         }
     }
 
+    // default for testing
+    boolean isShutdown() {
+        return shutdown.get();
+    }
+
     private void scheduleAuthorizationRequest() {
         try {
             if (elasticInferenceServiceSettings.isPeriodicAuthorizationEnabled() == false) {
@@ -177,6 +197,8 @@ private void scheduleAuthorizationRequest() {
             );
         } catch (Exception e) {
             logger.warn("Failed scheduling authorization request", e);
+            // Shutdown and complete the task so it will be restarted
+            onCancelled();
         }
     }
 

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.core.FixForMultiProject;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.persistent.AllocatedPersistentTask;
 import org.elasticsearch.persistent.ClusterPersistentTasksCustomMetadata;
@@ -93,6 +94,10 @@ protected void nodeOperation(AllocatedPersistentTask task, AuthorizationTaskPara
         authPoller.start();
     }
 
+    @FixForMultiProject(
+        description = "A single cluster can have multiple projects, "
+            + "we'll need to either make a call per project/org or use a bulk authorization api that EIS provides"
+    )
     @Override
     public Scope scope() {
         return Scope.CLUSTER;

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

@@ -12,12 +12,15 @@
 import org.elasticsearch.common.util.concurrent.DeterministicTaskQueue;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.core.inference.action.StoreInferenceEndpointsAction;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceComponents;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettingsTests;
 import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
 import org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntity;
@@ -315,4 +318,64 @@ public void testSendsTwoAuthorizationRequests() throws InterruptedException {
         assertThat(callbackCount.get(), is(2));
         verify(mockClient, never()).execute(eq(StoreInferenceEndpointsAction.INSTANCE), any(), any());
     }
+
+    public void testCallsShutdownAndMarksTaskAsCompleted_WhenSchedulingFails() throws InterruptedException {
+        var mockRegistry = mock(ModelRegistry.class);
+        when(mockRegistry.isReady()).thenReturn(true);
+        when(mockRegistry.getInferenceIds()).thenReturn(Set.of("id1", "id2"));
+
+        var mockAuthHandler = mock(ElasticInferenceServiceAuthorizationRequestHandler.class);
+        doAnswer(invocation -> {
+            ActionListener<ElasticInferenceServiceAuthorizationModel> listener = invocation.getArgument(0);
+            listener.onResponse(
+                ElasticInferenceServiceAuthorizationModel.of(
+                    new ElasticInferenceServiceAuthorizationResponseEntity(
+                        List.of(
+                            new ElasticInferenceServiceAuthorizationResponseEntity.AuthorizedModel(
+                                // this is an unknown model id so it won't trigger storing an inference endpoint because
+                                // it doesn't map to a known one
+                                "abc",
+                                EnumSet.of(TaskType.SPARSE_EMBEDDING)
+                            )
+                        )
+                    )
+                )
+            );
+            return Void.TYPE;
+        }).when(mockAuthHandler).getAuthorization(any(), any());
+
+        var mockClient = mock(Client.class);
+
+        var callbackCount = new AtomicInteger(0);
+        var latch = new CountDownLatch(1);
+
+        Runnable callback = () -> {
+            callbackCount.incrementAndGet();
+            latch.countDown();
+        };
+
+        // Simulate scheduling failure by having the settings throw an exception when queried
+        // Throwing an exception should cause the poller to shutdown and mark itself as completed
+        var settingsMock = mock(ElasticInferenceServiceSettings.class);
+        when(settingsMock.isPeriodicAuthorizationEnabled()).thenThrow(new IllegalStateException("failing"));
+
+        var poller = new AuthorizationPoller(
+            new AuthorizationPoller.TaskFields(0, "abc", "abc", "abc", new TaskId("abc", 0), Map.of()),
+            createWithEmptySettings(taskQueue.getThreadPool()),
+            mockAuthHandler,
+            mock(Sender.class),
+            settingsMock,
+            mockRegistry,
+            mockClient,
+            callback
+        );
+        poller.init(mock(PersistentTasksService.class), mock(TaskManager.class), "id", 0);
+        poller.start();
+        taskQueue.runAllRunnableTasks();
+        latch.await(TimeValue.THIRTY_SECONDS.getSeconds(), TimeUnit.SECONDS);
+
+        assertThat(callbackCount.get(), is(1));
+        assertTrue(poller.isShutdown());
+        verify(mockClient, never()).execute(eq(StoreInferenceEndpointsAction.INSTANCE), any(), any());
+    }
 }