Examples with StandardFlowFileQueue org.apache.nifi.controller.StandardFlowFileQueue used on opensource projects

Example 1 with StandardFlowFileQueue

use of org.apache.nifi.controller.StandardFlowFileQueue in project nifi by apache.

the class TestWriteAheadFlowFileRepository method testResourceClaimsIncremented.

@Test
public void testResourceClaimsIncremented() throws IOException {
    final ResourceClaimManager claimManager = new StandardResourceClaimManager();
    final TestQueueProvider queueProvider = new TestQueueProvider();
    final Connection connection = Mockito.mock(Connection.class);
    when(connection.getIdentifier()).thenReturn("1234");
    when(connection.getDestination()).thenReturn(Mockito.mock(Connectable.class));
    final FlowFileSwapManager swapMgr = new MockFlowFileSwapManager();
    final FlowFileQueue queue = new StandardFlowFileQueue("1234", connection, null, null, claimManager, null, swapMgr, null, 10000);
    when(connection.getFlowFileQueue()).thenReturn(queue);
    queueProvider.addConnection(connection);
    final ResourceClaim resourceClaim1 = claimManager.newResourceClaim("container", "section", "1", false, false);
    final ContentClaim claim1 = new StandardContentClaim(resourceClaim1, 0L);
    final ResourceClaim resourceClaim2 = claimManager.newResourceClaim("container", "section", "2", false, false);
    final ContentClaim claim2 = new StandardContentClaim(resourceClaim2, 0L);
    // resource claims' counts should be updated for both the swapped out FlowFile and the non-swapped out FlowFile
    try (final WriteAheadFlowFileRepository repo = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null))) {
        repo.initialize(claimManager);
        repo.loadFlowFiles(queueProvider, -1L);
        // Create a Repository Record that indicates that a FlowFile was created
        final FlowFileRecord flowFile1 = new StandardFlowFileRecord.Builder().id(1L).addAttribute("uuid", "11111111-1111-1111-1111-111111111111").contentClaim(claim1).build();
        final StandardRepositoryRecord rec1 = new StandardRepositoryRecord(queue);
        rec1.setWorking(flowFile1);
        rec1.setDestination(queue);
        // Create a Record that we can swap out
        final FlowFileRecord flowFile2 = new StandardFlowFileRecord.Builder().id(2L).addAttribute("uuid", "11111111-1111-1111-1111-111111111112").contentClaim(claim2).build();
        final StandardRepositoryRecord rec2 = new StandardRepositoryRecord(queue);
        rec2.setWorking(flowFile2);
        rec2.setDestination(queue);
        final List<RepositoryRecord> records = new ArrayList<>();
        records.add(rec1);
        records.add(rec2);
        repo.updateRepository(records);
        final String swapLocation = swapMgr.swapOut(Collections.singletonList(flowFile2), queue);
        repo.swapFlowFilesOut(Collections.singletonList(flowFile2), queue, swapLocation);
    }
    final ResourceClaimManager recoveryClaimManager = new StandardResourceClaimManager();
    try (final WriteAheadFlowFileRepository repo = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null))) {
        repo.initialize(recoveryClaimManager);
        final long largestId = repo.loadFlowFiles(queueProvider, 0L);
        // largest ID known is 1 because this doesn't take into account the FlowFiles that have been swapped out
        assertEquals(1, largestId);
    }
    // resource claim 1 will have a single claimant count while resource claim 2 will have no claimant counts
    // because resource claim 2 is referenced only by flowfiles that are swapped out.
    assertEquals(1, recoveryClaimManager.getClaimantCount(resourceClaim1));
    assertEquals(0, recoveryClaimManager.getClaimantCount(resourceClaim2));
    final SwapSummary summary = queue.recoverSwappedFlowFiles();
    assertNotNull(summary);
    assertEquals(2, summary.getMaxFlowFileId().intValue());
    assertEquals(new QueueSize(1, 0L), summary.getQueueSize());
    final List<ResourceClaim> swappedOutClaims = summary.getResourceClaims();
    assertNotNull(swappedOutClaims);
    assertEquals(1, swappedOutClaims.size());
    assertEquals(claim2.getResourceClaim(), swappedOutClaims.get(0));
}

Example 2 with StandardFlowFileQueue

use of org.apache.nifi.controller.StandardFlowFileQueue in project nifi by apache.

the class TestStandardProcessSession method createFlowFileQueueSpy.

private FlowFileQueue createFlowFileQueueSpy(Connection connection) {
    final FlowFileSwapManager swapManager = Mockito.mock(FlowFileSwapManager.class);
    final ProcessScheduler processScheduler = Mockito.mock(ProcessScheduler.class);
    final StandardFlowFileQueue actualQueue = new StandardFlowFileQueue("1", connection, flowFileRepo, provenanceRepo, null, processScheduler, swapManager, null, 10000);
    return Mockito.spy(actualQueue);
}