[WIP] ES-11278 Add option for byte-based capacity for geoip cache

modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/DatabaseNodeServiceIT.java

@@ -9,13 +9,11 @@
 
 package org.elasticsearch.ingest.geoip;
 
-import com.maxmind.geoip2.model.CountryResponse;
-import com.maxmind.geoip2.record.Country;
-
 import org.elasticsearch.action.admin.indices.flush.FlushRequest;
 import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.common.hash.MessageDigests;
+import org.elasticsearch.ingest.geoip.MaxmindIpDataLookups.CacheableCountryResponse;
 import org.elasticsearch.xcontent.XContentType;
 
 import java.io.BufferedInputStream;
@@ -96,11 +94,9 @@ private void assertValidDatabase(DatabaseNodeService databaseNodeService, String
         IpDatabase database = databaseNodeService.getDatabase(databaseFileName);
         assertNotNull(database);
         assertThat(database.getDatabaseType(), equalTo(databaseType));
-        CountryResponse countryResponse = database.getResponse("89.160.20.128", GeoIpTestUtils::getCountry);
+        CacheableCountryResponse countryResponse = database.getResponse("89.160.20.128", GeoIpTestUtils::getCountry);
         assertNotNull(countryResponse);
-        Country country = countryResponse.getCountry();
-        assertNotNull(country);
-        assertThat(country.getName(), equalTo("Sweden"));
+        assertThat(countryResponse.countryName(), equalTo("Sweden"));
     }
 
     /*

modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/ReloadingDatabasesWhilePerformingGeoLookupsIT.java

@@ -195,7 +195,7 @@ public void test() throws Exception {
 
     private static DatabaseNodeService createRegistry(Path geoIpConfigDir, Path geoIpTmpDir, ClusterService clusterService)
         throws IOException {
-        GeoIpCache cache = new GeoIpCache(0);
+        GeoIpCache cache = GeoIpCache.createGeoIpCacheWithMaxCount(0);
         ConfigDatabases configDatabases = new ConfigDatabases(geoIpConfigDir, cache);
         copyDefaultDatabases(geoIpConfigDir, configDatabases);
         DatabaseNodeService databaseNodeService = new DatabaseNodeService(

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

@@ -105,7 +105,10 @@ int current() {
 
     @Override
     @Nullable
-    public <RESPONSE> RESPONSE getResponse(String ipAddress, CheckedBiFunction<Reader, String, RESPONSE, Exception> responseProvider) {
+    public <RESPONSE extends GeoIpCache.CacheableValue> RESPONSE getResponse(
+        String ipAddress,
+        CheckedBiFunction<Reader, String, RESPONSE, Exception> responseProvider
+    ) {
         return cache.putIfAbsent(ipAddress, cachedDatabasePathToString, ip -> {
             try {
                 return responseProvider.apply(get(), ipAddress);

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

@@ -10,8 +10,10 @@
 
 import com.maxmind.db.NodeCache;
 
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.cache.Cache;
 import org.elasticsearch.common.cache.CacheBuilder;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.ingest.geoip.stats.CacheStats;
 
@@ -28,44 +30,69 @@
  */
 public final class GeoIpCache {
 
+    public interface CacheableValue {
+
+        // TODO PETE: Remove this default implementation and implement in all implementing classes instead
+        default long sizeInBytes() {
+            return 0;
+        }
+    }
+
+    static <V> GeoIpCache createGeoIpCacheWithMaxCount(long maxSize) {
+        if (maxSize < 0) {
+            throw new IllegalArgumentException("geoip max cache size must be 0 or greater");
+        }
+        return new GeoIpCache(System::nanoTime, CacheBuilder.<CacheKey, CacheableValue>builder().setMaximumWeight(maxSize).build());
+    }
+
+    static GeoIpCache createGeoIpCacheWithMaxBytes(ByteSizeValue maxByteSize) {
+        if (maxByteSize.getBytes() < 0) {
+            throw new IllegalArgumentException("geoip max cache size in bytes must be 0 or greater");
+        }
+        return new GeoIpCache(
+            System::nanoTime,
+            CacheBuilder.<CacheKey, CacheableValue>builder()
+                .setMaximumWeight(maxByteSize.getBytes())
+                .weigher((key, value) -> key.sizeInBytes() + value.sizeInBytes())
+                .build()
+        );
+    }
+
+    // package private for testing
+    static GeoIpCache createGeoIpCacheWithMaxCountAndCustomTimeProvider(long maxSize, LongSupplier relativeNanoTimeProvider) {
+        return new GeoIpCache(relativeNanoTimeProvider, CacheBuilder.<CacheKey, CacheableValue>builder().setMaximumWeight(maxSize).build());
+    }
+
     /**
      * Internal-only sentinel object for recording that a result from the geoip database was null (i.e. there was no result). By caching
      * this no-result we can distinguish between something not being in the cache because we haven't searched for that data yet, versus
      * something not being in the cache because the data doesn't exist in the database.
      */
     // visible for testing
-    static final Object NO_RESULT = new Object() {
+    static final CacheableValue NO_RESULT = new CacheableValue() {
         @Override
         public String toString() {
             return "NO_RESULT";
         }
     };
 
     private final LongSupplier relativeNanoTimeProvider;
-    private final Cache<CacheKey, Object> cache;
+    private final Cache<CacheKey, CacheableValue> cache;
     private final AtomicLong hitsTimeInNanos = new AtomicLong(0);
     private final AtomicLong missesTimeInNanos = new AtomicLong(0);
 
-    // package private for testing
-    GeoIpCache(long maxSize, LongSupplier relativeNanoTimeProvider) {
-        if (maxSize < 0) {
-            throw new IllegalArgumentException("geoip max cache size must be 0 or greater");
-        }
+    private GeoIpCache(LongSupplier relativeNanoTimeProvider, Cache<CacheKey, CacheableValue> cache) {
         this.relativeNanoTimeProvider = relativeNanoTimeProvider;
-        this.cache = CacheBuilder.<CacheKey, Object>builder().setMaximumWeight(maxSize).build();
-    }
-
-    GeoIpCache(long maxSize) {
-        this(maxSize, System::nanoTime);
+        this.cache = cache;
     }
 
     @SuppressWarnings("unchecked")
-    <RESPONSE> RESPONSE putIfAbsent(String ip, String databasePath, Function<String, RESPONSE> retrieveFunction) {
+    <RESPONSE extends CacheableValue> RESPONSE putIfAbsent(String ip, String databasePath, Function<String, RESPONSE> retrieveFunction) {
         // can't use cache.computeIfAbsent due to the elevated permissions for the jackson (run via the cache loader)
         CacheKey cacheKey = new CacheKey(ip, databasePath);
         long cacheStart = relativeNanoTimeProvider.getAsLong();
         // intentionally non-locking for simplicity...it's OK if we re-put the same key/value in the cache during a race condition.
-        Object response = cache.get(cacheKey);
+        CacheableValue response = cache.get(cacheKey);
         long cacheRequestTime = relativeNanoTimeProvider.getAsLong() - cacheStart;
 
         // populate the cache for this key, if necessary
@@ -135,5 +162,17 @@ public CacheStats getCacheStats() {
      * path is needed to be included in the cache key. For example, if we only used the IP address as the key the City and ASN the same
      * IP may be in both with different values and we need to cache both.
      */
-    private record CacheKey(String ip, String databasePath) {}
+    private record CacheKey(String ip, String databasePath) {
+
+        private static final long BASE_BYTES = RamUsageEstimator.shallowSizeOfInstance(CacheKey.class);
+
+        private long sizeInBytes() {
+            return keySizeInBytes(ip, databasePath);
+        }
+    }
+
+    // visible for testing
+    static long keySizeInBytes(String ip, String databasePath) {
+        return CacheKey.BASE_BYTES + RamUsageEstimator.sizeOf(ip) + RamUsageEstimator.sizeOf(databasePath);
+    }
 }

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

@@ -24,6 +24,8 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.SettingsFilter;
 import org.elasticsearch.common.settings.SettingsModule;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.indices.SystemIndexDescriptor;
 import org.elasticsearch.ingest.EnterpriseGeoIpTask.EnterpriseGeoIpTaskParams;
@@ -86,7 +88,17 @@ public class IngestGeoIpPlugin extends Plugin
         PersistentTaskPlugin,
         ActionPlugin,
         ReloadablePlugin {
-    public static final Setting<Long> CACHE_SIZE = Setting.longSetting("ingest.geoip.cache_size", 1000, 0, Setting.Property.NodeScope);
+    private static final Setting<Long> CACHE_SIZE_COUNT = Setting.longSetting(
+        "ingest.geoip.cache_size",
+        1000,
+        0,
+        Setting.Property.NodeScope
+    );
+    private static final Setting<ByteSizeValue> CACHE_SIZE_BYTES = Setting.byteSizeSetting(
+        "ingest.geoip.cache_memory_size",
+        ByteSizeValue.MINUS_ONE,
+        Setting.Property.NodeScope
+    );
     private static final int GEOIP_INDEX_MAPPINGS_VERSION = 1;
     /**
      * No longer used for determining the age of mappings, but system index descriptor
@@ -105,7 +117,8 @@ public class IngestGeoIpPlugin extends Plugin
     @Override
     public List<Setting<?>> getSettings() {
         return List.of(
-            CACHE_SIZE,
+            CACHE_SIZE_COUNT,
+            CACHE_SIZE_BYTES,
             GeoIpDownloaderTaskExecutor.EAGER_DOWNLOAD_SETTING,
             GeoIpDownloaderTaskExecutor.ENABLED_SETTING,
             GeoIpDownloader.ENDPOINT_SETTING,
@@ -119,8 +132,7 @@ public List<Setting<?>> getSettings() {
     public Map<String, Processor.Factory> getProcessors(Processor.Parameters parameters) {
         ingestService.set(parameters.ingestService);
 
-        long cacheSize = CACHE_SIZE.get(parameters.env.settings());
-        GeoIpCache geoIpCache = new GeoIpCache(cacheSize);
+        GeoIpCache geoIpCache = createGeoIpCache(parameters.env.settings());
         DatabaseNodeService registry = new DatabaseNodeService(
             parameters.env,
             parameters.client,
@@ -135,6 +147,30 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
         );
     }
 
+    private static GeoIpCache createGeoIpCache(Settings settings) {
+        if (settings.hasValue(CACHE_SIZE_BYTES.getKey())) {
+            if (settings.hasValue(CACHE_SIZE_COUNT.getKey())) {
+                // Both CACHE_SIZE_COUNT and CACHE_SIZE_BYTES are set, which is an error:
+                throw new IllegalArgumentException(
+                    Strings.format(
+                        "Both %s and %s are set: "
+                            + "please use either %s to set a size based on count or %s to set a size based on bytes of memory",
+                        CACHE_SIZE_COUNT.getKey(),
+                        CACHE_SIZE_BYTES.getKey(),
+                        CACHE_SIZE_COUNT.getKey(),
+                        CACHE_SIZE_BYTES.getKey()
+                    )
+                );
+            } else {
+                // Only CACHE_SIZE_BYTES is set, so use that:
+                return GeoIpCache.createGeoIpCacheWithMaxBytes(CACHE_SIZE_BYTES.get(settings));
+            }
+        } else {
+            // CACHE_SIZE_BYTES is not set, so use either the explicit or default value of CACHE_SIZE_COUNT:
+            return GeoIpCache.createGeoIpCacheWithMaxCount(CACHE_SIZE_COUNT.get(settings));
+        }
+    }
+
     @Override
     public Collection<?> createComponents(PluginServices services) {
         try {

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

@@ -34,5 +34,5 @@ interface IpDataLookup {
      * as a network for which the  record applies. Having a helper record prevents each individual response record from needing to
      * track these bits of information.
      */
-    record Result<T>(T result, String ip, String network) {}
+    record Result<T extends GeoIpCache.CacheableValue>(T result, String ip, String network) implements GeoIpCache.CacheableValue {}
 }

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

@@ -33,10 +33,12 @@ public interface IpDatabase extends AutoCloseable {
      * @param ipAddress the address to lookup
      * @param responseProvider a method for extracting a response from a {@link Reader}, usually this will be a method reference
      * @return a possibly-null response
-     * @param <RESPONSE> the type of response that will be returned
      */
     @Nullable
-    <RESPONSE> RESPONSE getResponse(String ipAddress, CheckedBiFunction<Reader, String, RESPONSE, Exception> responseProvider);
+    <RESPONSE extends GeoIpCache.CacheableValue> RESPONSE getResponse(
+        String ipAddress,
+        CheckedBiFunction<Reader, String, RESPONSE, Exception> responseProvider
+    );
 
     /**
      * Releases the current database object. Called after processing a single document. Databases should be closed or returned to a

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

@@ -191,7 +191,7 @@ public record AsnResult(
         String domain,
         String name,
         @Nullable String type // not present in the free asn database
-    ) {
+    ) implements GeoIpCache.CacheableValue {
         @SuppressWarnings("checkstyle:RedundantModifier")
         @MaxMindDbConstructor
         public AsnResult(
@@ -210,20 +210,14 @@ public record CountryResult(
         @MaxMindDbParameter(name = "continent_name") String continentName,
         @MaxMindDbParameter(name = "country") String country,
         @MaxMindDbParameter(name = "country_name") String countryName
-    ) {
+    ) implements GeoIpCache.CacheableValue {
         @MaxMindDbConstructor
         public CountryResult {}
     }
 
-    public record GeolocationResult(
-        String city,
-        String country,
-        Double lat,
-        Double lng,
-        String postalCode,
-        String region,
-        String timezone
-    ) {
+    public record GeolocationResult(String city, String country, Double lat, Double lng, String postalCode, String region, String timezone)
+        implements
+            GeoIpCache.CacheableValue {
         @SuppressWarnings("checkstyle:RedundantModifier")
         @MaxMindDbConstructor
         public GeolocationResult(
@@ -241,7 +235,9 @@ public GeolocationResult(
         }
     }
 
-    public record PrivacyDetectionResult(Boolean hosting, Boolean proxy, Boolean relay, String service, Boolean tor, Boolean vpn) {
+    public record PrivacyDetectionResult(Boolean hosting, Boolean proxy, Boolean relay, String service, Boolean tor, Boolean vpn)
+        implements
+            GeoIpCache.CacheableValue {
         @SuppressWarnings("checkstyle:RedundantModifier")
         @MaxMindDbConstructor
         public PrivacyDetectionResult(
@@ -466,7 +462,7 @@ protected Map<String, Object> transform(final Result<PrivacyDetectionResult> res
      *
      * @param <RESPONSE> the record type that will be wrapped and returned
      */
-    private abstract static class AbstractBase<RESPONSE> implements IpDataLookup {
+    private abstract static class AbstractBase<RESPONSE extends GeoIpCache.CacheableValue> implements IpDataLookup {
 
         protected final Set<Database.Property> properties;
         protected final Class<RESPONSE> clazz;