/*
 * 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;

import org.apache.nifi.controller.MockFlowFileRecord;
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.PollStrategy;
import org.apache.nifi.controller.queue.QueueSize;
import org.apache.nifi.controller.queue.SwappablePriorityQueue;
import org.apache.nifi.controller.repository.FlowFileRecord;
import org.apache.nifi.events.EventReporter;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.flowfile.FlowFilePrioritizer;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.StringUtils;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.mockito.Mockito;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class TestSwappablePriorityQueue {

    private MockSwapManager swapManager;
    private final List<String> events = new ArrayList<>();
    private EventReporter eventReporter;

    private final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    private final DropFlowFileAction dropAction = (flowFiles, requestor) -> {
        return new QueueSize(flowFiles.size(), flowFiles.stream().mapToLong(FlowFileRecord::getSize).sum());
    };

    private SwappablePriorityQueue queue;

    @BeforeEach
    public void setup() {
        swapManager = new MockSwapManager();

        events.clear();
        eventReporter = (severity, category, message) -> events.add(message);

        when(flowFileQueue.getIdentifier()).thenReturn("unit-test");
        queue = new SwappablePriorityQueue(swapManager, 10000, eventReporter, flowFileQueue, dropAction, "local");
    }

    @Test
    public void testPrioritizersBigQueue() {
        final FlowFilePrioritizer iAttributePrioritizer = new FlowFilePrioritizer() {
            @Override
            public int compare(final FlowFile o1, final FlowFile o2) {
                final int i1 = Integer.parseInt(o1.getAttribute("i"));
                final int i2 = Integer.parseInt(o2.getAttribute("i"));
                return Integer.compare(i1, i2);
            }
        };

        queue.setPriorities(Collections.singletonList(iAttributePrioritizer));
        final int iterations = 29000;

        for (int i=0; i < iterations; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i)));
            queue.put(flowFile);
        }

        for (int i=0; i < iterations; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i + iterations);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i + iterations)));

            final FlowFileRecord polled = queue.poll(Collections.emptySet(), 0L);
            assertEquals(polled.getAttribute("i"), String.valueOf(i));

            queue.put(flowFile);
        }

        // Make sure that the data is pulled from the queue and added back a couple of times.
        // This will trigger swapping to occur, but also leave a lot of data in memory on the queue.
        // This specifically tests the edge case where data is swapped out, and we want to make sure that
        // when we read from the queue, that we swap the data back in before processing anything on the
        // pending 'swap queue' internally.
        repopulateQueue();
        repopulateQueue();

        int i=iterations;
        FlowFileRecord flowFile;
        while ((flowFile = queue.poll(Collections.emptySet(), 0)) != null) {
            assertEquals(String.valueOf(i), flowFile.getAttribute("i"));
            i++;
        }
    }


    @Test
    public void testOrderingWithCornerCases() {
        final FlowFilePrioritizer iAttributePrioritizer = new FlowFilePrioritizer() {
            @Override
            public int compare(final FlowFile o1, final FlowFile o2) {
                final int i1 = Integer.parseInt(o1.getAttribute("i"));
                final int i2 = Integer.parseInt(o2.getAttribute("i"));
                return Integer.compare(i1, i2);
            }
        };

        queue.setPriorities(Collections.singletonList(iAttributePrioritizer));

        for (final int queueSize : new int[] {1, 9999, 10_000, 10_001, 19_999, 20_000, 20_001}) {
            System.out.println("Queue Size: " + queueSize);

            for (int i=0; i < queueSize; i++) {
                final MockFlowFile flowFile = new MockFlowFile(i);
                flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i)));
                queue.put(flowFile);
            }

            for (int i=0; i < queueSize; i++) {
                final FlowFileRecord flowFile = queue.poll(Collections.emptySet(), 0);
                assertEquals(String.valueOf(i), flowFile.getAttribute("i"));
            }

            assertNull(queue.poll(Collections.emptySet(), 0));
        }
    }

    @Test
    public void testPrioritizerWhenOutOfOrderDataEntersSwapQueue() {
        final FlowFilePrioritizer iAttributePrioritizer = new FlowFilePrioritizer() {
            @Override
            public int compare(final FlowFile o1, final FlowFile o2) {
                final int i1 = Integer.parseInt(o1.getAttribute("i"));
                final int i2 = Integer.parseInt(o2.getAttribute("i"));
                return Integer.compare(i1, i2);
            }
        };

        queue.setPriorities(Collections.singletonList(iAttributePrioritizer));

        // Add 10,000 FlowFiles to the queue. These will all go to the active queue.
        final int iterations = 10000;
        for (int i=0; i < iterations; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i)));
            queue.put(flowFile);
        }

        // Added 3 FlowFiles to the queue. These will all go to the Swap Queue.
        for (final String iValue : new String[] {"10000", "-5", "8000"}) {
            final MockFlowFile swapQueueFlowFile1 = new MockFlowFile(10_000);
            swapQueueFlowFile1.putAttributes(Collections.singletonMap("i", iValue));
            queue.put(swapQueueFlowFile1);
        }

        // 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);
            assertEquals(String.valueOf(i), flowFile.getAttribute("i"));
        }

        for (final String iValue : new String[] {"-5", "8000", "10000"}) {
            final FlowFileRecord flowFile = queue.poll(Collections.emptySet(), 0);
            assertEquals(iValue, flowFile.getAttribute("i"));
        }
    }

    @Test
    public void testPrioritizersDataAddedAfterSwapOccurs() {
        final FlowFilePrioritizer iAttributePrioritizer = new FlowFilePrioritizer() {
            @Override
            public int compare(final FlowFile o1, final FlowFile o2) {
                final int i1 = Integer.parseInt(o1.getAttribute("i"));
                final int i2 = Integer.parseInt(o2.getAttribute("i"));
                return Integer.compare(i1, i2);
            }
        };

        queue.setPriorities(Collections.singletonList(iAttributePrioritizer));
        final int iterations = 29000;

        for (int i=0; i < iterations; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i)));
            queue.put(flowFile);
        }

        for (int i=0; i < iterations; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i + iterations);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(i + iterations)));

            final FlowFileRecord polled = queue.poll(Collections.emptySet(), 0L);
            assertEquals(polled.getAttribute("i"), String.valueOf(i));

            queue.put(flowFile);
        }

        // Make sure that the data is pulled from the queue and added back a couple of times.
        // This will trigger swapping to occur, but also leave a lot of data in memory on the queue.
        // This specifically tests the edge case where data is swapped out, and we want to make sure that
        // when we read from the queue, that we swap the data back in before processing anything on the
        // pending 'swap queue' internally.
        repopulateQueue();
        repopulateQueue();

        // Add enough data for another swap file to get created.
        final int baseI = iterations * 2;
        for (int i=0; i < 10_000; i++) {
            final MockFlowFile flowFile = new MockFlowFile(i);
            flowFile.putAttributes(Collections.singletonMap("i", String.valueOf(baseI + i)));
            queue.put(flowFile);
        }

        repopulateQueue();

        int i=iterations;
        FlowFileRecord flowFile;
        while ((flowFile = queue.poll(Collections.emptySet(), 0)) != null) {
            assertEquals(String.valueOf(i), flowFile.getAttribute("i"));
            i++;
        }
    }

    private void repopulateQueue() {
        final List<String> attrs = new ArrayList<>();
        final List<FlowFileRecord> ffs = new ArrayList<>();
        FlowFileRecord ff;
        while ((ff = queue.poll(Collections.emptySet(), 0L)) != null) {
            ffs.add(ff);
            attrs.add(ff.getAttribute("i"));
        }

        ffs.forEach(queue::put);
        System.out.println(StringUtils.join(attrs, ", "));
    }


    @Test
    public void testSwapOutFailureLeavesCorrectQueueSize() {
        swapManager.setSwapOutFailureOnNthIteration(1, null);

        for (int i = 0; i < 19999; i++) {
            queue.put(new MockFlowFile(i));
        }

        assertEquals(19999, queue.size().getObjectCount());
        assertEquals(0, events.size());

        queue.put(new MockFlowFile(20000));
        assertEquals(20000, queue.size().getObjectCount());
        assertEquals(1, events.size()); // Expect a single failure event to be emitted
    }

    @Test
    public void testPrioritizer() {
        final FlowFilePrioritizer prioritizer = (o1, o2) -> Long.compare(o1.getId(), o2.getId());
        queue.setPriorities(Collections.singletonList(prioritizer));

        for (int i = 0; i < 5000; i++) {
            queue.put(new MockFlowFile(i));
        }

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        for (int i = 0; i < 5000; i++) {
            final FlowFileRecord polled = queue.poll(expiredRecords, 500000L);
            assertEquals(i, polled.getId());
        }

        // We can add flowfiles in reverse order (highest ID first) and we should still get the same order back when polling
        for (int i = 0; i < 5000; i++) {
            queue.put(new MockFlowFile(5000 - i));
        }
        for (int i = 0; i < 5000; i++) {
            final FlowFileRecord polled = queue.poll(expiredRecords, 500000L);
            // ID's will start at 1, since the last FlowFile added will have ID of 5000 - 4999
            assertEquals(i + 1, polled.getId());
        }

        // Add FlowFiles again, then change prioritizer and ensure that the order is updated
        for (int i = 0; i < 5000; i++) {
            queue.put(new MockFlowFile(i));
        }

        final FlowFilePrioritizer reversePrioritizer = (o1, o2) -> Long.compare(o2.getId(), o1.getId());
        queue.setPriorities(Collections.singletonList(reversePrioritizer));

        for (int i = 0; i < 5000; i++) {
            final FlowFileRecord polled = queue.poll(expiredRecords, 500000L);
            // ID's will start at 4999, since the last FlowFile added will have ID of 4999 (i < 5000, not i <= 5000).
            assertEquals(5000 - i - 1, polled.getId());
        }
    }

    @Test
    public void testPollWithPenalizedFlowFile() {
        final FlowFileRecord penalizedFlowFile = mock(FlowFileRecord.class);
        when(penalizedFlowFile.isPenalized()).thenReturn(true);
        assertTrue(queue.isEmpty());
        queue.put(penalizedFlowFile);

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        FlowFileRecord polled = queue.poll(expiredRecords, 0, PollStrategy.UNPENALIZED_FLOWFILES);
        assertNull(polled);

        assertFalse(queue.isEmpty());

        polled = queue.poll(expiredRecords, 0, PollStrategy.ALL_FLOWFILES);
        assertNotNull(polled);
        assertSame(penalizedFlowFile, polled);

        // queue is still not empty because FlowFile has not yet been acknowledged.
        queue.acknowledge(polled);

        assertTrue(queue.isEmpty());
    }

    @Test
    public void testPollWithOnlyExpiredFlowFile() {
        final FlowFileRecord expiredFlowFile = mock(FlowFileRecord.class);
        when(expiredFlowFile.getEntryDate()).thenReturn(System.currentTimeMillis() - 5000L);
        queue.put(expiredFlowFile);

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        final FlowFileRecord polled = queue.poll(expiredRecords, 4999);
        assertNull(polled);

        assertEquals(1, expiredRecords.size());
        final FlowFileRecord expired = expiredRecords.iterator().next();
        assertSame(expiredFlowFile, expired);
    }

    @Test
    public void testPollWithExpiredAndUnexpired() {
        final SwappablePriorityQueue queue = new SwappablePriorityQueue(swapManager, 100, eventReporter, flowFileQueue, dropAction, "local");

        final FlowFileRecord expiredFlowFile = mock(FlowFileRecord.class);
        when(expiredFlowFile.getEntryDate()).thenReturn(System.currentTimeMillis() - 5000L);
        queue.put(expiredFlowFile);

        final FlowFileRecord unexpiredFlowFile = mock(FlowFileRecord.class);
        when(unexpiredFlowFile.getEntryDate()).thenReturn(System.currentTimeMillis() + 500000L);
        queue.put(unexpiredFlowFile);

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        final FlowFileRecord polled = queue.poll(expiredRecords, 4999);
        assertSame(unexpiredFlowFile, polled);

        assertEquals(1, expiredRecords.size());
        final FlowFileRecord expired = expiredRecords.iterator().next();
        assertSame(expiredFlowFile, expired);
    }

    @Test
    public void testEmpty() {
        assertTrue(queue.isEmpty());
        assertTrue(queue.isActiveQueueEmpty());

        for (int i = 0; i < 9; i++) {
            queue.put(new MockFlowFileRecord());
            assertFalse(queue.isEmpty());
            assertFalse(queue.isActiveQueueEmpty());
        }

        queue.put(new MockFlowFileRecord());
        assertFalse(queue.isEmpty());
        assertFalse(queue.isActiveQueueEmpty());

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        final FlowFileRecord polled = queue.poll(expiredRecords, 500000);
        assertNotNull(polled);
        assertTrue(expiredRecords.isEmpty());

        assertFalse(queue.isEmpty());
        assertFalse(queue.isActiveQueueEmpty());

        // queue is still full because FlowFile has not yet been acknowledged.
        queue.acknowledge(polled);

        // FlowFile has been acknowledged; queue should no longer be full.
        assertFalse(queue.isEmpty());
        assertFalse(queue.isActiveQueueEmpty());
    }

    @Test
    public void testSwapOutOccurs() {
        for (int i = 0; i < 10000; i++) {
            queue.put(new MockFlowFileRecord());
            assertEquals(0, swapManager.swapOutCalledCount);
            assertEquals(i + 1, queue.size().getObjectCount());
            assertEquals(i + 1, queue.size().getByteCount());
        }

        for (int i = 0; i < 9999; i++) {
            queue.put(new MockFlowFileRecord());
            assertEquals(0, swapManager.swapOutCalledCount);
            assertEquals(i + 10001, queue.size().getObjectCount());
            assertEquals(i + 10001, queue.size().getByteCount());
        }

        queue.put(new MockFlowFileRecord(1000));
        assertEquals(1, swapManager.swapOutCalledCount);
        assertEquals(20000, queue.size().getObjectCount());
        assertEquals(20999, queue.size().getByteCount());

        assertEquals(10000, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());
    }

    @Test
    public void testLowestPrioritySwappedOutFirst() {
        final List<FlowFilePrioritizer> prioritizers = new ArrayList<>();
        prioritizers.add((o1, o2) -> Long.compare(o1.getSize(), o2.getSize()));
        queue.setPriorities(prioritizers);

        long maxSize = 20000;
        for (int i = 1; i <= 20000; i++) {
            queue.put(new MockFlowFileRecord(maxSize - i));
        }

        assertEquals(1, swapManager.swapOutCalledCount);
        assertEquals(20000, queue.size().getObjectCount());

        assertEquals(10000, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        // The first 10,000 FlowFiles to be added to the queue will be sorted by size (first 10,000 because that's the swap threshold, by size because of the prioritizer).
        // The next 10,000 spill over to the swap queue. So we expect the first 10,000 FlowFiles to be size 10,000 to 20,000. Then the next 10,000 to be sized 0 to 9,999.
        final List<FlowFileRecord> firstBatch = queue.poll(Integer.MAX_VALUE, Collections.emptySet(), 0);
        assertEquals(10000, firstBatch.size());
        for (int i = 0; i < 10000; i++) {
            assertEquals(10_000 + i, firstBatch.get(i).getSize());
        }

        final List<FlowFileRecord> secondBatch = queue.poll(Integer.MAX_VALUE, Collections.emptySet(), 0);
        assertEquals(10000, secondBatch.size());
        for (int i = 0; i < 10000; i++) {
            assertEquals(i, secondBatch.get(i).getSize());
        }

    }

    @Test
    public void testPrioritiesKeptIntactBeforeSwap() {
        final List<FlowFilePrioritizer> prioritizers = new ArrayList<>();
        prioritizers.add((o1, o2) -> Long.compare(o1.getSize(), o2.getSize()));
        queue.setPriorities(prioritizers);

        int maxSize = 9999;
        for (int i = 1; i <= maxSize; i++) {
            queue.put(new MockFlowFileRecord(maxSize - i));
        }

        assertEquals(0, swapManager.swapOutCalledCount);
        assertEquals(maxSize, queue.size().getObjectCount());

        assertEquals(9999, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        FlowFileRecord flowFile;
        int i=0;
        while ((flowFile = queue.poll(Collections.emptySet(), 0L)) != null) {
            assertEquals(i++, flowFile.getSize());
        }
    }

    @Test
    public void testSwapIn() {
        for (int i = 1; i <= 20000; i++) {
            queue.put(new MockFlowFileRecord());
        }

        assertEquals(1, swapManager.swappedOut.size());
        queue.put(new MockFlowFileRecord());
        assertEquals(1, swapManager.swappedOut.size());

        final Set<FlowFileRecord> exp = new HashSet<>();
        for (int i = 0; i < 9999; i++) {
            final FlowFileRecord flowFile = queue.poll(exp, 500000);
            assertNotNull(flowFile);
            assertEquals(1, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());
            assertEquals(1, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getByteCount());

            queue.acknowledge(Collections.singleton(flowFile));
            assertEquals(0, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());
            assertEquals(0, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getByteCount());
        }

        assertEquals(0, swapManager.swapInCalledCount);
        assertEquals(1, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());
        assertNotNull(queue.poll(exp, 500000));

        assertEquals(0, swapManager.swapInCalledCount);
        assertEquals(0, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        assertEquals(1, swapManager.swapOutCalledCount);

        assertNotNull(queue.poll(exp, 500000)); // this should trigger a swap-in of 10,000 records, and then pull 1 off the top.
        assertEquals(1, swapManager.swapInCalledCount);
        assertEquals(9999, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        assertTrue(swapManager.swappedOut.isEmpty());

        queue.poll(exp, 500000);
    }

    @Test
    public void testSwapInWhenThresholdIsLessThanSwapSize() {
        // create a queue where the swap threshold is less than 10k
        queue = new SwappablePriorityQueue(swapManager, 1000, eventReporter, flowFileQueue, dropAction, null);

        for (int i = 1; i <= 20000; i++) {
            queue.put(new MockFlowFileRecord());
        }

        assertEquals(1, swapManager.swappedOut.size());
        queue.put(new MockFlowFileRecord());
        assertEquals(1, swapManager.swappedOut.size());

        final Set<FlowFileRecord> exp = new HashSet<>();

        // At this point there should be:
        // 1k flow files in the active queue
        // 9,001 flow files in the swap queue
        // 10k flow files swapped to disk

        for (int i = 0; i < 999; i++) { //
            final FlowFileRecord flowFile = queue.poll(exp, 500000);
            assertNotNull(flowFile);
            assertEquals(1, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());
            assertEquals(1, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getByteCount());

            queue.acknowledge(Collections.singleton(flowFile));
            assertEquals(0, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getObjectCount());
            assertEquals(0, queue.getQueueDiagnostics().getUnacknowledgedQueueSize().getByteCount());
        }

        assertEquals(0, swapManager.swapInCalledCount);
        assertEquals(1, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());
        assertNotNull(queue.poll(exp, 500000));

        assertEquals(0, swapManager.swapInCalledCount);
        assertEquals(0, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        assertEquals(1, swapManager.swapOutCalledCount);

        assertNotNull(queue.poll(exp, 500000)); // this should trigger a swap-in of 10,000 records, and then pull 1 off the top.
        assertEquals(1, swapManager.swapInCalledCount);
        assertEquals(9999, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());

        assertTrue(swapManager.swappedOut.isEmpty());

        queue.poll(exp, 500000);
    }

    @Test
    public void testQueueCountsUpdatedWhenIncompleteSwapFile() {
        for (int i = 1; i <= 20000; i++) {
            queue.put(new MockFlowFileRecord());
        }

        assertEquals(20000, queue.size().getObjectCount());
        assertEquals(20000, queue.size().getByteCount());

        assertEquals(1, swapManager.swappedOut.size());

        // when we swap in, cause an IncompleteSwapFileException to be
        // thrown and contain only 9,999 of the 10,000 FlowFiles
        swapManager.enableIncompleteSwapFileException(9999);
        final Set<FlowFileRecord> expired = Collections.emptySet();
        FlowFileRecord flowFile;

        for (int i = 0; i < 10000; i++) {
            flowFile = queue.poll(expired, 500000);
            assertNotNull(flowFile);
            queue.acknowledge(Collections.singleton(flowFile));
        }

        // 10,000 FlowFiles on queue - all swapped out
        assertEquals(10000, queue.size().getObjectCount());
        assertEquals(10000, queue.size().getByteCount());
        assertEquals(1, swapManager.swappedOut.size());
        assertEquals(0, swapManager.swapInCalledCount);

        // Trigger swap in. This will remove 1 FlowFile from queue, leaving 9,999 but
        // on swap in, we will get only 9,999 FlowFiles put onto the queue, and the queue size will
        // be decremented by 10,000 (because the Swap File's header tells us that there are 10K
        // FlowFiles, even though only 9999 are in the swap file)
        flowFile = queue.poll(expired, 500000);
        assertNotNull(flowFile);
        queue.acknowledge(Collections.singleton(flowFile));

        // size should be 9,998 because we lost 1 on Swap In, and then we pulled one above.
        assertEquals(9998, queue.size().getObjectCount());
        assertEquals(9998, queue.size().getByteCount());
        assertEquals(0, swapManager.swappedOut.size());
        assertEquals(1, swapManager.swapInCalledCount);

        for (int i = 0; i < 9998; i++) {
            flowFile = queue.poll(expired, 500000);
            assertNotNull(flowFile, "Null FlowFile when i = " + i);
            queue.acknowledge(Collections.singleton(flowFile));

            final QueueSize queueSize = queue.size();
            assertEquals(9998 - i - 1, queueSize.getObjectCount());
            assertEquals(9998 - i - 1, queueSize.getByteCount());
        }

        final QueueSize queueSize = queue.size();
        assertEquals(0, queueSize.getObjectCount());
        assertEquals(0L, queueSize.getByteCount());

        flowFile = queue.poll(expired, 500000);
        assertNull(flowFile);
    }

    @Test
    @Timeout(120)
    public void testDropSwappedFlowFiles() {
        for (int i = 1; i <= 30000; i++) {
            queue.put(new MockFlowFileRecord());
        }

        assertEquals(2, swapManager.swappedOut.size());
        final DropFlowFileRequest request = new DropFlowFileRequest("Unit Test");

        queue.dropFlowFiles(request, "Unit Test");

        assertEquals(0, queue.size().getObjectCount());
        assertEquals(0, queue.size().getByteCount());
        assertEquals(0, swapManager.swappedOut.size());
        assertEquals(2, swapManager.swapInCalledCount);
    }


    @Test
    @Timeout(5)
    public void testGetActiveFlowFilesReturnsAllActiveFlowFiles() throws InterruptedException {
        for (int i = 0; i < 9999; i++) {
            queue.put(new MockFlowFileRecord());
        }

        final List<FlowFileRecord> active = queue.getActiveFlowFiles();
        assertNotNull(active);
        assertEquals(9999, active.size());
    }


    @Test
    @Timeout(5)
    public void testListFlowFilesResultsLimited() throws InterruptedException {
        for (int i = 0; i < 30050; i++) {
            queue.put(new MockFlowFileRecord());
        }

        final List<FlowFileRecord> activeFlowFiles = queue.getActiveFlowFiles();
        assertNotNull(activeFlowFiles);
        assertEquals(10000, activeFlowFiles.size());
    }


    @Test
    public void testOOMEFollowedBySuccessfulSwapIn() {
        final List<FlowFileRecord> flowFiles = new ArrayList<>();
        for (int i = 0; i < 50000; i++) {
            flowFiles.add(new MockFlowFileRecord());
        }

        queue.putAll(flowFiles);

        swapManager.failSwapInAfterN = 2;
        swapManager.setSwapInFailure(new OutOfMemoryError("Intentional OOME for unit test"));

        final Set<FlowFileRecord> expiredRecords = new HashSet<>();
        for (int i = 0; i < 30000; i++) {
            final FlowFileRecord polled = queue.poll(expiredRecords, 500000);
            assertNotNull(polled);
        }

        // verify that unexpected ERROR's are handled in such a way that we keep retrying
        for (int i = 0; i < 3; i++) {
            assertThrows(OutOfMemoryError.class, () -> queue.poll(expiredRecords, 500000));
        }

        // verify that unexpected Runtime Exceptions are handled in such a way that we keep retrying
        swapManager.setSwapInFailure(new NullPointerException("Intentional OOME for unit test"));

        for (int i = 0; i < 3; i++) {
            assertThrows(NullPointerException.class, () -> queue.poll(expiredRecords, 500000));
        }

        swapManager.failSwapInAfterN = -1;

        for (int i = 0; i < 20000; i++) {
            final FlowFileRecord polled = queue.poll(expiredRecords, 500000);
            assertNotNull(polled);
        }

        queue.acknowledge(flowFiles);
        assertNull(queue.poll(expiredRecords, 500000));
        assertEquals(0, queue.getQueueDiagnostics().getActiveQueueSize().getObjectCount());
        assertEquals(0, queue.size().getObjectCount());

        assertTrue(swapManager.swappedOut.isEmpty());
    }

    @Test
    // The purpose of this test is to populate a SwappablePriorityQueue and to ensure that the minLastQueueDate and TotalQueueDuration are correct.
    // To truly test this we need to get both the in-memory swap queue and swap "on disk" involved.
    public void testLastQueueDateMetrics() throws IOException {
        Set<FlowFileRecord> flowFileRecords = new HashSet<>(11001);
        queue = new SwappablePriorityQueue(swapManager, 1000, eventReporter, flowFileQueue, dropAction, "testGetMinLastQueueDate");
        long minQueueDate = Long.MAX_VALUE;
        long totalQueueDate = 0L;
        // Put enough files in the queue to swap to disk
        for (int i = 1; i <= 11001; i++) {
            FlowFileRecord flowFileRecord = new MockFlowFileRecord();
            queue.put(flowFileRecord);
            flowFileRecords.add(flowFileRecord);
            totalQueueDate += flowFileRecord.getLastQueueDate();
            minQueueDate = Long.min(minQueueDate, flowFileRecord.getLastQueueDate());
        }

        // Assert the queue has a max of active, in-memory swap, and on-disk swap
        assertEquals(1000, queue.getActiveFlowFiles().size());
        assertEquals(10001, queue.getFlowFileQueueSize().getSwappedCount());
        assertEquals(1, queue.getFlowFileQueueSize().getSwapFileCount());
        assertEquals(10000, swapManager.getSwapSummary(swapManager.recoverSwapLocations(flowFileQueue, "testGetMinLastQueueDate").get(0)).getQueueSize().getObjectCount());

        // Ensure that the min and totals are correct
        long now = System.currentTimeMillis();
        long totalNow = now * flowFileRecords.size();
        assertEquals(totalNow - totalQueueDate, queue.getTotalQueuedDuration(now));
        assertEquals(minQueueDate, queue.getMinLastQueueDate());

        List<FlowFileRecord> polledRecords = queue.poll(1000, Collections.emptySet(), -1);
        polledRecords.addAll(queue.poll(2, Collections.emptySet(), -1));

        // Assert that the lone swap file was recovered into memory and that all numbers are as we still expect them to be.
        assertEquals(9998, queue.getActiveFlowFiles().size());
        assertEquals(1, queue.getFlowFileQueueSize().getSwappedCount());
        assertEquals(0, queue.getFlowFileQueueSize().getSwapFileCount());
        assert(swapManager.recoverSwapLocations(flowFileQueue, "testGetMinLastQueueDate").isEmpty());

        // Ensure that the min and total are still correct
        flowFileRecords.removeAll(polledRecords);
        totalQueueDate = flowFileRecords.stream().mapToLong(FlowFileRecord::getLastQueueDate).sum();
        minQueueDate = flowFileRecords.stream().mapToLong(FlowFileRecord::getLastQueueDate).min().getAsLong();
        now = System.currentTimeMillis();
        totalNow = now * flowFileRecords.size();
        assertEquals(totalNow - totalQueueDate, queue.getTotalQueuedDuration(now));
        assertEquals(minQueueDate, queue.getMinLastQueueDate());
    }
}
