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 <[email protected]>
opensearch-project · Mar 28, 2024 · b9bf6e2 · b9bf6e2
1 parent 4f3ab96
commit b9bf6e2
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 1 deletion.
diff --git a/...ain/java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReader.java b/...ain/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<GeoDataInput, Map<String, Object>> geoDataCache = new HashMap<>();
     private Boolean isExpired = null;
 
     BatchGeoIPDatabaseReader(final GeoIPDatabaseReader delegate) {
@@ -32,7 +35,8 @@ static BatchGeoIPDatabaseReader decorate(final GeoIPDatabaseReader delegate) {
 
     @Override
     public Map<String, Object> getGeoData(final InetAddress inetAddress, final Collection<GeoIPField> fields, final Collection<GeoIPDatabase> 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<GeoIPField> fields;
+        final Collection<GeoIPDatabase> geoIPDatabases;
+
+        private GeoDataInput(InetAddress inetAddress, Collection<GeoIPField> fields, Collection<GeoIPDatabase> 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/...java/org/opensearch/dataprepper/plugins/geoip/processor/BatchGeoIPDatabaseReaderTest.java b/...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<GeoIPField> fields = List.of(mock(GeoIPField.class));
+        final Set<GeoIPDatabase> geoIPDatabases = Set.of(mock(GeoIPDatabase.class));
+
+        final Map<String, Object> 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<GeoIPField> fields = List.of(mock(GeoIPField.class));
+        final Set<GeoIPDatabase> geoIPDatabases = Set.of(mock(GeoIPDatabase.class));
+
+        final Map<String, Object> 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();