Browse Source

A small tidiness refactor of the GeoIpTaskState's Metadata (#110553)

Joe Gallo 1 year ago
parent
commit
00e744ef5c

+ 1 - 1
modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/GeoIpDownloaderIT.java

@@ -242,7 +242,7 @@ public class GeoIpDownloaderIT extends AbstractGeoIpIT {
                         Set.of("GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb", "MyCustomGeoLite2-City.mmdb"),
                         state.getDatabases().keySet()
                     );
-                    GeoIpTaskState.Metadata metadata = state.get(id);
+                    GeoIpTaskState.Metadata metadata = state.getDatabases().get(id);
                     int size = metadata.lastChunk() - metadata.firstChunk() + 1;
                     assertResponse(
                         prepareSearch(GeoIpDownloader.DATABASES_INDEX).setSize(size)

+ 12 - 7
modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/GeoIpDownloader.java

@@ -170,23 +170,28 @@ public class GeoIpDownloader extends AllocatedPersistentTask {
     }
 
     // visible for testing
-    void processDatabase(Map<String, Object> databaseInfo) {
+    void processDatabase(final Map<String, Object> databaseInfo) {
         String name = databaseInfo.get("name").toString().replace(".tgz", "") + ".mmdb";
         String md5 = (String) databaseInfo.get("md5_hash");
-        if (state.contains(name) && Objects.equals(md5, state.get(name).md5())) {
-            updateTimestamp(name, state.get(name));
-            return;
-        }
-        logger.debug("downloading geoip database [{}]", name);
         String url = databaseInfo.get("url").toString();
         if (url.startsWith("http") == false) {
             // relative url, add it after last slash (i.e. resolve sibling) or at the end if there's no slash after http[s]://
             int lastSlash = endpoint.substring(8).lastIndexOf('/');
             url = (lastSlash != -1 ? endpoint.substring(0, lastSlash + 8) : endpoint) + "/" + url;
         }
+        processDatabase(name, md5, url);
+    }
+
+    private void processDatabase(final String name, final String md5, final String url) {
+        Metadata metadata = state.getDatabases().getOrDefault(name, Metadata.EMPTY);
+        if (Objects.equals(metadata.md5(), md5)) {
+            updateTimestamp(name, metadata);
+            return;
+        }
+        logger.debug("downloading geoip database [{}]", name);
         long start = System.currentTimeMillis();
         try (InputStream is = httpClient.get(url)) {
-            int firstChunk = state.contains(name) ? state.get(name).lastChunk() + 1 : 0;
+            int firstChunk = metadata.lastChunk() + 1; // if there is no metadata, then Metadata.EMPTY.lastChunk() + 1 = 0
             int lastChunk = indexChunks(name, is, firstChunk, md5, start);
             if (lastChunk > firstChunk) {
                 state = state.put(name, new Metadata(start, firstChunk, lastChunk - 1, md5, start));

+ 7 - 9
modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/GeoIpTaskState.java

@@ -84,14 +84,6 @@ class GeoIpTaskState implements PersistentTaskState, VersionedNamedWriteable {
         return databases;
     }
 
-    public boolean contains(String name) {
-        return databases.containsKey(name);
-    }
-
-    public Metadata get(String name) {
-        return databases.get(name);
-    }
-
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -142,7 +134,13 @@ class GeoIpTaskState implements PersistentTaskState, VersionedNamedWriteable {
 
     record Metadata(long lastUpdate, int firstChunk, int lastChunk, String md5, long lastCheck) implements ToXContentObject {
 
-        static final String NAME = GEOIP_DOWNLOADER + "-metadata";
+        /**
+         * An empty Metadata object useful for getOrDefault -type calls. Crucially, the 'lastChunk' is -1, so it's safe to use
+         * with logic that says the new firstChunk is the old lastChunk + 1.
+         */
+        static Metadata EMPTY = new Metadata(-1, -1, -1, "", -1);
+
+        private static final String NAME = GEOIP_DOWNLOADER + "-metadata";
         private static final ParseField LAST_CHECK = new ParseField("last_check");
         private static final ParseField LAST_UPDATE = new ParseField("last_update");
         private static final ParseField FIRST_CHUNK = new ParseField("first_chunk");

+ 4 - 4
modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/GeoIpDownloaderTests.java

@@ -290,8 +290,8 @@ public class GeoIpDownloaderTests extends ESTestCase {
 
             @Override
             void updateTaskState() {
-                assertEquals(0, state.get("test.mmdb").firstChunk());
-                assertEquals(10, state.get("test.mmdb").lastChunk());
+                assertEquals(0, state.getDatabases().get("test.mmdb").firstChunk());
+                assertEquals(10, state.getDatabases().get("test.mmdb").lastChunk());
             }
 
             @Override
@@ -341,8 +341,8 @@ public class GeoIpDownloaderTests extends ESTestCase {
 
             @Override
             void updateTaskState() {
-                assertEquals(9, state.get("test.mmdb").firstChunk());
-                assertEquals(10, state.get("test.mmdb").lastChunk());
+                assertEquals(9, state.getDatabases().get("test.mmdb").firstChunk());
+                assertEquals(10, state.getDatabases().get("test.mmdb").lastChunk());
             }
 
             @Override