Skip to content

NIFI-16180 - Ensure Rebalancing Worker Stops Before Returning #11522

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pvillard31 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

NIFI-16180 - Ensure Rebalancing Worker Stops Before Returning #11522

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions ...k-core/src/main/java/org/apache/nifi/controller/queue/BlockingSwappablePriorityQueue.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,15 @@ public void putAll(final Collection<FlowFileRecord> flowFiles) {
}
}

@Override
public void putBack(final Collection<FlowFileRecord> flowFiles) {
super.putBack(flowFiles);

synchronized (monitor) {
monitor.notifyAll();
}
}

public FlowFileRecord poll(final Set<FlowFileRecord> expiredRecords, final long expirationMillis, final long waitMillis, final PollStrategy pollStrategy) throws InterruptedException {
final long maxTimestamp = System.currentTimeMillis() + waitMillis;

Expand Down
28 changes: 28 additions & 0 deletions ...framework-core/src/main/java/org/apache/nifi/controller/queue/SwappablePriorityQueue.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,34 @@ public void putAll(final Collection<FlowFileRecord> flowFiles) {
}
}

public void putBack(final Collection<FlowFileRecord> flowFiles) {
final int count = flowFiles.size();
final long bytes = flowFiles.stream().mapToLong(FlowFileRecord::getSize).sum();

writeLock.lock();
try {
activeQueue.addAll(flowFiles);

boolean updated;
do {
final FlowFileQueueSize currentSize = getFlowFileQueueSize();
final FlowFileQueueSize updatedSize = new FlowFileQueueSize(
currentSize.getActiveCount() + count,
currentSize.getActiveBytes() + bytes,
currentSize.getSwappedCount(),
currentSize.getSwappedBytes(),
currentSize.getSwapFileCount(),
currentSize.getUnacknowledgedCount() - count,
currentSize.getUnacknowledgedBytes() - bytes);
updated = updateSize(currentSize, updatedSize);
} while (!updated);

updateTopPenaltyExpiration();
} finally {
writeLock.unlock("putBack");
}
}

