WIP: make GeoIp search project aware

Author: samxbr

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

@@ -39,6 +39,7 @@ final class ConfigDatabases implements Closeable {
     private final GeoIpCache cache;
     private final Path geoipConfigDir;
 
+//    private final ConcurrentMap<ProjectId, ConcurrentMap<String, DatabaseReaderLazyLoader>> configDatabasesByProject;
     private final ConcurrentMap<String, DatabaseReaderLazyLoader> configDatabases;
 
     ConfigDatabases(Environment environment, GeoIpCache cache) {

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

@@ -88,6 +88,13 @@ public final String getDatabaseType() throws IOException {
         return databaseType.get();
     }
 
+    /**
+     * Prepares the database for lookup by incrementing the usage count.
+     * If the usage count is already negative, it indicates that the database is being closed,
+     * and this method will return false to indicate that no lookup should be performed.
+     *
+     * @return true if the database is ready for lookup, false if it is being closed
+     */
     boolean preLookup() {
         return currentUsages.updateAndGet(current -> current < 0 ? current : current + 1) > 0;
     }

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

@@ -146,6 +146,7 @@ public void writeTo(StreamOutput out) throws IOException {
      * @return the geoip downloader's task state or null if there is not a state to read
      */
     @Nullable
+    @Deprecated(forRemoval = true)
     static EnterpriseGeoIpTaskState getEnterpriseGeoIpTaskState(ClusterState state) {
         PersistentTasksCustomMetadata.PersistentTask<?> task = getTaskWithId(state, EnterpriseGeoIpTask.ENTERPRISE_GEOIP_DOWNLOADER);
         return (task == null) ? null : (EnterpriseGeoIpTaskState) task.getState();

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

@@ -53,6 +53,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -106,9 +107,9 @@ public final class GeoIpDownloaderTaskExecutor extends PersistentTasksExecutor<G
     private volatile TimeValue pollInterval;
     private volatile boolean eagerDownload;
 
-    private final ConcurrentHashMap<ProjectId, AtomicBoolean> atLeastOneGeoipProcessorByProject = new ConcurrentHashMap<>();
-    private final ConcurrentHashMap<ProjectId, AtomicBoolean> taskIsBootstrappedByProject = new ConcurrentHashMap<>();
-    private final ConcurrentHashMap<ProjectId, AtomicReference<GeoIpDownloader>> tasks = new ConcurrentHashMap<>();
+    private final ConcurrentMap<ProjectId, AtomicBoolean> atLeastOneGeoipProcessorByProject = new ConcurrentHashMap<>();
+    private final ConcurrentMap<ProjectId, AtomicBoolean> taskIsBootstrappedByProject = new ConcurrentHashMap<>();
+    private final ConcurrentMap<ProjectId, AtomicReference<GeoIpDownloader>> tasks = new ConcurrentHashMap<>();
     private final ProjectResolver projectResolver;
 
     GeoIpDownloaderTaskExecutor(
@@ -409,7 +410,7 @@ private static boolean isForeachProcessorWithGeoipProcessor(Map<String, Object>
     private void startTask(Runnable onFailure) {
         assert projectResolver.getProjectId() != null : "projectId must be set before starting geoIp download task";
         persistentTasksService.sendStartRequest(
-            getTaskId(projectResolver.getProjectId()),
+            getTaskId(projectResolver.getProjectId(), projectResolver.supportsMultipleProjects()),
             GEOIP_DOWNLOADER,
             new GeoIpTaskParams(),
             MasterNodeRequest.INFINITE_MASTER_NODE_TIMEOUT,
@@ -438,7 +439,7 @@ private void stopTask(Runnable onFailure) {
             }
         );
         persistentTasksService.sendRemoveRequest(
-            getTaskId(projectId),
+            getTaskId(projectId, projectResolver.supportsMultipleProjects()),
             MasterNodeRequest.INFINITE_MASTER_NODE_TIMEOUT,
             ActionListener.runAfter(listener, () -> {
                 IndexAbstraction databasesAbstraction = clusterService.state()
@@ -471,7 +472,7 @@ public GeoIpDownloader getTask(ProjectId projectId) {
         return null;
     }
 
-    private String getTaskId(ProjectId projectId) {
-        return projectResolver.supportsMultipleProjects() ? projectId + "/" + GEOIP_DOWNLOADER : GEOIP_DOWNLOADER;
+    public static String getTaskId(ProjectId projectId, boolean supportsMultipleProjects) {
+        return supportsMultipleProjects ? projectId + "/" + GEOIP_DOWNLOADER : GEOIP_DOWNLOADER;
     }
 }

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

@@ -180,16 +180,19 @@ public static final class DatabaseVerifyingSupplier implements CheckedSupplier<I
         private final IpDatabaseProvider ipDatabaseProvider;
         private final String databaseFile;
         private final String databaseType;
+        private final ProjectId projectId;
 
-        public DatabaseVerifyingSupplier(IpDatabaseProvider ipDatabaseProvider, String databaseFile, String databaseType) {
+        public DatabaseVerifyingSupplier(IpDatabaseProvider ipDatabaseProvider, String databaseFile, String databaseType,
+                                         ProjectId projectId) {
             this.ipDatabaseProvider = ipDatabaseProvider;
             this.databaseFile = databaseFile;
             this.databaseType = databaseType;
+            this.projectId = projectId;
         }
 
         @Override
         public IpDatabase get() throws IOException {
-            IpDatabase loader = ipDatabaseProvider.getDatabase(databaseFile);
+            IpDatabase loader = ipDatabaseProvider.getDatabase(projectId, databaseFile);
             if (loader == null) {
                 return null;
             }
@@ -242,7 +245,7 @@ public Processor create(
             readBooleanProperty(type, processorTag, config, "download_database_on_pipeline_creation", true);
 
             final String databaseType;
-            try (IpDatabase ipDatabase = ipDatabaseProvider.getDatabase(databaseFile)) {
+            try (IpDatabase ipDatabase = ipDatabaseProvider.getDatabase(projectId, databaseFile)) {
                 if (ipDatabase == null) {
                     // It's possible that the database could be downloaded via the GeoipDownloader process and could become available
                     // at a later moment, so a processor impl is returned that tags documents instead. If a database cannot be sourced
@@ -302,8 +305,8 @@ public Processor create(
                 processorTag,
                 description,
                 ipField,
-                new DatabaseVerifyingSupplier(ipDatabaseProvider, databaseFile, databaseType),
-                () -> ipDatabaseProvider.isValid(databaseFile),
+                new DatabaseVerifyingSupplier(ipDatabaseProvider, databaseFile, databaseType, projectId),
+                () -> ipDatabaseProvider.isValid(projectId, databaseFile),
                 targetField,
                 ipDataLookup,
                 ignoreMissing,

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.VersionedNamedWriteable;
@@ -252,16 +253,31 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         }
     }
 
+//    /**
+//     * Retrieves the geoip downloader's task state from the cluster state. This may return null in some circumstances,
+//     * for example if the geoip downloader task hasn't been created yet (which it wouldn't be if it's disabled).
+//     *
+//     * @param state the cluster state to read the task state from
+//     * @return the geoip downloader's task state or null if there is not a state to read
+//     */
+//    @Nullable
+//    @Deprecated(forRemoval = true)
+//    static GeoIpTaskState getGeoIpTaskState(ClusterState state) {
+//        PersistentTasksCustomMetadata.PersistentTask<?> task = getTaskWithId(state, GeoIpDownloader.GEOIP_DOWNLOADER);
+//        return (task == null) ? null : (GeoIpTaskState) task.getState();
+//    }
+
     /**
-     * Retrieves the geoip downloader's task state from the cluster state. This may return null in some circumstances,
+     * Retrieves the geoip downloader's task state from the project metadata. This may return null in some circumstances,
      * for example if the geoip downloader task hasn't been created yet (which it wouldn't be if it's disabled).
      *
-     * @param state the cluster state to read the task state from
+     * @param projectMetadata the project metatdata to read the task state from.
+     * @param taskId the task ID of the geoip downloader task to read the state for.
      * @return the geoip downloader's task state or null if there is not a state to read
      */
     @Nullable
-    static GeoIpTaskState getGeoIpTaskState(ClusterState state) {
-        PersistentTasksCustomMetadata.PersistentTask<?> task = getTaskWithId(state, GeoIpDownloader.GEOIP_DOWNLOADER);
+    static GeoIpTaskState getGeoIpTaskState(ProjectMetadata projectMetadata, String taskId) {
+        PersistentTasksCustomMetadata.PersistentTask<?> task = getTaskWithId(projectMetadata, taskId);
         return (task == null) ? null : (GeoIpTaskState) task.getState();
     }
 

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

@@ -139,7 +139,7 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
     public Collection<?> createComponents(PluginServices services) {
         try {
             String nodeId = services.nodeEnvironment().nodeId();
-            databaseRegistry.get().initialize(nodeId, services.resourceWatcherService(), ingestService.get());
+            databaseRegistry.get().initialize(nodeId, services.resourceWatcherService(), ingestService.get(), services.projectResolver());
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }

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

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.ingest.geoip;
 
+import org.elasticsearch.cluster.metadata.ProjectId;
+
 /**
  * Provides construction and initialization logic for {@link IpDatabase} instances.
  */
@@ -20,15 +22,17 @@ public interface IpDatabaseProvider {
      * Verifying database expiration is left to each provider implementation to determine. A return value of <code>false</code> does not
      * preclude the possibility of a provider returning <code>true</code> in the future.
      *
+     * @param projectId projectId to look for database.
      * @param name the name of the database to provide.
      * @return <code>false</code> IFF the requested database file is expired,
      *         <code>true</code> for all other cases (including unknown file name, file missing, wrong database type, etc).
      */
-    Boolean isValid(String name);
+    Boolean isValid(ProjectId projectId, String name);
 
     /**
+     * @param projectId projectId to look for database.
      * @param name the name of the database to provide.
      * @return a ready-to-use database instance, or <code>null</code> if no database could be loaded.
      */
-    IpDatabase getDatabase(String name);
+    IpDatabase getDatabase(ProjectId projectId, String name);
 }

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.nodes.TransportNodesAction;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -77,15 +78,16 @@ protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throw
 
     @Override
     protected NodeResponse nodeOperation(NodeRequest request, Task task) {
-        GeoIpDownloader geoIpTask = geoIpDownloaderTaskExecutor.getTask(projectResolver.getProjectId());
+        ProjectId projectId = projectResolver.getProjectId();
+        GeoIpDownloader geoIpTask = geoIpDownloaderTaskExecutor.getTask(projectId);
         GeoIpDownloaderStats downloaderStats = geoIpTask == null || geoIpTask.getStatus() == null ? null : geoIpTask.getStatus();
         CacheStats cacheStats = registry.getCacheStats();
         return new NodeResponse(
             transportService.getLocalNode(),
             downloaderStats,
             cacheStats,
-            registry.getAvailableDatabases(),
-            registry.getFilesInTemp(),
+            registry.getAvailableDatabases(projectId),
+            registry.getFilesInTemp(projectId),
             registry.getConfigDatabases()
         );
     }