Examples with FlowFileQueue org.apache.nifi.controller.queue.FlowFileQueue used on opensource projects

Example 36 with FlowFileQueue

use of org.apache.nifi.controller.queue.FlowFileQueue in project nifi by apache.

the class TestWriteAheadFlowFileRepository method testRestartWithOneRecord.

@Test
public void testRestartWithOneRecord() throws IOException {
    final Path path = Paths.get("target/test-repo");
    if (Files.exists(path)) {
        FileUtils.deleteFile(path.toFile(), true);
    }
    final WriteAheadFlowFileRepository repo = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null));
    repo.initialize(new StandardResourceClaimManager());
    final TestQueueProvider queueProvider = new TestQueueProvider();
    repo.loadFlowFiles(queueProvider, 0L);
    final List<FlowFileRecord> flowFileCollection = new ArrayList<>();
    final Connection connection = Mockito.mock(Connection.class);
    when(connection.getIdentifier()).thenReturn("1234");
    final FlowFileQueue queue = Mockito.mock(FlowFileQueue.class);
    when(queue.getIdentifier()).thenReturn("1234");
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(final InvocationOnMock invocation) throws Throwable {
            flowFileCollection.add((FlowFileRecord) invocation.getArguments()[0]);
            return null;
        }
    }).when(queue).put(any(FlowFileRecord.class));
    when(connection.getFlowFileQueue()).thenReturn(queue);
    queueProvider.addConnection(connection);
    StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder();
    ffBuilder.id(1L);
    ffBuilder.addAttribute("abc", "xyz");
    ffBuilder.size(0L);
    final FlowFileRecord flowFileRecord = ffBuilder.build();
    final List<RepositoryRecord> records = new ArrayList<>();
    final StandardRepositoryRecord record = new StandardRepositoryRecord(null);
    record.setWorking(flowFileRecord);
    record.setDestination(connection.getFlowFileQueue());
    records.add(record);
    repo.updateRepository(records);
    // update to add new attribute
    ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord).addAttribute("hello", "world");
    final FlowFileRecord flowFileRecord2 = ffBuilder.build();
    record.setWorking(flowFileRecord2);
    repo.updateRepository(records);
    // update size but no attribute
    ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord2).size(40L);
    final FlowFileRecord flowFileRecord3 = ffBuilder.build();
    record.setWorking(flowFileRecord3);
    repo.updateRepository(records);
    repo.close();
    // restore
    final WriteAheadFlowFileRepository repo2 = new WriteAheadFlowFileRepository(NiFiProperties.createBasicNiFiProperties(null, null));
    repo2.initialize(new StandardResourceClaimManager());
    repo2.loadFlowFiles(queueProvider, 0L);
    assertEquals(1, flowFileCollection.size());
    final FlowFileRecord flowFile = flowFileCollection.get(0);
    assertEquals(1L, flowFile.getId());
    assertEquals("xyz", flowFile.getAttribute("abc"));
    assertEquals(40L, flowFile.getSize());
    assertEquals("world", flowFile.getAttribute("hello"));
    repo2.close();
}

Example 37 with FlowFileQueue

use of org.apache.nifi.controller.queue.FlowFileQueue in project nifi by apache.

the class TestSchemaSwapSerializerDeserializer method testRoundTripSerializeDeserializeSummary.

