diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/SwappablePriorityQueue.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/SwappablePriorityQueue.java index d2f13071c792..57f18cc76406 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/SwappablePriorityQueue.java +++ b/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 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 predicate) th droppedFlowFiles.addAll(result.droppedFlowFiles()); } + // Replace rewritten swap files in place so that retained FlowFiles remain ahead of later swap files. + final List 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 newActiveQueue = new PriorityQueue<>(Math.max(20, activeQueue.size()), new QueuePrioritizer(getPriorities())); int droppedFromActiveCount = 0; diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/TestSwappablePriorityQueue.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/TestSwappablePriorityQueue.java index f8035fd24510..822e5406d907 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/TestSwappablePriorityQueue.java +++ b/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 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 {