geoip: managed databases need to be closed when updated (#178)

Author: yaauie

src/main/java/co/elastic/logstash/filters/elasticintegration/geoip/ConstantIpDatabaseHolder.java

@@ -19,7 +19,7 @@ public boolean isValid() {
     }
 
     @Override
-    public IpDatabase getDatabase() {
+    public IpDatabaseAdapter getDatabase() {
         return this.ipDatabase;
     }
 

src/main/java/co/elastic/logstash/filters/elasticintegration/geoip/IpDatabaseAdapter.java

@@ -27,6 +27,8 @@ public class IpDatabaseAdapter implements IpDatabase {
     private final Reader databaseReader;
     private final String databaseType;
 
+    private volatile boolean isReaderClosed = false;
+
     public IpDatabaseAdapter(final Reader databaseReader) {
         this.databaseReader = databaseReader;
         this.databaseType = databaseReader.getMetadata().getDatabaseType();
@@ -56,6 +58,12 @@ public void close() throws IOException {
     public void closeReader() throws IOException {
         LOGGER.debug("Closing the database adapter");
         this.databaseReader.close();
+        this.isReaderClosed = true;
+    }
+
+    // visible for test
+    boolean isReaderClosed() {
+        return this.isReaderClosed;
     }
 
     public static IpDatabaseAdapter defaultForPath(final Path database) throws IOException {

src/main/java/co/elastic/logstash/filters/elasticintegration/geoip/IpDatabaseHolder.java

@@ -5,7 +5,7 @@
 interface IpDatabaseHolder {
     boolean isValid();
 
-    IpDatabase getDatabase();
+    IpDatabaseAdapter getDatabase();
 
     String getTypeIdentifier();
 

src/main/java/co/elastic/logstash/filters/elasticintegration/geoip/IpDatabaseProvider.java

@@ -34,20 +34,25 @@ public class IpDatabaseProvider implements org.elasticsearch.ingest.geoip.IpData
 
     @Override
     public Boolean isValid(String databaseIdentifierFileName) {
-        final IpDatabaseHolder holder = databaseMap.get(databaseIdentifierFileName);
+        final IpDatabaseHolder holder = getDatabaseHolder(databaseIdentifierFileName);
         return Objects.nonNull(holder) && holder.isValid();
     }
 
     @Override
     public IpDatabase getDatabase(String databaseIdentifierFileName) {
-        final IpDatabaseHolder holder = databaseMap.get(databaseIdentifierFileName);
+        final IpDatabaseHolder holder = getDatabaseHolder(databaseIdentifierFileName);
         if (Objects.isNull(holder)) {
             return null;
         } else {
             return holder.getDatabase();
         }
     }
 
+    // visible for test
+    IpDatabaseHolder getDatabaseHolder(String databaseIdentifierFileName) {
+        return databaseMap.get(databaseIdentifierFileName);
+    }
+
     @Override
     public void close() throws IOException {
         databaseMap.forEach((name, holder) -> {

src/main/java/co/elastic/logstash/filters/elasticintegration/geoip/ManagedIpDatabaseHolder.java

@@ -2,6 +2,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.ingest.geoip.IpDatabase;
 
 import java.io.Closeable;
@@ -30,7 +31,14 @@ public class ManagedIpDatabaseHolder implements IpDatabaseHolder, Closeable {
     }
 
     public ManagedIpDatabaseHolder(final String databaseTypeIdentifier) {
+        this(databaseTypeIdentifier, null);
+    }
+
+    ManagedIpDatabaseHolder(final String databaseTypeIdentifier, final Path databaseLocation) {
         this.databaseTypeIdentifier = databaseTypeIdentifier;
+        if (Objects.nonNull(databaseLocation)) {
+            setDatabasePath(databaseLocation.toAbsolutePath().toString());
+        }
     }
 
     @Override
@@ -39,7 +47,7 @@ public boolean isValid() {
     }
 
     @Override
-    public IpDatabase getDatabase() {
+    public IpDatabaseAdapter getDatabase() {
         return withLock(readLock, () -> this.currentDatabase);
     }
 
@@ -54,12 +62,18 @@ public String info() {
     }
 
     public void setDatabasePath(final String newDatabasePath) {
-        withLock(writeLock, () -> {
+        final IpDatabaseAdapter previousDatabase = withLock(writeLock, () -> {
+            final IpDatabaseAdapter localPreviousDatabase = this.currentDatabase;
             this.currentDatabase = Optional.ofNullable(newDatabasePath)
                     .map(Paths::get)
                     .map(this::loadDatabase)
                     .orElse(null);
+            return localPreviousDatabase;
         });
+
+        if (previousDatabase != null) {
+            IOUtils.closeWhileHandlingException(previousDatabase::closeReader);
+        }
     }
 
     @Override

src/test/java/co/elastic/logstash/filters/elasticintegration/geoip/ManagedIpDatabaseHolderTest.java

@@ -0,0 +1,78 @@
+package co.elastic.logstash.filters.elasticintegration.geoip;
+
+import co.elastic.logstash.filters.elasticintegration.util.ResourcesUtil;
+import org.apache.commons.io.FilenameUtils;
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Test;
+
+import java.nio.file.Path;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.*;
+
+
+class ManagedIpDatabaseHolderTest {
+    static final Path GEOLITE2_ASN_MMDB_PATH;
+    static final Path GEOLITE2_CITY_MMDB_PATH;
+
+    static final String GEOLITE2_ASN_TYPE = "GeoLite2-ASN.mmdb";
+    static final String GEOLITE2_CITY_TYPE = "GeoLite2-City.mmdb";
+
+    static {
+        try {
+            final Path databases = ResourcesUtil.getResourcePath(ManagedIpDatabaseHolderTest.class, "databases").orElseThrow();
+            GEOLITE2_ASN_MMDB_PATH = databases.resolve(GEOLITE2_ASN_TYPE).toRealPath();
+            GEOLITE2_CITY_MMDB_PATH = databases.resolve(GEOLITE2_CITY_TYPE).toRealPath();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    IpDatabaseProvider buildIpDatabaseProvider(final Path... paths) {
+        final IpDatabaseProvider.Builder builder = new IpDatabaseProvider.Builder();
+        for (Path path : paths) {
+            final String databaseFileName = path.getFileName().toString();
+            final String databaseTypeIdentifier = FilenameUtils.removeExtension(databaseFileName);
+            @SuppressWarnings("resource") final ManagedIpDatabaseHolder dbh = new ManagedIpDatabaseHolder(databaseTypeIdentifier, path);
+            builder.setDatabaseHolder(databaseFileName, dbh);
+        }
+        return builder.build();
+    }
+
+    @Test
+    void releaseOnReload() throws Exception {
+        try (IpDatabaseProvider ipDatabaseProvider = buildIpDatabaseProvider(GEOLITE2_CITY_MMDB_PATH)) {
+            final IpDatabaseHolder cityIpDatabaseHolder = ipDatabaseProvider.getDatabaseHolder(GEOLITE2_CITY_TYPE);
+            assert cityIpDatabaseHolder instanceof ManagedIpDatabaseHolder;
+
+            final IpDatabaseAdapter first = cityIpDatabaseHolder.getDatabase();
+            assertFalse(first.isReaderClosed());
+
+            // update the holder's path. It's okay to be the same underlying file.
+            ((ManagedIpDatabaseHolder) cityIpDatabaseHolder).setDatabasePath(GEOLITE2_CITY_MMDB_PATH.toString());
+
+            // get the database again, and make sure we have a new open one, and that the old one has been closed
+            final IpDatabaseAdapter second = cityIpDatabaseHolder.getDatabase();
+            assertNotEquals(first, second);
+            assertTrue(first.isReaderClosed());
+            assertFalse(second.isReaderClosed());
+        }
+    }
+
+    @Test
+    void rejectTypeChangeOnReload() throws Exception {
+        try (IpDatabaseProvider ipDatabaseProvider = buildIpDatabaseProvider(GEOLITE2_CITY_MMDB_PATH)) {
+            final IpDatabaseHolder cityIpDatabaseHolder = ipDatabaseProvider.getDatabaseHolder(GEOLITE2_CITY_TYPE);
+            assert cityIpDatabaseHolder instanceof ManagedIpDatabaseHolder;
+
+            final IpDatabaseAdapter first = cityIpDatabaseHolder.getDatabase();
+            assertFalse(first.isReaderClosed());
+
+            // update the holder's path to the wrong type
+            IllegalStateException illegalStateException = assertThrows(IllegalStateException.class, () -> {
+                ((ManagedIpDatabaseHolder) cityIpDatabaseHolder).setDatabasePath(GEOLITE2_ASN_MMDB_PATH.toString());
+            });
+            assertThat(illegalStateException.getMessage(), Matchers.containsString("Incompatible database type `GeoLite2-ASN` (expected `GeoLite2-City`)"));
+        }
+    }
+}