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

This adds the option to have a geoip cache with a byte-based rather
than a count-based capacity.

Author: PeteGillinElastic

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

@@ -96,7 +96,7 @@ 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);
+        CountryResponse countryResponse = database.getResponse("89.160.20.128", GeoIpTestUtils::getCountry).result();
         assertNotNull(countryResponse);
         Country country = countryResponse.getCountry();
         assertNotNull(country);

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,71 @@
  */
 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());
+    }
+
+    // TODO PETE: Add tests for this
+    // TODO PETE: Make plugin use this instead of the other factory method when the settings require it
+    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 +164,12 @@ 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);
+
+        public long sizeInBytes() {
+            return BASE_BYTES + RamUsageEstimator.sizeOf(ip) + RamUsageEstimator.sizeOf(databasePath);
+        }
+    }
 }

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

@@ -120,7 +120,7 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
         ingestService.set(parameters.ingestService);
 
         long cacheSize = CACHE_SIZE.get(parameters.env.settings());
-        GeoIpCache geoIpCache = new GeoIpCache(cacheSize);
+        GeoIpCache geoIpCache = GeoIpCache.createGeoIpCacheWithMaxCount(cacheSize);
         DatabaseNodeService registry = new DatabaseNodeService(
             parameters.env,
             parameters.client,

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>(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/MaxmindIpDataLookups.java

@@ -109,7 +109,23 @@ static Function<Set<Database.Property>, IpDataLookup> getMaxmindLookup(final Dat
         };
     }
 
-    static class AnonymousIp extends AbstractBase<AnonymousIpResponse, AnonymousIpResponse> {
+    static class CacheableAnonymousIpResponse extends AnonymousIpResponse implements GeoIpCache.CacheableValue {
+
+        CacheableAnonymousIpResponse(AnonymousIpResponse response) {
+            super(
+                response.getIpAddress(),
+                response.isAnonymous(),
+                response.isAnonymousVpn(),
+                response.isHostingProvider(),
+                response.isPublicProxy(),
+                response.isResidentialProxy(),
+                response.isTorExitNode(),
+                response.getNetwork()
+            );
+        }
+    }
+
+    static class AnonymousIp extends AbstractBase<AnonymousIpResponse, CacheableAnonymousIpResponse> {
         AnonymousIp(final Set<Database.Property> properties) {
             super(
                 properties,
@@ -119,12 +135,12 @@ static class AnonymousIp extends AbstractBase<AnonymousIpResponse, AnonymousIpRe
         }
 
         @Override
-        protected AnonymousIpResponse cacheableRecord(AnonymousIpResponse response) {
-            return response;
+        protected CacheableAnonymousIpResponse cacheableRecord(AnonymousIpResponse response) {
+            return new CacheableAnonymousIpResponse(response);
         }
 
         @Override
-        protected Map<String, Object> transform(final AnonymousIpResponse response) {
+        protected Map<String, Object> transform(final CacheableAnonymousIpResponse response) {
             boolean isHostingProvider = response.isHostingProvider();
             boolean isTorExitNode = response.isTorExitNode();
             boolean isAnonymousVpn = response.isAnonymousVpn();
@@ -160,18 +176,30 @@ protected Map<String, Object> transform(final AnonymousIpResponse response) {
         }
     }
 
-    static class Asn extends AbstractBase<AsnResponse, AsnResponse> {
+    static class CacheableAsnResponse extends AsnResponse implements GeoIpCache.CacheableValue {
+
+        CacheableAsnResponse(AsnResponse response) {
+            super(
+                response.getAutonomousSystemNumber(),
+                response.getAutonomousSystemOrganization(),
+                response.getIpAddress(),
+                response.getNetwork()
+            );
+        }
+    }
+
+    static class Asn extends AbstractBase<AsnResponse, CacheableAsnResponse> {
         Asn(Set<Database.Property> properties) {
             super(properties, AsnResponse.class, (response, ipAddress, network, locales) -> new AsnResponse(response, ipAddress, network));
         }
 
         @Override
-        protected AsnResponse cacheableRecord(AsnResponse response) {
-            return response;
+        protected CacheableAsnResponse cacheableRecord(AsnResponse response) {
+            return new CacheableAsnResponse(response);
         }
 
         @Override
-        protected Map<String, Object> transform(final AsnResponse response) {
+        protected Map<String, Object> transform(final CacheableAsnResponse response) {
             Long asn = response.getAutonomousSystemNumber();
             String organizationName = response.getAutonomousSystemOrganization();
             Network network = response.getNetwork();
@@ -354,7 +382,14 @@ protected Map<String, Object> transform(final Result<CacheableCityResponse> resu
         }
     }
 
-    static class ConnectionType extends AbstractBase<ConnectionTypeResponse, ConnectionTypeResponse> {
+    static class CacheableConnectionTypeResponse extends ConnectionTypeResponse implements GeoIpCache.CacheableValue {
+
+        CacheableConnectionTypeResponse(ConnectionTypeResponse response) {
+            super(response.getConnectionType(), response.getIpAddress(), response.getNetwork());
+        }
+    }
+
+    static class ConnectionType extends AbstractBase<ConnectionTypeResponse, CacheableConnectionTypeResponse> {
         ConnectionType(final Set<Database.Property> properties) {
             super(
                 properties,
@@ -364,12 +399,12 @@ static class ConnectionType extends AbstractBase<ConnectionTypeResponse, Connect
         }
 
         @Override
-        protected ConnectionTypeResponse cacheableRecord(ConnectionTypeResponse response) {
-            return response;
+        protected CacheableConnectionTypeResponse cacheableRecord(ConnectionTypeResponse response) {
+            return new CacheableConnectionTypeResponse(response);
         }
 
         @Override
-        protected Map<String, Object> transform(final ConnectionTypeResponse response) {
+        protected Map<String, Object> transform(final CacheableConnectionTypeResponse response) {
             ConnectionTypeResponse.ConnectionType connectionType = response.getConnectionType();
 
             Map<String, Object> data = new HashMap<>();
@@ -479,7 +514,14 @@ protected Map<String, Object> transform(final Result<CacheableCountryResponse> r
         }
     }
 
-    static class Domain extends AbstractBase<DomainResponse, DomainResponse> {
+    static class CacheableDomainResponse extends DomainResponse implements GeoIpCache.CacheableValue {
+
+        CacheableDomainResponse(DomainResponse response) {
+            super(response.getDomain(), response.getIpAddress(), response.getNetwork());
+        }
+    }
+
+    static class Domain extends AbstractBase<DomainResponse, CacheableDomainResponse> {
         Domain(final Set<Database.Property> properties) {
             super(
                 properties,
@@ -489,12 +531,12 @@ static class Domain extends AbstractBase<DomainResponse, DomainResponse> {
         }
 
         @Override
-        protected DomainResponse cacheableRecord(DomainResponse response) {
-            return response;
+        protected CacheableDomainResponse cacheableRecord(DomainResponse response) {
+            return new CacheableDomainResponse(response);
         }
 
         @Override
-        protected Map<String, Object> transform(final DomainResponse response) {
+        protected Map<String, Object> transform(final CacheableDomainResponse response) {
             String domain = response.getDomain();
 
             Map<String, Object> data = new HashMap<>();
@@ -784,18 +826,25 @@ protected Map<String, Object> transform(final Result<CacheableEnterpriseResponse
         }
     }
 
-    static class Isp extends AbstractBase<IspResponse, IspResponse> {
+    static class CacheableIspResponse extends IspResponse implements GeoIpCache.CacheableValue {
+
+        CacheableIspResponse(IspResponse response) {
+            super(response, response.getIpAddress(), response.getNetwork());
+        }
+    }
+
+    static class Isp extends AbstractBase<IspResponse, CacheableIspResponse> {
         Isp(final Set<Database.Property> properties) {
             super(properties, IspResponse.class, (response, ipAddress, network, locales) -> new IspResponse(response, ipAddress, network));
         }
 
         @Override
-        protected IspResponse cacheableRecord(IspResponse response) {
-            return response;
+        protected CacheableIspResponse cacheableRecord(IspResponse response) {
+            return new CacheableIspResponse(response);
         }
 
         @Override
-        protected Map<String, Object> transform(final IspResponse response) {
+        protected Map<String, Object> transform(final CacheableIspResponse response) {
             String isp = response.getIsp();
             String ispOrganization = response.getOrganization();
             String mobileNetworkCode = response.getMobileNetworkCode();
@@ -867,7 +916,9 @@ private interface ResponseBuilder<RESPONSE extends AbstractResponse> {
      *
      * @param <RESPONSE> the intermediate type of {@link AbstractResponse}
      */
-    private abstract static class AbstractBase<RESPONSE extends AbstractResponse, RECORD> implements IpDataLookup {
+    private abstract static class AbstractBase<RESPONSE extends AbstractResponse, RECORD extends GeoIpCache.CacheableValue>
+        implements
+            IpDataLookup {
 
         protected final Set<Database.Property> properties;
         protected final Class<RESPONSE> clazz;