Add trie copy method and refactor CSVFileDataAdapter parseCSVFile logic

graylog2-server/src/main/java/org/graylog2/lookup/adapters/CSVFileDataAdapter.java

@@ -118,7 +118,7 @@ public void doStart() throws Exception {
 
         // Set file info before parsing the data for the first time
         fileInfo = getNewFileInfo();
-        lookupRef.set(parseCSVFile());
+        setLookupRefFromCSV();
     }
 
     @Override
@@ -152,7 +152,7 @@ protected void doRefresh(LookupCachePurge cachePurge) throws Exception {
             }
 
             LOG.debug("CSV file {} has changed, updating data", config.path());
-            lookupRef.set(parseCSVFile());
+            setLookupRefFromCSV();
             cachePurge.purgeAll();
             // If the file has been moved, then moved back, the fileInfo might have been disconnected.
             // In this case, create a new fileInfo.
@@ -164,7 +164,7 @@ protected void doRefresh(LookupCachePurge cachePurge) throws Exception {
         }
     }
 
-    private Map<String, String> parseCSVFile() throws IOException {
+    private void setLookupRefFromCSV() throws IOException {
         final InputStream inputStream = Files.newInputStream(Paths.get(config.path()));
         final InputStreamReader fileReader = new InputStreamReader(inputStream, StandardCharsets.UTF_8);
         final ImmutableMap.Builder<String, String> newLookupBuilder = ImmutableMap.builder();
@@ -237,6 +237,7 @@ private Map<String, String> parseCSVFile() throws IOException {
                     }
                 }
             }
+
         } catch (Exception e) {
             LOG.error("Couldn't parse CSV file {} (settings separator=<{}> quotechar=<{}> key_column=<{}> value_column=<{}>)", config.path(),
                     config.separator(), config.quotechar(), config.keyColumn(), config.valueColumn(), e);
@@ -246,8 +247,9 @@ private Map<String, String> parseCSVFile() throws IOException {
 
         if (config.isCidrLookup()) {
             cidrLookupRef.set(cidrLookupTrie);
+        } else {
+            lookupRef.set(newLookupBuilder.build());
         }
-        return newLookupBuilder.build();
     }
 
     private String ipAddressToCIDR(String ip) {

graylog2-server/src/main/java/org/graylog2/utilities/CIDRLookupTrie.java

@@ -23,6 +23,8 @@
 
 /**
  * Binary tree used to for efficient lookups in CIDR data adapters.
+ * NOTE: This class is NOT thread-safe. Use {@link #copy()} to clone the trie, make modifications, and then atomically
+ * replace the in-use copy if needed.
  */
 public class CIDRLookupTrie {
     private static class TrieNode {
@@ -34,9 +36,30 @@ private static class TrieNode {
         boolean rangeIsIPv6 = false;
         // Time in millis after which the node is considered expired
         long expireAfter = 0L;
+
+        TrieNode deepCopy() {
+            TrieNode newNode = new TrieNode();
+            newNode.rangeName = this.rangeName;
+            newNode.rangeIsIPv6 = this.rangeIsIPv6;
+            newNode.expireAfter = this.expireAfter;
+            if (this.children[0] != null) {
+                newNode.children[0] = this.children[0].deepCopy();
+            }
+            if (this.children[1] != null) {
+                newNode.children[1] = this.children[1].deepCopy();
+            }
+
+            return newNode;
+        }
     }
 
-    private final TrieNode root = new TrieNode();
+    private TrieNode root = new TrieNode();
+
+    public CIDRLookupTrie copy() {
+        final CIDRLookupTrie copy = new CIDRLookupTrie();
+        copy.root = this.root.deepCopy();
+        return copy;
+    }
 
     public void insertCIDR(String cidr, String rangeName) {
         insertCIDR(cidr, rangeName, 0L);
@@ -45,8 +68,8 @@ public void insertCIDR(String cidr, String rangeName) {
     /**
      * Insert a CIDR range into the trie with a time-to-live
      *
-     * @param cidr properly formatted CIDR address (must include '/rangePrefix' even if it is a single address
-     * @param rangeName the name of the CIDR range
+     * @param cidr        properly formatted CIDR address (must include '/rangePrefix' even if it is a single address
+     * @param rangeName   the name of the CIDR range
      * @param expireAfter epoch time in millis after which the CIDR should be expired
      */
     public void insertCIDR(String cidr, String rangeName, long expireAfter) {
@@ -84,7 +107,7 @@ public String longestPrefixRangeLookup(String ip) {
      * Returns the rangeName of the range with the longest prefix that contains the IP address or null if one does not
      * exist.
      *
-     * @param ip IP address to check against the collection of ranges
+     * @param ip               IP address to check against the collection of ranges
      * @param lookupTimeMillis time lookup was performed in epoch time milliseconds or 0 if node expiry is not a concern
      * @return the name of the range with the longest prefix that contains the IP if it exists, null otherwise
      */