NIFI-16344 Preserve Swap File Ordering during Selective Drop

nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/SwappablePriorityQueue.java

@@ -1014,11 +1014,9 @@ public SelectiveDropResult dropFlowFiles(final Predicate<FlowFile> predicate) th
                 final String newSwapLocation = result.newSwapLocation();
 
                 swapLocationUpdates.put(oldSwapLocation, newSwapLocation);
-                swapLocations.remove(oldSwapLocation);
 
                 if (newSwapLocation != null) {
                     // Some FlowFiles remain in new swap file
-                    swapLocations.add(newSwapLocation);
                     incrementSwapQueueSize(-result.droppedFlowFiles().size(), -result.droppedBytes(), 0);
 
                     // Update metrics for the new swap location
@@ -1040,6 +1038,21 @@ public SelectiveDropResult dropFlowFiles(final Predicate<FlowFile> predicate) th
                 droppedFlowFiles.addAll(result.droppedFlowFiles());
             }
 
+            // Replace rewritten swap files in place so that retained FlowFiles remain ahead of later swap files.
+            final List<String> rewrittenSwapLocations = new ArrayList<>(swapLocations.size());
+            for (final String swapLocation : swapLocations) {
+                if (swapLocationUpdates.containsKey(swapLocation)) {
+                    final String newSwapLocation = swapLocationUpdates.get(swapLocation);
+                    if (newSwapLocation != null) {
+                        rewrittenSwapLocations.add(newSwapLocation);
+                    }
+                } else {
+                    rewrittenSwapLocations.add(swapLocation);
+                }
+            }
+            swapLocations.clear();
+            swapLocations.addAll(rewrittenSwapLocations);
+
             // Filter the active queue
             final Queue<FlowFileRecord> newActiveQueue = new PriorityQueue<>(Math.max(20, activeQueue.size()), new QueuePrioritizer(getPriorities()));
             int droppedFromActiveCount = 0;

nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/TestSwappablePriorityQueue.java

@@ -734,6 +734,34 @@ public void testSelectiveDropFromSwapFile() throws IOException {
         assertEquals(2, result.getSwapLocationUpdates().size());
     }
 
+    @Test
+    public void testSelectiveDropPreservesSwapFileOrdering() throws IOException {
+        queue = new SwappablePriorityQueue(swapManager, 10, eventReporter, flowFileQueue, dropAction, "local");
+
+        for (int i = 0; i < 30; i++) {
+            queue.put(new MockFlowFileRecord(Map.of("index", Integer.toString(i)), 1L));
+        }
+
+        assertEquals(2, swapManager.swappedOut.size());
+
+        final SelectiveDropResult result = queue.dropFlowFiles(flowFile -> "10".equals(flowFile.getAttribute("index")));
+
+        assertEquals(1, result.getDroppedCount());
+        assertEquals(1, result.getSwapLocationUpdates().size());
+
+        final Set<FlowFileRecord> expired = new HashSet<>();
+        for (int expectedIndex = 0; expectedIndex < 30; expectedIndex++) {
+            if (expectedIndex == 10) {
+                continue;
+            }
+
+            final FlowFileRecord flowFile = queue.poll(expired, 0L);
+            assertNotNull(flowFile);
+            assertEquals(Integer.toString(expectedIndex), flowFile.getAttribute("index"));
+        }
+        assertNull(queue.poll(expired, 0L));
+    }
+
     @Test
     @Timeout(120)
     public void testSelectiveDropRemovesEntireSwapFileWhenAllMatch() throws IOException {