@Test
public void testRoundTripSerializeDeserializeSummary() throws IOException {
    final ResourceClaimManager resourceClaimManager = new StandardResourceClaimManager();
    final ResourceClaim firstResourceClaim = resourceClaimManager.newResourceClaim("container", "section", "id", true, false);
    resourceClaimManager.incrementClaimantCount(firstResourceClaim);
    final List<FlowFileRecord> toSwap = new ArrayList<>(10000);
    final Map<String, String> attrs = new HashMap<>();
    long size = 0L;
    final ContentClaim firstClaim = MockFlowFile.createContentClaim("id", resourceClaimManager);
    for (int i = 0; i < 10000; i++) {
        attrs.put("i", String.valueOf(i));
        final FlowFileRecord ff = i < 2 ? new MockFlowFile(attrs, i, firstClaim) : new MockFlowFile(attrs, i, resourceClaimManager);
        toSwap.add(ff);
        size += i;
    }
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    Mockito.when(flowFileQueue.getIdentifier()).thenReturn("87bb99fe-412c-49f6-a441-d1b0af4e20b4");
    final String swapLocation = "target/testRoundTrip.swap";
    final File swapFile = new File(swapLocation);
    Files.deleteIfExists(swapFile.toPath());
    final SwapSerializer serializer = new SchemaSwapSerializer();
    try (final FileOutputStream fos = new FileOutputStream(swapFile)) {
        serializer.serializeFlowFiles(toSwap, flowFileQueue, swapLocation, fos);
    }
    final SwapDeserializer deserializer = new SchemaSwapDeserializer();
    final SwapSummary swapSummary;
    try (final FileInputStream fis = new FileInputStream(swapFile);
        final DataInputStream dis = new DataInputStream(fis)) {
        swapSummary = deserializer.getSwapSummary(dis, swapLocation, resourceClaimManager);
    }
    assertEquals(10000, swapSummary.getQueueSize().getObjectCount());
    assertEquals(size, swapSummary.getQueueSize().getByteCount());
    assertEquals(9999, swapSummary.getMaxFlowFileId().intValue());
    final List<ResourceClaim> resourceClaims = swapSummary.getResourceClaims();
    assertEquals(10000, resourceClaims.size());
    assertFalse(resourceClaims.stream().anyMatch(claim -> claim == null));
    assertEquals(2, resourceClaims.stream().filter(claim -> claim.getId().equals("id")).collect(Collectors.counting()).intValue());
    final Set<ResourceClaim> uniqueClaims = new HashSet<>(resourceClaims);
    assertEquals(9999, uniqueClaims.size());
}

Example 38 with FlowFileQueue

use of org.apache.nifi.controller.queue.FlowFileQueue in project nifi by apache.

the class TestSimpleSwapSerializerDeserializer method testWritePerformance.

@Test
@Ignore("For manual testing only. Not intended to be run as part of the automated unit tests but can " + "be convenient for determining a baseline for performance if making modifications.")
public void testWritePerformance() throws IOException, InterruptedException {
    final ResourceClaimManager resourceClaimManager = new StandardResourceClaimManager();
    final List<FlowFileRecord> toSwap = new ArrayList<>(10000);
    final Map<String, String> attrs = new HashMap<>();
    for (int i = 0; i < 10000; i++) {
        attrs.put("i", String.valueOf(i));
        final FlowFileRecord ff = new MockFlowFile(attrs, i, resourceClaimManager);
        toSwap.add(ff);
    }
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    Mockito.when(flowFileQueue.getIdentifier()).thenReturn("87bb99fe-412c-49f6-a441-d1b0af4e20b4");
    final String swapLocation = "target/testRoundTrip.swap";
    final int iterations = 1000;
    final long start = System.nanoTime();
    final SwapSerializer serializer = new SimpleSwapSerializer();
    for (int i = 0; i < iterations; i++) {
        try (final OutputStream out = new NullOutputStream()) {
            serializer.serializeFlowFiles(toSwap, flowFileQueue, swapLocation, out);
        }
    }
    final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
    System.out.println("Wrote " + iterations + " Swap Files in " + millis + " millis");
}

Example 39 with FlowFileQueue

use of org.apache.nifi.controller.queue.FlowFileQueue in project nifi by apache.

the class TestConnectableTask method testIsWorkToDo.

