This is an automated email from the ASF dual-hosted git repository.
pvillard31 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new c8197c24020 NIFI-16468: Return rolled-back FlowFiles to the active
queue (#11815)
c8197c24020 is described below
commit c8197c2402046f78e8ec31e70d0f363f4aadd42a
Author: Mark Payne <[email protected]>
AuthorDate: Fri Oct 23 11:52:38 2026 -0400
NIFI-16468: Return rolled-back FlowFiles to the active queue (#11815)
* NIFI-16468: Return rolled-back FlowFiles to the active queue
Session rollback used FlowFileQueue.put(), which places a FlowFile on the
swap queue when the queue is swapping. That puts it behind FlowFiles that were
already in the active queue.
---
.../nifi/controller/queue/FlowFileQueue.java | 10 ++++
.../repository/StandardProcessSession.java | 9 ++--
.../controller/queue/StandardFlowFileQueue.java | 5 ++
.../controller/queue/SwappablePriorityQueue.java | 20 +++++++
.../clustered/SocketLoadBalancedFlowFileQueue.java | 14 +++++
.../queue/clustered/partition/QueuePartition.java | 7 +++
.../clustered/partition/RemoteQueuePartition.java | 5 ++
.../partition/StandardRebalancingPartition.java | 6 +++
.../SwappablePriorityQueueLocalPartition.java | 5 ++
.../nifi/controller/tasks/StatelessFlowTask.java | 9 ++--
.../clustered/TestSwappablePriorityQueue.java | 13 +++++
.../repository/StandardProcessSessionIT.java | 23 ++++++++
.../TestWriteAheadFlowFileRepository.java | 4 ++
.../controller/tasks/TestStatelessFlowTask.java | 63 ++++++++++++++++++++++
.../stateless/queue/StatelessFlowFileQueue.java | 6 +++
15 files changed, 191 insertions(+), 8 deletions(-)
diff --git
a/nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/FlowFileQueue.java
b/nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/FlowFileQueue.java
index d6cd4811539..10f7290f813 100644
---
a/nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/FlowFileQueue.java
+++
b/nifi-framework-api/src/main/java/org/apache/nifi/controller/queue/FlowFileQueue.java
@@ -161,6 +161,16 @@ public interface FlowFileQueue {
*/
void put(FlowFileRecord file);
+ /**
+ * Places the given FlowFile onto the active portion of the queue, never
onto the swap queue, even if the queue is
+ * currently swapping FlowFiles out. This is intended for FlowFiles that
were previously polled from this queue and are
+ * being returned to it, such as on session rollback. Because such a
FlowFile was polled from the active queue, it was
+ * ahead of any FlowFiles that are swapped out, and placing it back onto
the active queue preserves that ordering.
+ *
+ * @param file the FlowFile to place onto the active queue
+ */
+ void putToActive(FlowFileRecord file);
+
/**
* places the given files into the queue
*
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
index b1db067a05d..0a6d214fa06 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java
@@ -1441,7 +1441,8 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
/**
* Rolls back the Record in a manner that is appropriate for the context.
The default implementation
- * is to place the queue back on its original queue, if it exists, or just
ignore it if it has no original queue.
+ * is to place the FlowFile back on the active portion of its original
queue, if it exists, so that it stays ahead of
+ * any FlowFiles that have been swapped out, or just ignore it if it has
no original queue.
* However, subclasses may wish to change the behavior for how Records are
handled when a rollback occurs.
* @param record the Record that is to be rolled back
* @param penalize whether or not the Record should be penalized
@@ -1453,9 +1454,9 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
if (penalize) {
final long expirationEpochMillis =
System.currentTimeMillis() +
context.getConnectable().getPenalizationPeriod(TimeUnit.MILLISECONDS);
final FlowFileRecord newFile = new
StandardFlowFileRecord.Builder().fromFlowFile(record.getOriginal()).penaltyExpirationTime(expirationEpochMillis).build();
- originalQueue.put(newFile);
+ originalQueue.putToActive(newFile);
} else {
- originalQueue.put(record.getOriginal());
+ originalQueue.putToActive(record.getOriginal());
}
}
}
@@ -2014,7 +2015,7 @@ public class StandardProcessSession implements
ProcessSession, ProvenanceEventEn
LOG.error("Attempted to pull {} from {} but the Session already has a
FlowFile with the same ID ({}): {}, which was pulled from {}. This means that
the system has two FlowFiles with the" +
" same ID, which should not happen.", flowFile, connection,
flowFile.getId(), conflict.getCurrent(), conflict.getOriginalQueue());
- connection.getFlowFileQueue().put(flowFile);
+ connection.getFlowFileQueue().putToActive(flowFile);
rollback(true, false);
throw new FlowFileAccessException("Attempted to pull a FlowFile with
ID " + flowFile.getId() + " from Connection "
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/StandardFlowFileQueue.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/StandardFlowFileQueue.java
index cbfe04cab2d..c9759a6907c 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/StandardFlowFileQueue.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/StandardFlowFileQueue.java
@@ -119,6 +119,11 @@ public class StandardFlowFileQueue extends
AbstractFlowFileQueue implements Flow
queue.put(file);
}
+ @Override
+ public void putToActive(final FlowFileRecord file) {
+ queue.putToActive(file);
+ }
+
@Override
public void putAll(final Collection<FlowFileRecord> files) {
queue.putAll(files);
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 8778683713e..171113afb4a 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
@@ -547,6 +547,26 @@ public class SwappablePriorityQueue {
}
}
+ /**
+ * Adds the given FlowFile to the active queue regardless of whether this
queue is in swap mode. This is intended only for FlowFiles that
+ * were previously polled from the active queue and are being returned.
Such a FlowFile was ahead of everything in the swap queue and swap files,
+ * so adding it back to the active queue keeps it ahead of them. This may
cause the active queue to temporarily exceed the swap threshold.
+ *
+ * @param flowFile the FlowFile to add to the active queue
+ */
+ public void putToActive(final FlowFileRecord flowFile) {
+ writeLock.lock();
+ try {
+ incrementActiveQueueSize(1, flowFile.getSize());
+ activeQueue.add(flowFile);
+
+ updateTopPenaltyExpiration();
+ logger.trace("{} put to active queue of {}", flowFile, this);
+ } finally {
+ writeLock.unlock("putToActive");
+ }
+ }
+
public void putAll(final Collection<FlowFileRecord> flowFiles) {
final int numFiles = flowFiles.size();
long bytes = 0L;
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/SocketLoadBalancedFlowFileQueue.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/SocketLoadBalancedFlowFileQueue.java
index dde99caaabd..db05a350d8a 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/SocketLoadBalancedFlowFileQueue.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/SocketLoadBalancedFlowFileQueue.java
@@ -864,6 +864,20 @@ public class SocketLoadBalancedFlowFileQueue extends
AbstractFlowFileQueue imple
return partition;
}
+ @Override
+ public void putToActive(final FlowFileRecord flowFile) {
+ partitionReadLock.lock();
+ try {
+ // As with put(), the size must be adjusted before the FlowFile is
made available on a partition
+ adjustSize(1, flowFile.getSize());
+
+ final QueuePartition partition = getPartition(flowFile);
+ partition.putToActive(flowFile);
+ } finally {
+ partitionReadLock.unlock();
+ }
+ }
+
@Override
public void receiveFromPeer(final Collection<FlowFileRecord> flowFiles)
throws IllegalClusterStateException {
partitionReadLock.lock();
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/QueuePartition.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/QueuePartition.java
index e60401c89d4..683dc0e4b78 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/QueuePartition.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/QueuePartition.java
@@ -60,6 +60,13 @@ public interface QueuePartition {
*/
void put(FlowFileRecord flowFile);
+ /**
+ * Adds the given FlowFile to the active portion of this partition, even
if the partition is currently swapping FlowFiles out.
+ * This is intended for FlowFiles that were previously polled and are
being returned, so that they stay ahead of any swapped-out FlowFiles.
+ * @param flowFile the FlowFile to add
+ */
+ void putToActive(FlowFileRecord flowFile);
+
/**
* Adds the given FlowFiles to this partition
* @param flowFiles the FlowFiles to add
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/RemoteQueuePartition.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/RemoteQueuePartition.java
index c0829a5ab83..71b52193150 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/RemoteQueuePartition.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/RemoteQueuePartition.java
@@ -139,6 +139,11 @@ public class RemoteQueuePartition implements
QueuePartition {
priorityQueue.put(flowFile);
}
+ @Override
+ public void putToActive(final FlowFileRecord flowFile) {
+ priorityQueue.putToActive(flowFile);
+ }
+
@Override
public void putAll(final Collection<FlowFileRecord> flowFiles) {
priorityQueue.putAll(flowFiles);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/StandardRebalancingPartition.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/StandardRebalancingPartition.java
index 5c819f5fcdf..2b45c5a0a9f 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/StandardRebalancingPartition.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/StandardRebalancingPartition.java
@@ -155,6 +155,12 @@ public class StandardRebalancingPartition implements
RebalancingPartition {
startRebalanceTask();
}
+ @Override
+ public void putToActive(final FlowFileRecord flowFile) {
+ queue.putToActive(flowFile);
+ startRebalanceTask();
+ }
+
@Override
public void putAll(final Collection<FlowFileRecord> flowFiles) {
queue.putAll(flowFiles);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/SwappablePriorityQueueLocalPartition.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/SwappablePriorityQueueLocalPartition.java
index e6a57bca32f..8d1b13dc6be 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/SwappablePriorityQueueLocalPartition.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/partition/SwappablePriorityQueueLocalPartition.java
@@ -97,6 +97,11 @@ public class SwappablePriorityQueueLocalPartition implements
LocalQueuePartition
priorityQueue.put(flowFile);
}
+ @Override
+ public void putToActive(final FlowFileRecord flowFile) {
+ priorityQueue.putToActive(flowFile);
+ }
+
@Override
public void putAll(final Collection<FlowFileRecord> flowFiles) {
priorityQueue.putAll(flowFiles);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java
index b133bf74c02..deae10e4ff2 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java
@@ -390,18 +390,19 @@ public class StatelessFlowTask {
// this.outputRepositoryRecords and this.cloneResults
if (!inputFlowFiles.isEmpty()) {
if (destinationPort == null) {
- // There is no destination port. Create a penalized FlowFile
and place it back on its original queue.
- // Do not penalize FlowFiles if the group was stopped, though.
That is not a true failure.
+ // There is no destination port. Create a penalized FlowFile
and place it back on the active portion of its original queue,
+ // so that it stays ahead of any FlowFiles that have been
swapped out. Do not penalize FlowFiles if the group was stopped, though.
+ // That is not a true failure.
for (final PolledFlowFile polledFlowFile : inputFlowFiles) {
if (stopped) {
-
polledFlowFile.getOriginalQueue().put(polledFlowFile.getInputFlowFile());
+
polledFlowFile.getOriginalQueue().putToActive(polledFlowFile.getInputFlowFile());
} else {
final FlowFileRecord newFile = new
StandardFlowFileRecord.Builder()
.fromFlowFile(polledFlowFile.getInputFlowFile())
.penaltyExpirationTime(System.currentTimeMillis()
+ 30_000L)
.build();
- polledFlowFile.getOriginalQueue().put(newFile);
+ polledFlowFile.getOriginalQueue().putToActive(newFile);
}
}
} else {
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 1b0514ef64e..da91edd8ce6 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
@@ -22,6 +22,7 @@ import org.apache.nifi.controller.MockSwapManager;
import org.apache.nifi.controller.queue.DropFlowFileAction;
import org.apache.nifi.controller.queue.DropFlowFileRequest;
import org.apache.nifi.controller.queue.FlowFileQueue;
+import org.apache.nifi.controller.queue.FlowFileQueueSize;
import org.apache.nifi.controller.queue.FlowFileQueueSnapshot;
import org.apache.nifi.controller.queue.PollStrategy;
import org.apache.nifi.controller.queue.QueueSize;
@@ -174,6 +175,18 @@ public class TestSwappablePriorityQueue {
queue.put(swapQueueFlowFile1);
}
+ // Poll the highest priority FlowFile from the active queue and return
it, as a session rollback would. It must go back to the
+ // active queue, ahead of the FlowFiles on the swap queue, even though
the queue is in swap mode.
+ final FlowFileRecord returnedFlowFile =
queue.poll(Collections.emptySet(), 0);
+ assertEquals("0", returnedFlowFile.getAttribute("i"));
+ queue.putToActive(returnedFlowFile);
+ queue.acknowledge(returnedFlowFile);
+
+ final FlowFileQueueSize queueSize = queue.getFlowFileQueueSize();
+ assertEquals(iterations, queueSize.getActiveCount());
+ assertEquals(3, queueSize.getSwappedCount());
+ assertEquals(0, queueSize.getUnacknowledgedCount());
+
// The first 10,000 should be ordered. Then all FlowFiles on the swap
queue should be transferred over, as a single unit, just as they would be in a
swap file.
for (int i = 0; i < iterations; i++) {
final FlowFileRecord flowFile = queue.poll(Collections.emptySet(),
0);
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
index 6baab772b82..3a04fc386ff 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionIT.java
@@ -1919,6 +1919,29 @@ public class StandardProcessSessionIT {
assertFalse(flowFileQueue.isActiveQueueEmpty());
}
+ @Test
+ public void testRollbackReturnsFlowFileToFrontOfSwappingQueue() {
+ // The queue's swap threshold is 10,000, so the last FlowFile queued
goes to the swap queue and the queue enters swap mode
+ final int flowFileCount = 10_001;
+ for (int i = 0; i < flowFileCount; i++) {
+ final FlowFileRecord flowFileRecord = new
StandardFlowFileRecord.Builder()
+ .id(i)
+ .addAttribute("uuid",
String.format("00000000-0000-0000-0000-%012d", i))
+ .entryDate(System.currentTimeMillis())
+ .size(0L)
+ .build();
+ flowFileQueue.put(flowFileRecord);
+ }
+
+ final FlowFile firstFlowFile = session.get();
+ assertEquals(0, firstFlowFile.getId());
+ session.rollback();
+
+ final FlowFile flowFileAfterRollback = session.get();
+ assertEquals(0, flowFileAfterRollback.getId());
+ assertEquals(flowFileCount, flowFileQueue.size().getObjectCount());
+ }
+
@Test
public void testCreateEmitted() throws IOException {
final FlowFile newFlowFile = session.create();
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/TestWriteAheadFlowFileRepository.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/TestWriteAheadFlowFileRepository.java
index 7033e93dc52..fee21a802d7 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/TestWriteAheadFlowFileRepository.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/TestWriteAheadFlowFileRepository.java
@@ -238,6 +238,10 @@ public class TestWriteAheadFlowFileRepository {
public void put(FlowFileRecord file) {
}
+ @Override
+ public void putToActive(FlowFileRecord file) {
+ }
+
@Override
public void putAll(Collection<FlowFileRecord> files) {
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java
index 59d3292c9d4..a66e726da53 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java
@@ -24,13 +24,18 @@ import org.apache.nifi.connectable.FlowFileActivity;
import org.apache.nifi.connectable.FlowFileTransferCounts;
import org.apache.nifi.connectable.Port;
import org.apache.nifi.controller.MockFlowFileRecord;
+import org.apache.nifi.controller.ProcessScheduler;
+import org.apache.nifi.controller.ScheduledState;
import org.apache.nifi.controller.metrics.ProcessSessionEvent;
import org.apache.nifi.controller.queue.FlowFileQueue;
+import org.apache.nifi.controller.queue.LocalQueuePartitionDiagnostics;
+import org.apache.nifi.controller.queue.StandardFlowFileQueue;
import org.apache.nifi.controller.repository.ContentRepository;
import
org.apache.nifi.controller.repository.DeferredStatelessContentRepository;
import org.apache.nifi.controller.repository.FlowFileEventRepository;
import org.apache.nifi.controller.repository.FlowFileRecord;
import org.apache.nifi.controller.repository.FlowFileRepository;
+import org.apache.nifi.controller.repository.FlowFileSwapManager;
import org.apache.nifi.controller.repository.RepositoryRecord;
import org.apache.nifi.controller.repository.RepositoryRecordType;
import org.apache.nifi.controller.repository.StandardContentRepositoryContext;
@@ -50,10 +55,13 @@ import org.apache.nifi.groups.ProcessGroup;
import org.apache.nifi.groups.StatelessGroupNode;
import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.mock.MockComponentLogger;
+import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.provenance.ProvenanceEventRecord;
import org.apache.nifi.provenance.ProvenanceEventRepository;
import org.apache.nifi.provenance.ProvenanceEventType;
import org.apache.nifi.provenance.StandardProvenanceEventRecord;
+import org.apache.nifi.stateless.flow.DataflowTrigger;
+import org.apache.nifi.stateless.flow.DataflowTriggerContext;
import org.apache.nifi.stateless.flow.StatelessDataflow;
import org.apache.nifi.stateless.flow.TriggerResult;
import org.apache.nifi.stateless.repository.ByteArrayContentRepository;
@@ -72,8 +80,11 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.OptionalLong;
import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -494,6 +505,58 @@ public class TestStatelessFlowTask {
assertEquals(2,
resourceClaimManager.getClaimantCount(flowFile1.getContentClaim().getResourceClaim()));
}
+ @Test
+ public void testFailureReturnsInputFlowFileToFrontOfSwappingQueue() throws
InterruptedException {
+ // The source queue's swap threshold is 3, so the fourth FlowFile goes
to the swap queue and the queue enters swap mode
+ final StandardFlowFileQueue sourceQueue = new
StandardFlowFileQueue("source-queue", flowFileRepository, provenanceRepository,
mock(ProcessScheduler.class),
+ mock(FlowFileSwapManager.class), null, 3, "0 sec", 0L, "0 B");
+ final List<FlowFileRecord> queuedFlowFiles = new ArrayList<>();
+ for (int i = 0; i < 4; i++) {
+ final FlowFileRecord flowFile = new MockFlowFileRecord();
+ queuedFlowFiles.add(flowFile);
+ sourceQueue.put(flowFile);
+ }
+
+ final Connection sourceConnection = mock(Connection.class);
+ when(sourceConnection.getFlowFileQueue()).thenReturn(sourceQueue);
+
when(inputPort1.getIncomingConnections()).thenReturn(List.of(sourceConnection));
+
+ final AtomicReference<ScheduledState> desiredState = new
AtomicReference<>();
+ when(statelessGroupNode.getDesiredState()).thenAnswer(invocation ->
desiredState.get());
+
+ // The group is stopped while the flow is running, so the input
FlowFile is returned to the source queue without a penalty
+ final TriggerResult canceledResult = mock(TriggerResult.class);
+ when(canceledResult.isCanceled()).thenReturn(true);
+ triggerFailingFlow(canceledResult, desiredState,
ScheduledState.STOPPED);
+
+ final FlowFileRecord flowFileAfterCancel = sourceQueue.poll(new
HashSet<>());
+ assertEquals(queuedFlowFiles.getFirst().getId(),
flowFileAfterCancel.getId());
+ sourceQueue.acknowledge(flowFileAfterCancel);
+
+ // The flow fails, so the input FlowFile is penalized and returned to
the active queue, ahead of the FlowFile on the swap queue
+ final TriggerResult failedResult = mock(TriggerResult.class);
+ when(failedResult.getFailureCause()).thenReturn(Optional.of(new
ProcessException("Intentional failure")));
+ triggerFailingFlow(failedResult, desiredState, ScheduledState.RUNNING);
+
+ final LocalQueuePartitionDiagnostics diagnostics =
sourceQueue.getQueueDiagnostics().getLocalQueuePartitionDiagnostics();
+ assertEquals(2, diagnostics.getActiveQueueSize().getObjectCount());
+ assertEquals(1, diagnostics.getSwapQueueSize().getObjectCount());
+ }
+
+ private void triggerFailingFlow(final TriggerResult triggerResult, final
AtomicReference<ScheduledState> desiredState, final ScheduledState
stateAfterTrigger) throws InterruptedException {
+ final DataflowTrigger dataflowTrigger = mock(DataflowTrigger.class);
+ when(dataflowTrigger.getResult(anyLong(),
any(TimeUnit.class))).thenReturn(Optional.of(triggerResult));
+
when(statelessFlow.trigger(any(DataflowTriggerContext.class))).thenAnswer(invocation
-> {
+ final DataflowTriggerContext triggerContext =
invocation.getArgument(0);
+
triggerContext.getFlowFileSupplier().getFlowFile(inputPort1.getName());
+ desiredState.set(stateAfterTrigger);
+ return dataflowTrigger;
+ });
+
+ desiredState.set(ScheduledState.RUNNING);
+ task.trigger();
+ }
+
private Connection mockConnection() {
final FlowFileQueue queue = mock(FlowFileQueue.class);
final Connection connection = mock(Connection.class);
diff --git
a/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/queue/StatelessFlowFileQueue.java
b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/queue/StatelessFlowFileQueue.java
index 9d6a231363b..1bb137dbdf5 100644
---
a/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/queue/StatelessFlowFileQueue.java
+++
b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/queue/StatelessFlowFileQueue.java
@@ -181,6 +181,12 @@ public class StatelessFlowFileQueue implements
DrainableFlowFileQueue {
totalBytes.addAndGet(flowFile.getSize());
}
+ @Override
+ public void putToActive(final FlowFileRecord flowFile) {
+ // This queue never swaps FlowFiles out, so every FlowFile is already
on the active queue
+ put(flowFile);
+ }
+
@Override
public void putAll(final Collection<FlowFileRecord> flowFiles) {
// Order the FlowFiles in the same order they were transferred by the
Processor. This ensures that we keep the ordering provided by the Processor.
