diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/BlockingSwappablePriorityQueue.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/BlockingSwappablePriorityQueue.java index db0e30e04140..f5c5d91393ca 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/BlockingSwappablePriorityQueue.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/BlockingSwappablePriorityQueue.java @@ -51,6 +51,15 @@ public void putAll(final Collection flowFiles) { } } + @Override + public void putBack(final Collection flowFiles) { + super.putBack(flowFiles); + + synchronized (monitor) { + monitor.notifyAll(); + } + } + public FlowFileRecord poll(final Set expiredRecords, final long expirationMillis, final long waitMillis, final PollStrategy pollStrategy) throws InterruptedException { final long maxTimestamp = System.currentTimeMillis() + waitMillis; 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..88d29fb374f1 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 @@ -572,6 +572,34 @@ public void putAll(final Collection flowFiles) { } } + public void putBack(final Collection 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 expiredRecords, final long expirationMillis) { return poll(expiredRecords, expirationMillis, PollStrategy.UNPENALIZED_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 f30157343e5f..763de3a27481 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 @@ -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 + "]"; @@ -131,28 +138,91 @@ public void setPriorities(final List 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); @@ -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); @@ -192,13 +262,24 @@ 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 { @@ -206,53 +287,67 @@ private class RebalanceTask implements Runnable { private final Set 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 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 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 toDistribute = new ArrayList<>(); + toDistribute.add(polled); - final List additionalRecords = queue.poll(999, expiredRecords, -1, PollStrategy.ALL_FLOWFILES); - toDistribute.addAll(additionalRecords); + final List 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); } } } 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..f7470786130e 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 @@ -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 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) -> { diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/partition/TestStandardRebalancingPartition.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/partition/TestStandardRebalancingPartition.java new file mode 100644 index 000000000000..5b1d2592ac88 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/partition/TestStandardRebalancingPartition.java @@ -0,0 +1,227 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.nifi.controller.queue.clustered.partition; + +import org.apache.nifi.controller.MockFlowFileRecord; +import org.apache.nifi.controller.MockSwapManager; +import org.apache.nifi.controller.queue.BlockingSwappablePriorityQueue; +import org.apache.nifi.controller.queue.LoadBalancedFlowFileQueue; +import org.apache.nifi.controller.queue.PollStrategy; +import org.apache.nifi.controller.queue.QueueSize; +import org.apache.nifi.controller.repository.FlowFileRecord; +import org.apache.nifi.events.EventReporter; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class TestStandardRebalancingPartition { + private LoadBalancedFlowFileQueue flowFileQueue; + private ControllableBlockingQueue queue; + private StandardRebalancingPartition partition; + + @BeforeEach + void setUp() { + flowFileQueue = mock(LoadBalancedFlowFileQueue.class); + when(flowFileQueue.getIdentifier()).thenReturn("unit-test"); + queue = new ControllableBlockingQueue(flowFileQueue); + partition = new StandardRebalancingPartition(queue, flowFileQueue); + } + + @Test + @Timeout(10) + void testStopInterruptsBlockedPollAndWaitsForWorker() throws Exception { + final BlockingPollQueue blockingQueue = new BlockingPollQueue(flowFileQueue); + final StandardRebalancingPartition blockingPartition = new StandardRebalancingPartition(blockingQueue, flowFileQueue); + blockingPartition.start(mock(FlowFilePartitioner.class)); + assertTrue(blockingQueue.awaitPoll()); + + try (final ExecutorService executor = Executors.newSingleThreadExecutor()) { + final Future stopFuture = executor.submit(blockingPartition::stop); + assertTrue(blockingQueue.awaitStopInterrupt()); + assertFalse(stopFuture.isDone()); + + blockingQueue.releasePoll(); + stopFuture.get(5, TimeUnit.SECONDS); + } + + verify(flowFileQueue, never()).distributeToPartitions(org.mockito.ArgumentMatchers.anyCollection()); + } + + @Test + @Timeout(10) + void testStopRestoresPolledFlowFilesBeforeRestart() throws Exception { + final FlowFileRecord flowFile = new MockFlowFileRecord(10L); + partition.rebalance(List.of(flowFile)); + partition.start(mock(FlowFilePartitioner.class)); + + assertTrue(queue.awaitBatchPoll()); + + try (final ExecutorService executor = Executors.newSingleThreadExecutor()) { + final Future stopFuture = executor.submit(partition::stop); + assertTrue(queue.awaitStopInterrupt()); + queue.releaseBatchPoll(); + stopFuture.get(5, TimeUnit.SECONDS); + } + + verify(flowFileQueue, never()).distributeToPartitions(org.mockito.ArgumentMatchers.anyCollection()); + assertEquals(new QueueSize(1, 10L), partition.size()); + + partition.start(mock(FlowFilePartitioner.class)); + verify(flowFileQueue, timeout(5000).times(1)).distributeToPartitions(List.of(flowFile)); + partition.stop(); + assertEquals(new QueueSize(0, 0L), partition.size()); + } + + @Test + @Timeout(10) + void testStartWaitsForStopToComplete() throws Exception { + final FlowFileRecord flowFile = new MockFlowFileRecord(10L); + partition.rebalance(List.of(flowFile)); + partition.start(mock(FlowFilePartitioner.class)); + assertTrue(queue.awaitBatchPoll()); + + try (final ExecutorService executor = Executors.newFixedThreadPool(2)) { + final Future stopFuture = executor.submit(partition::stop); + assertTrue(queue.awaitStopInterrupt()); + + final CountDownLatch startInvoked = new CountDownLatch(1); + final Future startFuture = executor.submit(() -> { + startInvoked.countDown(); + partition.start(mock(FlowFilePartitioner.class)); + }); + assertTrue(startInvoked.await(5, TimeUnit.SECONDS)); + verify(flowFileQueue, never()).distributeToPartitions(org.mockito.ArgumentMatchers.anyCollection()); + assertFalse(startFuture.isDone()); + + queue.releaseBatchPoll(); + stopFuture.get(5, TimeUnit.SECONDS); + startFuture.get(5, TimeUnit.SECONDS); + } + + verify(flowFileQueue, timeout(5000).times(1)).distributeToPartitions(List.of(flowFile)); + partition.stop(); + assertEquals(new QueueSize(0, 0L), partition.size()); + } + + private static class ControllableBlockingQueue extends BlockingSwappablePriorityQueue { + private final CountDownLatch batchPollEntered = new CountDownLatch(1); + private final CountDownLatch stopInterruptObserved = new CountDownLatch(1); + private final CountDownLatch releaseBatchPoll = new CountDownLatch(1); + private final AtomicBoolean blockBatchPoll = new AtomicBoolean(true); + + ControllableBlockingQueue(final LoadBalancedFlowFileQueue flowFileQueue) { + super(new MockSwapManager(), 10000, EventReporter.NO_OP, flowFileQueue, (flowFiles, requestor) -> new QueueSize(0, 0L), "rebalance"); + } + + @Override + public List poll(final int maxResults, final Set expiredRecords, + final long expirationMillis, final PollStrategy pollStrategy) { + if (blockBatchPoll.compareAndSet(true, false)) { + batchPollEntered.countDown(); + boolean interrupted = false; + while (true) { + try { + releaseBatchPoll.await(); + break; + } catch (final InterruptedException e) { + interrupted = true; + stopInterruptObserved.countDown(); + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + return Collections.emptyList(); + } + + boolean awaitBatchPoll() throws InterruptedException { + return batchPollEntered.await(5, TimeUnit.SECONDS); + } + + boolean awaitStopInterrupt() throws InterruptedException { + return stopInterruptObserved.await(5, TimeUnit.SECONDS); + } + + void releaseBatchPoll() { + releaseBatchPoll.countDown(); + } + } + + private static class BlockingPollQueue extends BlockingSwappablePriorityQueue { + private final CountDownLatch pollEntered = new CountDownLatch(1); + private final CountDownLatch stopInterruptObserved = new CountDownLatch(1); + private final CountDownLatch releasePoll = new CountDownLatch(1); + + BlockingPollQueue(final LoadBalancedFlowFileQueue flowFileQueue) { + super(new MockSwapManager(), 10000, EventReporter.NO_OP, flowFileQueue, (flowFiles, requestor) -> new QueueSize(0, 0L), "rebalance"); + } + + @Override + public FlowFileRecord poll(final Set expiredRecords, final long expirationMillis, + final long waitMillis, final PollStrategy pollStrategy) throws InterruptedException { + pollEntered.countDown(); + boolean interrupted = false; + while (true) { + try { + releasePoll.await(); + break; + } catch (final InterruptedException e) { + interrupted = true; + stopInterruptObserved.countDown(); + } + } + if (interrupted) { + throw new InterruptedException(); + } + return null; + } + + boolean awaitPoll() throws InterruptedException { + return pollEntered.await(5, TimeUnit.SECONDS); + } + + boolean awaitStopInterrupt() throws InterruptedException { + return stopInterruptObserved.await(5, TimeUnit.SECONDS); + } + + void releasePoll() { + releasePoll.countDown(); + } + } +}