@Test
public void testIsWorkToDo() {
    final ProcessorNode procNode = Mockito.mock(ProcessorNode.class);
    Mockito.when(procNode.hasIncomingConnection()).thenReturn(false);
    final Processor processor = Mockito.mock(Processor.class);
    Mockito.when(procNode.getIdentifier()).thenReturn("123");
    Mockito.when(procNode.getRunnableComponent()).thenReturn(processor);
    final FlowController flowController = Mockito.mock(FlowController.class);
    Mockito.when(flowController.getStateManagerProvider()).thenReturn(Mockito.mock(StateManagerProvider.class));
    final RepositoryContext repoContext = Mockito.mock(RepositoryContext.class);
    Mockito.when(repoContext.getFlowFileEventRepository()).thenReturn(Mockito.mock(FlowFileEventRepository.class));
    final RepositoryContextFactory contextFactory = Mockito.mock(RepositoryContextFactory.class);
    Mockito.when(contextFactory.newProcessContext(Mockito.any(Connectable.class), Mockito.any(AtomicLong.class))).thenReturn(repoContext);
    final LifecycleState scheduleState = new LifecycleState();
    final StringEncryptor encryptor = Mockito.mock(StringEncryptor.class);
    ConnectableTask task = new ConnectableTask(Mockito.mock(SchedulingAgent.class), procNode, flowController, contextFactory, scheduleState, encryptor);
    // There is work to do because there are no incoming connections.
    assertFalse(task.invoke().isYield());
    // Test with only a single connection that is self-looping and empty
    final Connection selfLoopingConnection = Mockito.mock(Connection.class);
    when(selfLoopingConnection.getSource()).thenReturn(procNode);
    when(selfLoopingConnection.getDestination()).thenReturn(procNode);
    when(procNode.hasIncomingConnection()).thenReturn(true);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(selfLoopingConnection));
    assertFalse(task.invoke().isYield());
    // Test with only a single connection that is self-looping and empty
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    when(flowFileQueue.isActiveQueueEmpty()).thenReturn(true);
    final FlowFileQueue nonEmptyQueue = Mockito.mock(FlowFileQueue.class);
    when(nonEmptyQueue.isActiveQueueEmpty()).thenReturn(false);
    when(selfLoopingConnection.getFlowFileQueue()).thenReturn(nonEmptyQueue);
    assertFalse(task.invoke().isYield());
    // Test with only a non-looping Connection that has no FlowFiles
    final Connection emptyConnection = Mockito.mock(Connection.class);
    when(emptyConnection.getSource()).thenReturn(Mockito.mock(ProcessorNode.class));
    when(emptyConnection.getDestination()).thenReturn(procNode);
    when(emptyConnection.getFlowFileQueue()).thenReturn(flowFileQueue);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(emptyConnection));
    // Create a new ConnectableTask because we want to have a different value for the 'hasNonLoopConnection' value, which is calculated in the task's constructor.
    task = new ConnectableTask(Mockito.mock(SchedulingAgent.class), procNode, flowController, contextFactory, scheduleState, encryptor);
    assertTrue(task.invoke().isYield());
    // test when the queue has data
    final Connection nonEmptyConnection = Mockito.mock(Connection.class);
    when(nonEmptyConnection.getSource()).thenReturn(Mockito.mock(ProcessorNode.class));
    when(nonEmptyConnection.getDestination()).thenReturn(procNode);
    when(nonEmptyConnection.getFlowFileQueue()).thenReturn(nonEmptyQueue);
    when(procNode.getIncomingConnections()).thenReturn(Collections.singletonList(nonEmptyConnection));
    assertFalse(task.invoke().isYield());
}

Example 40 with FlowFileQueue

use of org.apache.nifi.controller.queue.FlowFileQueue in project nifi by apache.

the class FlowController method getAllQueues.

@Override
public Collection<FlowFileQueue> getAllQueues() {
    final Collection<Connection> connections = getGroup(getRootGroupId()).findAllConnections();
    final List<FlowFileQueue> queues = new ArrayList<>(connections.size());
    for (final Connection connection : connections) {
        queues.add(connection.getFlowFileQueue());
    }
    return queues;
}