From b9bf6e258a0eb62cff88ad55eccb0cbc299c15f8 Mon Sep 17 00:00:00 2001 From: David Venable Date: Wed, 27 Mar 2024 19:21:27 -0500 Subject: [PATCH] Cache geolocation data within any given batch of Data Prepper events to avoid extra calls to the MaxMind GeoIP code. Signed-off-by: David Venable --- .../processor/BatchGeoIPDatabaseReader.java | 31 +++++++++++++++- .../BatchGeoIPDatabaseReaderTest.java | 37 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReader.java index d4430bba44..74ec5608b0 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReader.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReader.java @@ -11,13 +11,16 @@ import java.net.InetAddress; import java.util.Collection; +import java.util.HashMap; import java.util.Map; +import java.util.Objects; /** * A {@link GeoIPDatabaseReader} useful for a single batch of processing. */ class BatchGeoIPDatabaseReader implements GeoIPDatabaseReader { private final GeoIPDatabaseReader delegate; + private final Map> geoDataCache = new HashMap<>(); private Boolean isExpired = null; BatchGeoIPDatabaseReader(final GeoIPDatabaseReader delegate) { @@ -32,7 +35,8 @@ static BatchGeoIPDatabaseReader decorate(final GeoIPDatabaseReader delegate) { @Override public Map getGeoData(final InetAddress inetAddress, final Collection fields, final Collection geoIPDatabases) { - return delegate.getGeoData(inetAddress, fields, geoIPDatabases); + final GeoDataInput geoDataInput = new GeoDataInput(inetAddress, fields, geoIPDatabases); + return geoDataCache.computeIfAbsent(geoDataInput, unused -> delegate.getGeoData(inetAddress, fields, geoIPDatabases)); } @Override @@ -52,4 +56,29 @@ public void retain() { public void close() throws Exception { delegate.close(); } + + private static class GeoDataInput { + final InetAddress inetAddress; + final Collection fields; + final Collection geoIPDatabases; + + private GeoDataInput(InetAddress inetAddress, Collection fields, Collection geoIPDatabases) { + this.inetAddress = inetAddress; + this.fields = fields; + this.geoIPDatabases = geoIPDatabases; + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + final GeoDataInput that = (GeoDataInput) o; + return Objects.equals(inetAddress, that.inetAddress) && Objects.equals(fields, that.fields) && Objects.equals(geoIPDatabases, that.geoIPDatabases); + } + + @Override + public int hashCode() { + return Objects.hash(inetAddress, fields, geoIPDatabases); + } + } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReaderTest.java index c0968f34da..2b14b38612 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReaderTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReaderTest.java @@ -25,6 +25,7 @@ import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -59,6 +60,42 @@ void getGeoData_returns_from_inner_geoIPDatabaseReader() { equalTo(geoData)); } + @Test + void getGeoData_returns_cached_value() { + final InetAddress inetAddress = mock(InetAddress.class); + final List fields = List.of(mock(GeoIPField.class)); + final Set geoIPDatabases = Set.of(mock(GeoIPDatabase.class)); + + final Map geoData = Map.of(UUID.randomUUID().toString(), UUID.randomUUID().toString()); + when(geoIPDatabaseReader.getGeoData(inetAddress, fields, geoIPDatabases)) + .thenReturn(geoData); + + final BatchGeoIPDatabaseReader objectUnderTest = objectUnderTestFromDecorate(); + assertThat(objectUnderTest.getGeoData(inetAddress, fields, geoIPDatabases), equalTo(geoData)); + assertThat(objectUnderTest.getGeoData(inetAddress, fields, geoIPDatabases), equalTo(geoData)); + + verify(geoIPDatabaseReader, times(1)).getGeoData(inetAddress, fields, geoIPDatabases); + } + + @Test + void getGeoData_caches_for_IP_fields_and_database() { + final InetAddress inetAddress = mock(InetAddress.class); + final List fields = List.of(mock(GeoIPField.class)); + final Set geoIPDatabases = Set.of(mock(GeoIPDatabase.class)); + + final Map geoData = Map.of(UUID.randomUUID().toString(), UUID.randomUUID().toString()); + when(geoIPDatabaseReader.getGeoData(any(), any(), any())) + .thenReturn(geoData); + + final BatchGeoIPDatabaseReader objectUnderTest = objectUnderTestFromDecorate(); + assertThat(objectUnderTest.getGeoData(inetAddress, fields, geoIPDatabases), equalTo(geoData)); + assertThat(objectUnderTest.getGeoData(mock(InetAddress.class), fields, geoIPDatabases), equalTo(geoData)); + assertThat(objectUnderTest.getGeoData(inetAddress, fields, Set.of(mock(GeoIPDatabase.class))), equalTo(geoData)); + assertThat(objectUnderTest.getGeoData(inetAddress, List.of(mock(GeoIPField.class)), geoIPDatabases), equalTo(geoData)); + + verify(geoIPDatabaseReader, times(4)).getGeoData(any(), any(), any()); + } + @Test void close_calls_inner_geoIPDataReader_close() throws Exception { objectUnderTestFromDecorate().close();