public FlowFileRecord poll(final Set<FlowFileRecord> expiredRecords, final long expirationMillis) {
return poll(expiredRecords, expirationMillis, PollStrategy.UNPENALIZED_FLOWFILES);
}
Expand Down
201 changes: 148 additions & 53 deletions ...va/org/apache/nifi/controller/queue/clustered/partition/StandardRebalancingPartition.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,20 @@ public class StandardRebalancingPartition implements RebalancingPartition {
private final LoadBalancedFlowFileQueue flowFileQueue;
private final String description;

private final Object lifecycleMonitor = new Object();

private volatile boolean stopped = true;
private RebalanceTask rebalanceTask;
private Thread rebalanceThread;
private boolean stopping;

public StandardRebalancingPartition(final FlowFileSwapManager swapManager, final int swapThreshold, final EventReporter eventReporter,
final LoadBalancedFlowFileQueue flowFileQueue, final DropFlowFileAction dropAction) {
this(new BlockingSwappablePriorityQueue(swapManager, swapThreshold, eventReporter, flowFileQueue, dropAction, SWAP_PARTITION_NAME), flowFileQueue);
}

this.queue = new BlockingSwappablePriorityQueue(swapManager, swapThreshold, eventReporter, flowFileQueue, dropAction, SWAP_PARTITION_NAME);
StandardRebalancingPartition(final BlockingSwappablePriorityQueue queue, final LoadBalancedFlowFileQueue flowFileQueue) {
this.queue = queue;
this.queueIdentifier = flowFileQueue.getIdentifier();
this.flowFileQueue = flowFileQueue;
this.description = "RebalancingPartition[queueId=" + queueIdentifier + "]";
Expand Down Expand Up @@ -131,28 +138,91 @@ public void setPriorities(final List<FlowFilePrioritizer> newPriorities) {
}

@Override
public synchronized void start(final FlowFilePartitioner partitionerUsed) {
stopped = false;
rebalanceFromQueue();
public void start(final FlowFilePartitioner partitionerUsed) {
boolean interrupted = false;
synchronized (lifecycleMonitor) {
while (stopping) {
try {
lifecycleMonitor.wait();
} catch (final InterruptedException e) {
interrupted = true;
}
}

stopped = false;
startRebalanceTask();
}

if (interrupted) {
Thread.currentThread().interrupt();
}
}

@Override
public synchronized void stop() {
stopped = true;
public void stop() {
final RebalanceTask task;
final Thread thread;
boolean interrupted = false;

synchronized (lifecycleMonitor) {
while (stopping) {
try {
lifecycleMonitor.wait();
} catch (final InterruptedException e) {
interrupted = true;
}
}

if (this.rebalanceTask != null) {
this.rebalanceTask.stop();
stopped = true;
task = rebalanceTask;
thread = rebalanceThread;
if (task == null || thread == null) {
if (interrupted) {
Thread.currentThread().interrupt();
}
return;
}

stopping = true;
}

this.rebalanceTask = null;
task.stop();
thread.interrupt();

while (thread.isAlive()) {
try {
thread.join();
} catch (final InterruptedException e) {
interrupted = true;
}
}

synchronized (lifecycleMonitor) {
if (rebalanceTask == task) {
rebalanceTask = null;
rebalanceThread = null;
}
stopping = false;
lifecycleMonitor.notifyAll();
}

if (interrupted) {
Thread.currentThread().interrupt();
}
}

private synchronized void rebalanceFromQueue() {
if (stopped) {
logger.debug("Will not rebalance from queue because {} is stopped", this);
return;
private void rebalanceFromQueue() {
synchronized (lifecycleMonitor) {
if (stopped) {
logger.debug("Will not rebalance from queue because {} is stopped", this);
return;
}

startRebalanceTask();
}
}

private void startRebalanceTask() {
// If a task is already defined, do nothing. There's already a thread running.
if (rebalanceTask != null) {
logger.debug("Rebalance Task already exists for {}", this);
Expand All @@ -161,7 +231,7 @@ private synchronized void rebalanceFromQueue() {

this.rebalanceTask = new RebalanceTask();

final Thread rebalanceThread = new Thread(this.rebalanceTask);
rebalanceThread = new Thread(this.rebalanceTask);
rebalanceThread.setName("Rebalance queued data for Connection " + queueIdentifier);
rebalanceThread.start();
logger.debug("No Rebalance Task currently exists for {}. Starting new Rebalance Thread {}", this, rebalanceThread);
Expand Down Expand Up @@ -192,67 +262,92 @@ public FlowFileQueueContents packageForRebalance(String newPartitionName) {
return queue.packageForRebalance(newPartitionName);
}

private synchronized boolean isComplete() {
if (!queue.isEmpty()) {
return false;
private boolean isQueueEmpty() {
synchronized (lifecycleMonitor) {
return queue.isEmpty();
}
}

private void taskCompleted(final RebalanceTask task) {
synchronized (lifecycleMonitor) {
if (rebalanceTask == task) {
rebalanceTask = null;
rebalanceThread = null;
}

this.rebalanceTask = null;
return true;
if (!stopped && !queue.isEmpty()) {
startRebalanceTask();
}
lifecycleMonitor.notifyAll();
}
}

private class RebalanceTask implements Runnable {
private volatile boolean stopped = false;
private final Set<FlowFileRecord> expiredRecords = new HashSet<>();
private final long pollWaitMillis = 100L;

public void stop() {
public synchronized void stop() {
stopped = true;
}

@Override
public void run() {
while (!stopped) {
final FlowFileRecord polled;
private synchronized boolean distribute(final Collection<FlowFileRecord> flowFiles) {
if (stopped) {
return false;
}

expiredRecords.clear();
flowFileQueue.distributeToPartitions(flowFiles);
queue.acknowledge(flowFiles);
return true;
}

// Wait up to #pollWaitMillis milliseconds to get a FlowFile. If none, then check if stopped
// and if not, poll again.
try {
polled = queue.poll(expiredRecords, -1, pollWaitMillis, PollStrategy.ALL_FLOWFILES);
} catch (final InterruptedException ie) {
Thread.currentThread().interrupt();
continue;
}
@Override
public void run() {
try {
while (!stopped) {
final FlowFileRecord polled;

expiredRecords.clear();

// Wait up to #pollWaitMillis milliseconds to get a FlowFile. If none, then check if stopped
// and if not, poll again.
try {
polled = queue.poll(expiredRecords, -1, pollWaitMillis, PollStrategy.ALL_FLOWFILES);
} catch (final InterruptedException ie) {
flowFileQueue.handleExpiredRecords(expiredRecords);
Thread.currentThread().interrupt();
return;
}

if (polled == null) {
flowFileQueue.handleExpiredRecords(expiredRecords);
if (polled == null) {
flowFileQueue.handleExpiredRecords(expiredRecords);

if (isComplete()) {
logger.debug("Rebalance Task completed for {}", this);
return;
} else {
continue;
if (isQueueEmpty()) {
logger.debug("Rebalance Task completed for {}", this);
return;
} else {
continue;
}
}
}

// We got 1 FlowFile. Try a second poll to obtain up to 999 more (for a total of 1,000).
final List<FlowFileRecord> toDistribute = new ArrayList<>();
toDistribute.add(polled);
// We got 1 FlowFile. Try a second poll to obtain up to 999 more (for a total of 1,000).
final List<FlowFileRecord> toDistribute = new ArrayList<>();
toDistribute.add(polled);

final List<FlowFileRecord> additionalRecords = queue.poll(999, expiredRecords, -1, PollStrategy.ALL_FLOWFILES);
toDistribute.addAll(additionalRecords);
final List<FlowFileRecord> additionalRecords = queue.poll(999, expiredRecords, -1, PollStrategy.ALL_FLOWFILES);
toDistribute.addAll(additionalRecords);

flowFileQueue.handleExpiredRecords(expiredRecords);
flowFileQueue.handleExpiredRecords(expiredRecords);

logger.debug("{} Rebalancing {}", this, toDistribute);
logger.debug("{} Rebalancing {}", this, toDistribute);

// Transfer all of the FlowFiles that we got back to the FlowFileQueue itself. This will cause the data to be
// re-partitioned and binned appropriately. We also then need to ensure that we acknowledge the data from our
// own SwappablePriorityQueue to ensure that the sizes are kept in check.
flowFileQueue.distributeToPartitions(toDistribute);
queue.acknowledge(toDistribute);
if (!distribute(toDistribute)) {
queue.putBack(toDistribute);
return;
}
}
} finally {
taskCompleted(this);
}
}
}
Expand Down
17 changes: 17 additions & 0 deletions .../src/test/java/org/apache/nifi/controller/queue/clustered/TestSwappablePriorityQueue.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,23 @@ public void setup() {
queue = new SwappablePriorityQueue(swapManager, 10000, eventReporter, flowFileQueue, dropAction, "local");
}

@Test
public void testPutBackRestoresPolledFlowFiles() {
final FlowFileRecord first = new MockFlowFileRecord(10L);
final FlowFileRecord second = new MockFlowFileRecord(20L);
queue.putAll(List.of(first, second));

final List<FlowFileRecord> polled = queue.poll(2, Collections.emptySet(), -1, PollStrategy.ALL_FLOWFILES);
assertEquals(new QueueSize(2, 30L), queue.size());
assertEquals(2, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());

queue.putBack(polled);

assertEquals(new QueueSize(2, 30L), queue.size());
assertEquals(0, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());
assertEquals(2, queue.getActiveFlowFiles().size());
}

@Test
public void testPrioritizersBigQueue() {
final FlowFilePrioritizer iAttributePrioritizer = (o1, o2) -> {
Expand Down
Loading
Loading