NIFI-15480 Added ability to drop FlowFiles selectively using a Predicate (#10848)

* NIFI-15480: Added ability to drop FlowFiles selectively using a Predicate
Added system tests to verify selective dropping of flowfiles from a Connector. Also made significant updates so that diagnostic bundles are captured when a system test fails and diagnostic bundle also contains details of each Connection's FlowFileQueue. This was necessary in aiding the debugging of the new feature but is also something that has been missing for some time.

* NIFI-15480: Ensure that when we perform a selective drop on a FlowFileQueue that we pass appropriate SWAP_FILE_DELETED / SWAP_FILE_RENAMED events to the FlowFile Repository

Author: markap14

nifi-commons/nifi-write-ahead-log/src/main/java/org/apache/nifi/wali/HashMapSnapshot.java

@@ -175,14 +175,18 @@ public void update(final Collection<T> records) {
 
             switch (updateType) {
                 case DELETE:
-                    recordMap.remove(recordId);
+                    if (recordId != null) {
+                        recordMap.remove(recordId);
+                    }
                     break;
                 case SWAP_OUT:
                     final String location = serdeFactory.getLocation(record);
                     if (location == null) {
                         logger.error(logMessage, recordId, UpdateType.SWAP_OUT, "Swapped Out to", "lost");
                     } else {
-                        recordMap.remove(recordId);
+                        if (recordId != null) {
+                            recordMap.remove(recordId);
+                        }
                         this.swapLocations.add(location);
                     }
                     break;
@@ -193,10 +197,30 @@ public void update(final Collection<T> records) {
                     } else {
                         swapLocations.remove(swapLocation);
                     }
-                    recordMap.put(recordId, record);
+                    if (recordId != null) {
+                        recordMap.put(recordId, record);
+                    }
+                    break;
+                case SWAP_FILE_DELETED:
+                    final String deletedSwapLocation = serdeFactory.getLocation(record);
+                    if (deletedSwapLocation != null) {
+                        swapLocations.remove(deletedSwapLocation);
+                    }
+                    break;
+                case SWAP_FILE_RENAMED:
+                    final String originalLocation = serdeFactory.getOriginalLocation(record);
+                    final String newLocation = serdeFactory.getLocation(record);
+                    if (originalLocation != null) {
+                        swapLocations.remove(originalLocation);
+                    }
+                    if (newLocation != null) {
+                        swapLocations.add(newLocation);
+                    }
                     break;
                 default:
-                    recordMap.put(recordId, record);
+                    if (recordId != null) {
+                        recordMap.put(recordId, record);
+                    }
                     break;
             }
         }
@@ -209,6 +233,9 @@ public int getRecordCount() {
 
     @Override
     public T lookup(final Object recordId) {
+        if (recordId == null) {
+            return null;
+        }
         return recordMap.get(recordId);
     }
 

nifi-commons/nifi-write-ahead-log/src/main/java/org/apache/nifi/wali/LengthDelimitedJournal.java

@@ -477,8 +477,10 @@ public JournalRecovery recoverRecords(final Map<Object, T> recordMap, final Set<
 
                         switch (updateType) {
                             case DELETE: {
-                                idsRemoved.add(recordId);
-                                transactionRecordMap.remove(recordId);
+                                if (recordId != null) {
+                                    idsRemoved.add(recordId);
+                                    transactionRecordMap.remove(recordId);
+                                }
                                 break;
                             }
                             case SWAP_IN: {
@@ -488,7 +490,9 @@ public JournalRecovery recoverRecords(final Map<Object, T> recordMap, final Set<
                                 } else {
                                     swapLocationsRemoved.add(location);
                                     swapLocationsAdded.remove(location);
-                                    transactionRecordMap.put(recordId, record);
+                                    if (recordId != null) {
+                                        transactionRecordMap.put(recordId, record);
+                                    }
                                 }
                                 break;
                             }
@@ -499,15 +503,40 @@ public JournalRecovery recoverRecords(final Map<Object, T> recordMap, final Set<
                                 } else {
                                     swapLocationsRemoved.remove(location);
                                     swapLocationsAdded.add(location);
-                                    idsRemoved.add(recordId);
-                                    transactionRecordMap.remove(recordId);
+                                    if (recordId != null) {
+                                        idsRemoved.add(recordId);
+                                        transactionRecordMap.remove(recordId);
+                                    }
                                 }
 
                                 break;
                             }
+                            case SWAP_FILE_DELETED: {
+                                final String location = serde.getLocation(record);
+                                if (location != null) {
+                                    swapLocationsRemoved.add(location);
+                                    swapLocationsAdded.remove(location);
+                                }
+                                break;
+                            }
+                            case SWAP_FILE_RENAMED: {
+                                final String originalLocation = serde.getOriginalLocation(record);
+                                final String newLocation = serde.getLocation(record);
+                                if (originalLocation != null) {
+                                    swapLocationsRemoved.add(originalLocation);
+                                    swapLocationsAdded.remove(originalLocation);
+                                }
+                                if (newLocation != null) {
+                                    swapLocationsAdded.add(newLocation);
+                                    swapLocationsRemoved.remove(newLocation);
+                                }
+                                break;
+                            }
                             default: {
-                                transactionRecordMap.put(recordId, record);
-                                idsRemoved.remove(recordId);
+                                if (recordId != null) {
+                                    transactionRecordMap.put(recordId, record);
+                                    idsRemoved.remove(recordId);
+                                }
                                 break;
                             }
                         }

nifi-commons/nifi-write-ahead-log/src/main/java/org/wali/SerDe.java

@@ -136,6 +136,19 @@ default void readHeader(DataInputStream in) throws IOException {
      */
     String getLocation(T record);
 
+    /**
+     * Returns the original external location of the given record; this is used when a
+     * swap file is renamed. For {@link UpdateType#SWAP_FILE_RENAMED} records, this
+     * returns the original location before the rename, while {@link #getLocation(Object)}
+     * returns the new location after the rename.
+     *
+     * @param record to get original location of
+     * @return original location, or null if not applicable
+     */
+    default String getOriginalLocation(T record) {
+        return null;
+    }
+
     /**
      * Returns the version that this SerDe will use when writing. This used used
      * when serializing/deserializing the edit logs so that if the version

nifi-commons/nifi-write-ahead-log/src/main/java/org/wali/SerDeFactory.java

@@ -57,4 +57,17 @@ public interface SerDeFactory<T> {
      * @return location
      */
     String getLocation(T record);
+
+    /**
+     * Returns the original external location of the given record; this is used when a
+     * swap file is renamed. For {@link UpdateType#SWAP_FILE_RENAMED} records, this
+     * returns the original location before the rename, while {@link #getLocation(Object)}
+     * returns the new location after the rename.
+     *
+     * @param record to get original location of
+     * @return original location, or null if not applicable
+     */
+    default String getOriginalLocation(T record) {
+        return null;
+    }
 }

nifi-commons/nifi-write-ahead-log/src/main/java/org/wali/UpdateType.java

@@ -45,5 +45,13 @@ public enum UpdateType {
      * Used to indicate that a Record that was previously Swapped Out is now
      * being Swapped In
      */
-    SWAP_IN;
+    SWAP_IN,
+    /**
+     * Used to indicate that a Swap File has been deleted
+     */
+    SWAP_FILE_DELETED,
+    /**
+     * Used to indicate that a Swap File has been renamed
+     */
+    SWAP_FILE_RENAMED;
 }

nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/FlowFileQueue.java

@@ -16,9 +16,11 @@
  */
 package org.apache.nifi.controller.queue;
 
+import org.apache.nifi.components.connector.DropFlowFileSummary;
 import org.apache.nifi.controller.repository.FlowFileRecord;
 import org.apache.nifi.controller.repository.SwapSummary;
 import org.apache.nifi.controller.status.FlowFileAvailability;
+import org.apache.nifi.flowfile.FlowFile;
 import org.apache.nifi.flowfile.FlowFilePrioritizer;
 import org.apache.nifi.processor.FlowFileFilter;
 
@@ -27,6 +29,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Predicate;
 
 public interface FlowFileQueue {
 
@@ -217,6 +220,18 @@ public interface FlowFileQueue {
      */
     DropFlowFileStatus cancelDropFlowFileRequest(String requestIdentifier);
 
+    /**
+     * Synchronously drops all FlowFiles in this queue that match the given predicate. This method filters
+     * FlowFiles in the active queue, swap queue, and any swapped-out swap files. The FlowFile Repository
+     * and Provenance Repository are updated atomically after all matching FlowFiles have been identified.
+     *
+     * @param predicate the predicate used to determine which FlowFiles should be dropped; FlowFiles for which
+     *                  the predicate returns <code>true</code> will be dropped
+     * @return a summary of the FlowFiles that were dropped, including the count and total size in bytes
+     * @throws IOException if an error occurs while reading or writing swap files
+     */
+    DropFlowFileSummary dropFlowFiles(Predicate<FlowFile> predicate) throws IOException;
+
     /**
      * <p>
      * Initiates a request to obtain a listing of FlowFiles in this queue. This method returns a

nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/LocalQueuePartitionDiagnostics.java

@@ -29,4 +29,14 @@ public interface LocalQueuePartitionDiagnostics {
     boolean isAnyActiveFlowFilePenalized();
 
     boolean isAllActiveFlowFilesPenalized();
+
+    /**
+     * @return the QueueSize representing the penalized FlowFiles (count and total content size)
+     */
+    QueueSize getPenalizedQueueSize();
+
+    /**
+     * @return the total QueueSize across all swap files (count and total content size)
+     */
+    QueueSize getTotalSwapFileQueueSize();
 }

nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/FlowFileSwapManager.java

@@ -126,6 +126,16 @@ public interface FlowFileSwapManager {
      */
     void purge();
 
+    /**
+     * Deletes the swap file at the given location without updating the FlowFile Repository.
+     * This method is intended for use after the FlowFile Repository has already been updated
+     * to reflect that the FlowFiles are no longer in this swap file.
+     *
+     * @param swapLocation the location of the swap file to delete
+     * @throws IOException if unable to delete the swap file
+     */
+    void deleteSwapFile(String swapLocation) throws IOException;
+
     /**
      * Returns the ID of the queue that the given swap file belongs to
      * @param swapLocation the swap location

nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/RepositoryRecord.java

@@ -82,6 +82,14 @@ public interface RepositoryRecord {
      */
     String getSwapLocation();
 
+    /**
+     * @return For SWAP_FILE_RENAMED records, provides the original swap location before the rename.
+     * For other record types, returns <code>null</code>.
+     */
+    default String getOriginalSwapLocation() {
+        return null;
+    }
+
     /**
      * @return a List of Content Claims that are "transient," meaning that they existed only for the
      *         life of the Process Session in which they were created and should not be persisted.

nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/RepositoryRecordType.java

@@ -21,5 +21,5 @@
  */
 public enum RepositoryRecordType {
 
-    UPDATE, CREATE, DELETE, CONTENTMISSING, SWAP_IN, SWAP_OUT, CLEANUP_TRANSIENT_CLAIMS;
+    UPDATE, CREATE, DELETE, CONTENTMISSING, SWAP_IN, SWAP_OUT, CLEANUP_TRANSIENT_CLAIMS, SWAP_FILE_DELETED, SWAP_FILE_RENAMED;
 }