Make GeoIp downloader multi-project aware

modules/ingest-geoip/qa/multi-project/build.gradle

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+apply plugin: 'elasticsearch.internal-java-rest-test'
+
+dependencies {
+  javaRestTestImplementation project(':modules:ingest-geoip')
+  javaRestTestImplementation project(':test:external-modules:test-multi-project')
+  javaRestTestImplementation project(':test:fixtures:geoip-fixture')
+
+  clusterModules project(':modules:ingest-geoip')
+  clusterModules project(':modules:reindex')  // needed for database cleanup
+  clusterModules project(':test:external-modules:test-multi-project')
+}
+
+tasks.withType(Test).configureEach {
+  it.systemProperty "tests.multi_project.enabled", true
+}

modules/ingest-geoip/qa/multi-project/src/javaRestTest/java/geoip/GeoIpMultiProjectIT.java

@@ -0,0 +1,150 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package geoip;
+
+import fixture.geoip.GeoIpHttpFixture;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.core.Booleans;
+import org.elasticsearch.ingest.geoip.GeoIpDownloader;
+import org.elasticsearch.ingest.geoip.GeoIpDownloaderTaskExecutor;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.elasticsearch.test.rest.ObjectPath;
+import org.junit.ClassRule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class GeoIpMultiProjectIT extends ESRestTestCase {
+    // default true
+    private static final boolean useFixture = Booleans.parseBoolean(System.getProperty("geoip_use_service", "false")) == false;
+
+    public static final GeoIpHttpFixture fixture = new GeoIpHttpFixture(useFixture);
+
+    public static final ElasticsearchCluster cluster = ElasticsearchCluster.local()
+        .module("ingest-geoip")
+        .module("reindex")  // for database cleanup
+        .module("test-multi-project")
+        .setting("test.multi_project.enabled", "true")
+        .setting(GeoIpDownloaderTaskExecutor.ENABLED_SETTING.getKey(), "true")
+        .setting(GeoIpDownloader.ENDPOINT_SETTING.getKey(), fixture::getAddress, (k) -> useFixture)
+        .build();
+
+    @ClassRule
+    public static TestRule ruleChain = RuleChain.outerRule(fixture).around(cluster);
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+
+    @Override
+    protected boolean shouldConfigureProjects() {
+        return false;
+    }
+
+    public void testGeoIpDownloader() throws Exception {
+        String project1 = randomUniqueProjectId().id();
+        String project2 = randomUniqueProjectId().id();
+        createProject(project1);
+        createProject(project2);
+
+        // download databases for project1
+        putGeoIpPipeline(project1);
+        assertBusy(() -> assertDatabases(project1, true), 30, TimeUnit.SECONDS);
+        assertBusy(() -> assertDatabases(project2, false), 30, TimeUnit.SECONDS);
+
+        // download databases for project2
+        putGeoIpPipeline(project2);
+        assertBusy(() -> assertDatabases(project2, true), 30, TimeUnit.SECONDS);
+    }
+
+    private void putGeoIpPipeline(String projectId) throws IOException {
+        Request putPipelineRequest = new Request("PUT", "/_ingest/pipeline/geoip-pipeline");
+        putPipelineRequest.setJsonEntity("""
+            {
+              "processors" : [
+                {
+                  "geoip" : {
+                    "field" : "ip",
+                    "target_field" : "geo",
+                    "database_file" : "GeoLite2-Country.mmdb"
+                  }
+                }
+              ]
+            }
+            """);
+        setRequestProjectId(projectId, putPipelineRequest);
+        assertOK(client().performRequest(putPipelineRequest));
+    }
+
+    private static Request setRequestProjectId(String projectId, Request request) {
+        RequestOptions.Builder options = request.getOptions().toBuilder();
+        options.removeHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER);
+        options.addHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER, projectId);
+        request.setOptions(options);
+        return request;
+    }
+
+    @SuppressWarnings("unchecked")
+    private void assertDatabases(String projectId, boolean shouldDownload) throws IOException {
+        Request getTaskState = new Request("GET", "/_cluster/state");
+        setRequestProjectId(projectId, getTaskState);
+
+        ObjectPath state = ObjectPath.createFromResponse(assertOK(client().performRequest(getTaskState)));
+
+        List<Map<String, ?>> tasks = state.evaluate("metadata.persistent_tasks.tasks");
+        // Short-circuit to avoid using steams if the list is empty
+        if (tasks.isEmpty()) {
+            fail("persistent tasks list is empty, expected at least one task for geoip-downloader");
+        }
+
+        // verify project task id
+        Set<Map<String, ?>> id = tasks.stream()
+            .filter(task -> String.format("%s/geoip-downloader", projectId).equals(task.get("id")))
+            .collect(Collectors.toSet());
+        assertThat(id.size(), equalTo(1));
+
+        // verify database download
+        Map<String, Object> databases = (Map<String, Object>) tasks.stream().map(task -> {
+            try {
+                return ObjectPath.evaluate(task, "task.geoip-downloader.state.databases");
+            } catch (IOException e) {
+                return null;
+            }
+        }).filter(Objects::nonNull).findFirst().orElse(null);
+
+        if (shouldDownload) {
+            // verify database downloaded
+            assertNotNull(databases);
+            for (String name : List.of("GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb")) {
+                Object database = databases.get(name);
+                assertNotNull(database);
+                assertNotNull(ObjectPath.evaluate(database, "md5"));
+            }
+        } else {
+            // verify database not downloaded
+            assertNull(databases);
+        }
+
+    }
+}

modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/GeoIpDownloader.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.hash.MessageDigests;
 import org.elasticsearch.common.settings.Setting;
@@ -95,6 +96,8 @@ public class GeoIpDownloader extends AllocatedPersistentTask {
      */
     private final Supplier<Boolean> atLeastOneGeoipProcessorSupplier;
 
+    private final ProjectId projectId;
+
     GeoIpDownloader(
         Client client,
         HttpClient httpClient,
@@ -109,17 +112,19 @@ public class GeoIpDownloader extends AllocatedPersistentTask {
         Map<String, String> headers,
         Supplier<TimeValue> pollIntervalSupplier,
         Supplier<Boolean> eagerDownloadSupplier,
-        Supplier<Boolean> atLeastOneGeoipProcessorSupplier
+        Supplier<Boolean> atLeastOneGeoipProcessorSupplier,
+        ProjectId projectId
     ) {
         super(id, type, action, description, parentTask, headers);
-        this.client = client;
+        this.client = client.projectClient(projectId);
         this.httpClient = httpClient;
         this.clusterService = clusterService;
         this.threadPool = threadPool;
         this.endpoint = ENDPOINT_SETTING.get(settings);
         this.pollIntervalSupplier = pollIntervalSupplier;
         this.eagerDownloadSupplier = eagerDownloadSupplier;
         this.atLeastOneGeoipProcessorSupplier = atLeastOneGeoipProcessorSupplier;
+        this.projectId = projectId;
     }
 
     void setState(GeoIpTaskState state) {
@@ -134,16 +139,17 @@ void setState(GeoIpTaskState state) {
     // visible for testing
     void updateDatabases() throws IOException {
         var clusterState = clusterService.state();
-        var geoipIndex = clusterState.getMetadata().getProject().getIndicesLookup().get(GeoIpDownloader.DATABASES_INDEX);
+        var geoipIndex = clusterState.getMetadata().getProject(projectId).getIndicesLookup().get(GeoIpDownloader.DATABASES_INDEX);
         if (geoipIndex != null) {
             logger.trace("The {} index is not null", GeoIpDownloader.DATABASES_INDEX);
-            if (clusterState.getRoutingTable().index(geoipIndex.getWriteIndex()).allPrimaryShardsActive() == false) {
+            if (clusterState.routingTable(projectId).index(geoipIndex.getWriteIndex()).allPrimaryShardsActive() == false) {
                 logger.debug(
                     "Not updating geoip database because not all primary shards of the [" + DATABASES_INDEX + "] index are active."
                 );
                 return;
             }
-            var blockException = clusterState.blocks().indexBlockedException(ClusterBlockLevel.WRITE, geoipIndex.getWriteIndex().getName());
+            var blockException = clusterState.blocks()
+                .indexBlockedException(projectId, ClusterBlockLevel.WRITE, geoipIndex.getWriteIndex().getName());
             if (blockException != null) {
                 logger.debug(
                     "Not updating geoip database because there is a write block on the " + geoipIndex.getWriteIndex().getName() + " index",
@@ -196,7 +202,7 @@ private void processDatabase(final String name, final String md5, final String u
             updateTimestamp(name, metadata);
             return;
         }
-        logger.debug("downloading geoip database [{}]", name);
+        logger.debug("downloading geoip database [{}] for project [{}]", name, projectId);
         long start = System.currentTimeMillis();
         try (InputStream is = httpClient.get(url)) {
             int firstChunk = metadata.lastChunk() + 1; // if there is no metadata, then Metadata.EMPTY.lastChunk() + 1 = 0
@@ -205,12 +211,12 @@ private void processDatabase(final String name, final String md5, final String u
                 state = state.put(name, new Metadata(start, firstChunk, lastChunk - 1, md5, start));
                 updateTaskState();
                 stats = stats.successfulDownload(System.currentTimeMillis() - start).databasesCount(state.getDatabases().size());
-                logger.info("successfully downloaded geoip database [{}]", name);
+                logger.info("successfully downloaded geoip database [{}] for project [{}]", name, projectId);
                 deleteOldChunks(name, firstChunk);
             }
         } catch (Exception e) {
             stats = stats.failedDownload();
-            logger.error(() -> "error downloading geoip database [" + name + "]", e);
+            logger.error(() -> "error downloading geoip database [" + name + "] for project [" + projectId + "]", e);
         }
     }
 
@@ -230,15 +236,15 @@ void deleteOldChunks(String name, int firstChunk) {
 
     // visible for testing
     protected void updateTimestamp(String name, Metadata old) {
-        logger.debug("geoip database [{}] is up to date, updated timestamp", name);
+        logger.debug("geoip database [{}] is up to date for project [{}], updated timestamp", name, projectId);
         state = state.put(name, new Metadata(old.lastUpdate(), old.firstChunk(), old.lastChunk(), old.md5(), System.currentTimeMillis()));
         stats = stats.skippedDownload();
         updateTaskState();
     }
 
     void updateTaskState() {
         PlainActionFuture<PersistentTask<?>> future = new PlainActionFuture<>();
-        updatePersistentTaskState(state, future);
+        updateProjectPersistentTaskState(projectId, state, future);
         state = ((GeoIpTaskState) future.actionGet().getState());
     }
 
@@ -360,5 +366,4 @@ private void scheduleNextRun(TimeValue time) {
             scheduled = threadPool.schedule(this::runDownloader, time, threadPool.generic());
         }
     }
-
 }