org.apache.nifi.controller.repository.FlowFileRepository Java Examples

private static FlowFileRepository createFlowFileRepository(final NiFiProperties properties, final ResourceClaimManager contentClaimManager) {
    final String implementationClassName = properties.getProperty(NiFiProperties.FLOWFILE_REPOSITORY_IMPLEMENTATION, DEFAULT_FLOWFILE_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create FlowFile Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.FLOWFILE_REPOSITORY_IMPLEMENTATION);
    }

    try {
        final FlowFileRepository created = NarThreadContextClassLoader.createInstance(implementationClassName, FlowFileRepository.class, properties);
        synchronized (created) {
            created.initialize(contentClaimManager);
        }
        return created;
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

public StandardExtensionDiscoveringManager() {
    definitionMap.put(Processor.class, new HashSet<>());
    definitionMap.put(FlowFilePrioritizer.class, new HashSet<>());
    definitionMap.put(ReportingTask.class, new HashSet<>());
    definitionMap.put(ControllerService.class, new HashSet<>());
    definitionMap.put(Authorizer.class, new HashSet<>());
    definitionMap.put(UserGroupProvider.class, new HashSet<>());
    definitionMap.put(AccessPolicyProvider.class, new HashSet<>());
    definitionMap.put(LoginIdentityProvider.class, new HashSet<>());
    definitionMap.put(ProvenanceRepository.class, new HashSet<>());
    definitionMap.put(ComponentStatusRepository.class, new HashSet<>());
    definitionMap.put(FlowFileRepository.class, new HashSet<>());
    definitionMap.put(FlowFileSwapManager.class, new HashSet<>());
    definitionMap.put(ContentRepository.class, new HashSet<>());
    definitionMap.put(StateProvider.class, new HashSet<>());
    definitionMap.put(StatusAnalyticsModel.class, new HashSet<>());
}
 

private FileSystemSwapManager createSwapManager(final FlowFileRepository flowFileRepo) {
    final FileSystemSwapManager swapManager = new FileSystemSwapManager();
    final ResourceClaimManager resourceClaimManager = new NopResourceClaimManager();
    swapManager.initialize(new SwapManagerInitializationContext() {
        @Override
        public ResourceClaimManager getResourceClaimManager() {
            return resourceClaimManager;
        }

        @Override
        public FlowFileRepository getFlowFileRepository() {
            return flowFileRepo;
        }

        @Override
        public EventReporter getEventReporter() {
            return EventReporter.NO_OP;
        }
    });

    return swapManager;
}
 

@Test
public void testFailureOnRepoSwapOut() throws IOException {
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    when(flowFileQueue.getIdentifier()).thenReturn("87bb99fe-412c-49f6-a441-d1b0af4e20b4");

    final FlowFileRepository flowFileRepo = Mockito.mock(FlowFileRepository.class);
    Mockito.doThrow(new IOException("Intentional IOException for unit test"))
        .when(flowFileRepo).updateRepository(anyCollection());

    final FileSystemSwapManager swapManager = createSwapManager();

    final List<FlowFileRecord> flowFileRecords = new ArrayList<>();
    for (int i=0; i < 10000; i++) {
        flowFileRecords.add(new MockFlowFileRecord(i));
    }

    try {
        swapManager.swapOut(flowFileRecords, flowFileQueue, "partition-1");
        Assert.fail("Expected IOException");
    } catch (final IOException ioe) {
        // expected
    }
}
 

private static FlowFileRepository createFlowFileRepository(final NiFiProperties properties, final ExtensionManager extensionManager, final ResourceClaimManager contentClaimManager) {
    final String implementationClassName = properties.getProperty(NiFiProperties.FLOWFILE_REPOSITORY_IMPLEMENTATION, DEFAULT_FLOWFILE_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create FlowFile Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.FLOWFILE_REPOSITORY_IMPLEMENTATION);
    }

    try {
        final FlowFileRepository created = NarThreadContextClassLoader.createInstance(extensionManager, implementationClassName,
                FlowFileRepository.class, properties);
        if (EncryptedSequentialAccessWriteAheadLog.class.getName().equals(properties.getProperty(NiFiProperties.FLOWFILE_REPOSITORY_WAL_IMPLEMENTATION))
                && created instanceof WriteAheadFlowFileRepository) {
            synchronized (created) {
                ((WriteAheadFlowFileRepository) created).initialize(contentClaimManager, new EncryptedRepositoryRecordSerdeFactory(contentClaimManager, properties));
            }
        } else {
            synchronized (created) {
                created.initialize(contentClaimManager);
            }
        }
        return created;
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

private FileSystemSwapManager createSwapManager() {
    final FileSystemSwapManager swapManager = new FileSystemSwapManager();
    final ResourceClaimManager resourceClaimManager = new NopResourceClaimManager();
    final FlowFileRepository flowfileRepo = Mockito.mock(FlowFileRepository.class);
    swapManager.initialize(new SwapManagerInitializationContext() {
        @Override
        public ResourceClaimManager getResourceClaimManager() {
            return resourceClaimManager;
        }

        @Override
        public FlowFileRepository getFlowFileRepository() {
            return flowfileRepo;
        }

        @Override
        public EventReporter getEventReporter() {
            return EventReporter.NO_OP;
        }
    });

    return swapManager;
}
 

public StandardFlowFileQueue(final String identifier, final Connection connection, final FlowFileRepository flowFileRepo, final ProvenanceEventRepository provRepo,
    final ResourceClaimManager resourceClaimManager, final ProcessScheduler scheduler, final FlowFileSwapManager swapManager, final EventReporter eventReporter, final int swapThreshold) {
    activeQueue = new PriorityQueue<>(20, new Prioritizer(new ArrayList<FlowFilePrioritizer>()));
    priorities = new ArrayList<>();
    swapQueue = new ArrayList<>();
    this.eventReporter = eventReporter;
    this.swapManager = swapManager;
    this.flowFileRepository = flowFileRepo;
    this.provRepository = provRepo;
    this.resourceClaimManager = resourceClaimManager;

    this.identifier = identifier;
    this.swapThreshold = swapThreshold;
    this.scheduler = scheduler;
    this.connection = connection;

    readLock = new TimedLock(this.lock.readLock(), identifier + " Read Lock", 100);
    writeLock = new TimedLock(this.lock.writeLock(), identifier + " Write Lock", 100);
}
 

public RepositoryContextFactory(final ContentRepository contentRepository, final FlowFileRepository flowFileRepository,
        final FlowFileEventRepository flowFileEventRepository, final CounterRepository counterRepository,
        final ProvenanceRepository provenanceRepository) {

    this.contentRepo = contentRepository;
    this.flowFileRepo = flowFileRepository;
    this.flowFileEventRepo = flowFileEventRepository;
    this.counterRepo = counterRepository;
    this.provenanceRepo = provenanceRepository;
}
 

public ProcessContextFactory(final ContentRepository contentRepository, final FlowFileRepository flowFileRepository,
        final FlowFileEventRepository flowFileEventRepository, final CounterRepository counterRepository,
        final ProvenanceEventRepository provenanceRepository) {

    this.contentRepo = contentRepository;
    this.flowFileRepo = flowFileRepository;
    this.flowFileEventRepo = flowFileEventRepository;
    this.counterRepo = counterRepository;
    this.provenanceRepo = provenanceRepository;
}
 

private FlowFileRepository createFlowFileRepository(final List<RepositoryRecord> repoRecords) throws IOException {
    final FlowFileRepository flowFileRepo = mock(FlowFileRepository.class);
    doAnswer(invocation -> {
        final Collection records = invocation.getArgument(0);
        repoRecords.addAll(records);
        return null;
    }).when(flowFileRepo).updateRepository(anyCollection());

    return flowFileRepo;
}
 

@Before
@SuppressWarnings("unchecked")
public void setup() {
    provRecords.clear();

    connection = Mockito.mock(Connection.class);
    Mockito.when(connection.getSource()).thenReturn(Mockito.mock(Connectable.class));
    Mockito.when(connection.getDestination()).thenReturn(Mockito.mock(Connectable.class));

    scheduler = Mockito.mock(ProcessScheduler.class);
    swapManager = new MockSwapManager();

    flowFileRepo = Mockito.mock(FlowFileRepository.class);
    provRepo = Mockito.mock(ProvenanceEventRepository.class);
    claimManager = Mockito.mock(ResourceClaimManager.class);

    Mockito.when(provRepo.eventBuilder()).thenReturn(new StandardProvenanceEventRecord.Builder());
    Mockito.doAnswer(new Answer<Object>() {
        @Override
        public Object answer(final InvocationOnMock invocation) throws Throwable {
            final Iterable<ProvenanceEventRecord> iterable = (Iterable<ProvenanceEventRecord>) invocation.getArguments()[0];
            for (final ProvenanceEventRecord record : iterable) {
                provRecords.add(record);
            }
            return null;
        }
    }).when(provRepo).registerEvents(Mockito.any(Iterable.class));

    queue = new StandardFlowFileQueue("id", new NopConnectionEventListener(), flowFileRepo, provRepo, claimManager, scheduler, swapManager, null, 10000, 0L, "0 B");
    MockFlowFileRecord.resetIdGenerator();
}
 

public StandardFlowFileQueue(final String identifier, final ConnectionEventListener eventListener, final FlowFileRepository flowFileRepo, final ProvenanceEventRepository provRepo,
                             final ResourceClaimManager resourceClaimManager, final ProcessScheduler scheduler, final FlowFileSwapManager swapManager, final EventReporter eventReporter,
                             final int swapThreshold, final long defaultBackPressureObjectThreshold, final String defaultBackPressureDataSizeThreshold) {

    super(identifier, scheduler, flowFileRepo, provRepo, resourceClaimManager);
    this.swapManager = swapManager;
    this.queue = new SwappablePriorityQueue(swapManager, swapThreshold, eventReporter, this, this::drop, null);
    this.eventListener = eventListener;

    writeLock = new TimedLock(this.lock.writeLock(), getIdentifier() + " Write Lock", 100);

    setBackPressureDataSizeThreshold(defaultBackPressureDataSizeThreshold);
    setBackPressureObjectThreshold(defaultBackPressureObjectThreshold);
}
 

public StandardLoadBalanceProtocol(final FlowFileRepository flowFileRepository, final ContentRepository contentRepository, final ProvenanceRepository provenanceRepository,
                                   final FlowController flowController, final LoadBalanceAuthorizer authorizer) {
    this.flowFileRepository = flowFileRepository;
    this.contentRepository = contentRepository;
    this.provenanceRepository = provenanceRepository;
    this.flowController = flowController;
    this.authorizer = authorizer;
}
 

public RemoteQueuePartition(final NodeIdentifier nodeId, final SwappablePriorityQueue priorityQueue, final TransferFailureDestination failureDestination,
                            final FlowFileRepository flowFileRepo, final ProvenanceEventRepository provRepo, final ContentRepository contentRepository,
                            final AsyncLoadBalanceClientRegistry clientRegistry, final LoadBalancedFlowFileQueue flowFileQueue) {

    this.nodeIdentifier = nodeId;
    this.priorityQueue = priorityQueue;
    this.flowFileQueue = flowFileQueue;
    this.failureDestination = failureDestination;
    this.flowFileRepo = flowFileRepo;
    this.provRepo = provRepo;
    this.contentRepo = contentRepository;
    this.clientRegistry = clientRegistry;
    this.description = "RemoteQueuePartition[queueId=" + flowFileQueue.getIdentifier() + ", nodeId=" + nodeIdentifier + "]";
}
 

public AbstractFlowFileQueue(final String identifier, final ProcessScheduler scheduler,
        final FlowFileRepository flowFileRepo, final ProvenanceEventRepository provRepo, final ResourceClaimManager resourceClaimManager) {
    this.identifier = identifier;
    this.scheduler = scheduler;
    this.flowFileRepository = flowFileRepo;
    this.provRepository = provRepo;
    this.resourceClaimManager = resourceClaimManager;
}
 

public FlowFileSwapManager createSwapManager() {
    final String implementationClassName = nifiProperties.getProperty(NiFiProperties.FLOWFILE_SWAP_MANAGER_IMPLEMENTATION, DEFAULT_SWAP_MANAGER_IMPLEMENTATION);
    if (implementationClassName == null) {
        return null;
    }

    try {
        final FlowFileSwapManager swapManager = NarThreadContextClassLoader.createInstance(extensionManager, implementationClassName, FlowFileSwapManager.class, nifiProperties);

        final EventReporter eventReporter = createEventReporter();
        try (final NarCloseable narCloseable = NarCloseable.withNarLoader()) {
            final SwapManagerInitializationContext initializationContext = new SwapManagerInitializationContext() {
                @Override
                public ResourceClaimManager getResourceClaimManager() {
                    return resourceClaimManager;
                }

                @Override
                public FlowFileRepository getFlowFileRepository() {
                    return flowFileRepository;
                }

                @Override
                public EventReporter getEventReporter() {
                    return eventReporter;
                }
            };

            swapManager.initialize(initializationContext);
        }

        return swapManager;
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

@Before
@SuppressWarnings("unchecked")
public void setup() {
    provRecords.clear();

    connection = Mockito.mock(Connection.class);
    Mockito.when(connection.getSource()).thenReturn(Mockito.mock(Connectable.class));
    Mockito.when(connection.getDestination()).thenReturn(Mockito.mock(Connectable.class));

    scheduler = Mockito.mock(ProcessScheduler.class);
    swapManager = new TestSwapManager();

    flowFileRepo = Mockito.mock(FlowFileRepository.class);
    provRepo = Mockito.mock(ProvenanceEventRepository.class);
    claimManager = Mockito.mock(ResourceClaimManager.class);

    Mockito.when(provRepo.eventBuilder()).thenReturn(new StandardProvenanceEventRecord.Builder());
    Mockito.doAnswer(new Answer<Object>() {
        @Override
        public Object answer(final InvocationOnMock invocation) throws Throwable {
            final Iterable<ProvenanceEventRecord> iterable = (Iterable<ProvenanceEventRecord>) invocation.getArguments()[0];
            for (final ProvenanceEventRecord record : iterable) {
                provRecords.add(record);
            }
            return null;
        }
    }).when(provRepo).registerEvents(Mockito.any(Iterable.class));

    queue = new StandardFlowFileQueue("id", connection, flowFileRepo, provRepo, claimManager, scheduler, swapManager, null, 10000);
    TestFlowFile.idGenerator.set(0L);
}
 

@Override
public FlowInheritability checkInheritability(final DataFlow existingFlow, final DataFlow proposedFlow, final FlowController flowController) {
    final Document flowDocument = proposedFlow.getFlowDocument();
    final Element rootGroupElement = (Element) flowDocument.getDocumentElement().getElementsByTagName("rootGroup").item(0);
    final FlowEncodingVersion encodingVersion = FlowEncodingVersion.parse(flowDocument.getDocumentElement());

    final ProcessGroupDTO rootGroupDto = FlowFromDOMFactory.getProcessGroup(null, rootGroupElement, null, encodingVersion);
    final Set<String> connectionIds = findAllConnectionIds(rootGroupDto);

    final FlowFileRepository flowFileRepository = flowController.getRepositoryContextFactory().getFlowFileRepository();

    final Set<String> queuesWithFlowFiles;
    try {
        queuesWithFlowFiles = flowFileRepository.findQueuesWithFlowFiles(flowController.createSwapManager());
    } catch (final IOException ioe) {
        throw new FlowSynchronizationException("Failed to determine which connections have FlowFiles queued", ioe);
    }

    logger.debug("The following {} Connections/Queues have data queued up currently: {}", queuesWithFlowFiles.size(), queuesWithFlowFiles);

    for (final String queueId : queuesWithFlowFiles) {
        if (!connectionIds.contains(queueId)) {
            return FlowInheritability.notInheritable("Proposed Flow does not contain a Connection with ID " + queueId + " but this instance has data queued in that connection");
        }
    }

    return FlowInheritability.inheritable();
}
 

public SocketLoadBalancedFlowFileQueue(final String identifier, final ConnectionEventListener eventListener, final ProcessScheduler scheduler, final FlowFileRepository flowFileRepo,
                                       final ProvenanceEventRepository provRepo, final ContentRepository contentRepo, final ResourceClaimManager resourceClaimManager,
                                       final ClusterCoordinator clusterCoordinator, final AsyncLoadBalanceClientRegistry clientRegistry, final FlowFileSwapManager swapManager,
                                       final int swapThreshold, final EventReporter eventReporter) {

    super(identifier, scheduler, flowFileRepo, provRepo, resourceClaimManager);
    this.eventListener = eventListener;
    this.eventReporter = eventReporter;
    this.swapManager = swapManager;
    this.flowFileRepo = flowFileRepo;
    this.provRepo = provRepo;
    this.contentRepo = contentRepo;
    this.clusterCoordinator = clusterCoordinator;
    this.clientRegistry = clientRegistry;

    localPartition = new SwappablePriorityQueueLocalPartition(swapManager, swapThreshold, eventReporter, this, this::drop);
    rebalancingPartition = new StandardRebalancingPartition(swapManager, swapThreshold, eventReporter, this, this::drop);

    // Create a RemoteQueuePartition for each node
    nodeIdentifiers = clusterCoordinator == null ? Collections.emptySet() : clusterCoordinator.getNodeIdentifiers();

    final List<NodeIdentifier> sortedNodeIdentifiers = new ArrayList<>(nodeIdentifiers);
    sortedNodeIdentifiers.sort(Comparator.comparing(NodeIdentifier::getApiAddress));

    if (sortedNodeIdentifiers.isEmpty()) {
        // No Node Identifiers are known yet. Just create the partitions using the local partition.
        queuePartitions = new QueuePartition[] { localPartition };
    } else {
        // The node identifiers are known. Create the partitions using the local partition and 1 Remote Partition for each node
        // that is not the local node identifier. If the Local Node Identifier is not yet known, that's okay. When it becomes known,
        // the queuePartitions array will be recreated with the appropriate partitions.
        final List<QueuePartition> partitionList = new ArrayList<>();

        final NodeIdentifier localNodeId = clusterCoordinator.getLocalNodeIdentifier();
        for (final NodeIdentifier nodeId : sortedNodeIdentifiers) {
            if (nodeId.equals(localNodeId)) {
                partitionList.add(localPartition);
            } else {
                partitionList.add(createRemotePartition(nodeId));
            }
        }

        // Ensure that our list of queue partitions always contains the local partition.
        if (!partitionList.contains(localPartition)) {
            partitionList.add(localPartition);
        }

        queuePartitions = partitionList.toArray(new QueuePartition[0]);
    }

    partitioner = new LocalPartitionPartitioner();

    if (clusterCoordinator != null) {
        clusterCoordinator.registerEventListener(new ClusterEventListener());
    }

    rebalancingPartition.start(partitioner);
}
 

/**
 * Creates a connection between two Connectable objects.
 *
 * @param id                required ID of the connection
 * @param name              the name of the connection, or <code>null</code> to leave the
 *                          connection unnamed
 * @param source            required source
 * @param destination       required destination
 * @param relationshipNames required collection of relationship names
 * @return the connection
 * @throws NullPointerException     if the ID, source, destination, or set of relationships is null.
 * @throws IllegalArgumentException if <code>relationships</code> is an empty collection
 */
public Connection createConnection(final String id, final String name, final Connectable source, final Connectable destination, final Collection<String> relationshipNames) {
    final StandardConnection.Builder builder = new StandardConnection.Builder(processScheduler);

    final List<Relationship> relationships = new ArrayList<>();
    for (final String relationshipName : requireNonNull(relationshipNames)) {
        relationships.add(new Relationship.Builder().name(relationshipName).build());
    }

    // Create and initialize a FlowFileSwapManager for this connection
    final FlowFileSwapManager swapManager = createSwapManager();
    final EventReporter eventReporter = createEventReporter();

    try (final NarCloseable narCloseable = NarCloseable.withNarLoader()) {
        final SwapManagerInitializationContext initializationContext = new SwapManagerInitializationContext() {
            @Override
            public ResourceClaimManager getResourceClaimManager() {
                return resourceClaimManager;
            }

            @Override
            public FlowFileRepository getFlowFileRepository() {
                return flowFileRepository;
            }

            @Override
            public EventReporter getEventReporter() {
                return eventReporter;
            }
        };

        swapManager.initialize(initializationContext);
    }

    final FlowFileQueueFactory flowFileQueueFactory = new FlowFileQueueFactory() {
        @Override
        public FlowFileQueue createFlowFileQueue(final LoadBalanceStrategy loadBalanceStrategy, final String partitioningAttribute, final ConnectionEventListener eventListener) {
            final FlowFileQueue flowFileQueue;

            if (clusterCoordinator == null) {
                flowFileQueue = new StandardFlowFileQueue(id, eventListener, flowFileRepository, provenanceRepository, resourceClaimManager, processScheduler, swapManager,
                        eventReporter, nifiProperties.getQueueSwapThreshold(), nifiProperties.getDefaultBackPressureObjectThreshold(), nifiProperties.getDefaultBackPressureDataSizeThreshold());
            } else {
                flowFileQueue = new SocketLoadBalancedFlowFileQueue(id, eventListener, processScheduler, flowFileRepository, provenanceRepository, contentRepository, resourceClaimManager,
                        clusterCoordinator, loadBalanceClientRegistry, swapManager, nifiProperties.getQueueSwapThreshold(), eventReporter);

                flowFileQueue.setBackPressureObjectThreshold(nifiProperties.getDefaultBackPressureObjectThreshold());
                flowFileQueue.setBackPressureDataSizeThreshold(nifiProperties.getDefaultBackPressureDataSizeThreshold());
            }

            return flowFileQueue;
        }
    };

    final Connection connection = builder.id(requireNonNull(id).intern())
            .name(name == null ? null : name.intern())
            .relationships(relationships)
            .source(requireNonNull(source))
            .destination(destination)
            .flowFileQueueFactory(flowFileQueueFactory)
            .build();

    return connection;
}
 

@Override
public DiagnosticsDumpElement captureDump(final boolean verbose) {
    if (!verbose) {
        // This task is very expensive, as it must scan the contents of the Content Repository. As such, it will not
        // run at all if verbose output is disabled.
        return null;
    }

    final ContentRepository contentRepository = flowController.getRepositoryContextFactory().getContentRepository();
    if (!contentRepository.isActiveResourceClaimsSupported()) {
        return new StandardDiagnosticsDumpElement("Content Repository Scan", Collections.singletonList("Current Content Repository does not support scanning for in-use content"));
    }

    final FlowFileRepository flowFileRepository = flowController.getRepositoryContextFactory().getFlowFileRepository();
    final ResourceClaimManager resourceClaimManager = flowController.getResourceClaimManager();
    final FlowFileSwapManager swapManager = flowController.createSwapManager();

    final List<String> details = new ArrayList<>();

    for (final String containerName : contentRepository.getContainerNames()) {
        try {
            final Set<ResourceClaim> resourceClaims = contentRepository.getActiveResourceClaims(containerName);

            final Map<ResourceClaim, Set<ResourceClaimReference>> referenceMap = flowFileRepository.findResourceClaimReferences(resourceClaims, swapManager);

            for (final ResourceClaim resourceClaim : resourceClaims) {
                final int claimCount = resourceClaimManager.getClaimantCount(resourceClaim);
                final boolean inUse = resourceClaim.isInUse();
                final boolean destructable = resourceClaimManager.isDestructable(resourceClaim);

                final Set<ResourceClaimReference> references = referenceMap == null ? Collections.emptySet() : referenceMap.getOrDefault(resourceClaim, Collections.emptySet());

                final String path = resourceClaim.getContainer() + "/" + resourceClaim.getSection() + "/" + resourceClaim.getId();
                details.add(String.format("%1$s, Claimant Count = %2$d, In Use = %3$b, Awaiting Destruction = %4$b, References (%5$d) = %6$s",
                    path, claimCount, inUse, destructable,  references.size(), references.toString()));
            }
        } catch (final Exception e) {
            logger.error("Failed to obtain listing of Active Resource Claims for container {}", containerName, e);
            details.add("Failed to obtain listing of Active Resource Claims in container " + containerName);
        }
    }


    return new StandardDiagnosticsDumpElement("Content Repository Scan", details);
}
 

private void captureDiagnostics(final FlowFileRepository repository, final List<String> details) throws IOException {
    details.add("FlowFile Repository Implementation: " + repository.getClass().getName());
    details.add("FlowFile Repository File Store: " + repository.getFileStoreName());
    details.add("FlowFile Repository Storage Capacity: " + FormatUtils.formatDataSize(repository.getStorageCapacity()));
    details.add("FlowFile Repository Usable Space: " + FormatUtils.formatDataSize(repository.getUsableStorageSpace()));
}
 

protected final void initialize(final NiFiProperties nifiProperties) throws IOException {
    this.nifiProperties = nifiProperties;

    final FlowFileEventRepository flowFileEventRepository = new RingBufferEventRepository(5);

    final BulletinRepository bulletinRepo = new VolatileBulletinRepository();
    flowEngine = new FlowEngine(4, "unit test flow engine");
    extensionManager = new DirectInjectionExtensionManager();

    extensionManager.injectExtensionType(FlowFileRepository.class, WriteAheadFlowFileRepository.class);
    extensionManager.injectExtensionType(ContentRepository.class, FileSystemRepository.class);
    extensionManager.injectExtensionType(ProvenanceRepository.class, WriteAheadProvenanceRepository.class);
    extensionManager.injectExtensionType(StateProvider.class, WriteAheadLocalStateProvider.class);
    extensionManager.injectExtensionType(ComponentStatusRepository.class, VolatileComponentStatusRepository.class);
    extensionManager.injectExtensionType(FlowFileSwapManager.class, FileSystemSwapManager.class);

    extensionManager.injectExtensionType(Processor.class, BiConsumerProcessor.class);
    extensionManager.injectExtensionType(Processor.class, GenerateProcessor.class);
    extensionManager.injectExtensionType(Processor.class, TerminateOnce.class);
    extensionManager.injectExtensionType(Processor.class, TerminateAll.class);
    extensionManager.injectExtensionType(Processor.class, NopProcessor.class);
    extensionManager.injectExtensionType(Processor.class, UsernamePasswordProcessor.class);

    injectExtensionTypes(extensionManager);
    systemBundle = SystemBundle.create(nifiProperties);
    extensionManager.discoverExtensions(systemBundle, Collections.emptySet());

    final StringEncryptor encryptor = StringEncryptor.createEncryptor("PBEWITHMD5AND256BITAES-CBC-OPENSSL", "BC", "unit-test");
    final Authorizer authorizer = new AlwaysAuthorizedAuthorizer();
    final AuditService auditService = new NopAuditService();

    if (isClusteredTest()) {
        final File zookeeperDir = new File("target/state/zookeeper");
        final File version2Dir =  new File(zookeeperDir, "version-2");

        if (!version2Dir.exists()) {
            assertTrue(version2Dir.mkdirs());
        }

        final File[] children = version2Dir.listFiles();
        if (children != null) {
            for (final File file : children) {
                FileUtils.deleteFile(file, true);
            }
        }

        clusterCoordinator = Mockito.mock(ClusterCoordinator.class);
        final HeartbeatMonitor heartbeatMonitor = Mockito.mock(HeartbeatMonitor.class);
        final NodeProtocolSender protocolSender = Mockito.mock(NodeProtocolSender.class);
        final LeaderElectionManager leaderElectionManager = new CuratorLeaderElectionManager(2, nifiProperties);

        final NodeIdentifier localNodeId = new NodeIdentifier(UUID.randomUUID().toString(), "localhost", 8111, "localhost", 8081,
            "localhost", 8082, "localhost", 8083, 8084, false, Collections.emptySet());
        final NodeIdentifier node2Id = new NodeIdentifier(UUID.randomUUID().toString(), "localhost", 8222, "localhost", 8081,
            "localhost", 8082, "localhost", 8083, 8084, false, Collections.emptySet());

        final Set<NodeIdentifier> nodeIdentifiers = new HashSet<>();
        nodeIdentifiers.add(localNodeId);
        nodeIdentifiers.add(node2Id);
        Mockito.when(clusterCoordinator.getNodeIdentifiers()).thenReturn(nodeIdentifiers);
        Mockito.when(clusterCoordinator.getLocalNodeIdentifier()).thenReturn(localNodeId);

        flowController = FlowController.createClusteredInstance(flowFileEventRepository, nifiProperties, authorizer, auditService, encryptor, protocolSender, bulletinRepo, clusterCoordinator,
            heartbeatMonitor, leaderElectionManager, VariableRegistry.ENVIRONMENT_SYSTEM_REGISTRY, flowRegistryClient, extensionManager);

        flowController.setClustered(true, UUID.randomUUID().toString());
        flowController.setNodeId(localNodeId);

        flowController.setConnectionStatus(new NodeConnectionStatus(localNodeId, NodeConnectionState.CONNECTED));
    } else {
        flowController = FlowController.createStandaloneInstance(flowFileEventRepository, nifiProperties, authorizer, auditService, encryptor, bulletinRepo,
            VariableRegistry.ENVIRONMENT_SYSTEM_REGISTRY, flowRegistryClient, extensionManager);
    }

    processScheduler = new StandardProcessScheduler(flowEngine, flowController, encryptor, flowController.getStateManagerProvider(), nifiProperties);

    final RepositoryContextFactory repositoryContextFactory = flowController.getRepositoryContextFactory();
    final SchedulingAgent timerDrivenSchedulingAgent = new TimerDrivenSchedulingAgent(flowController, flowEngine, repositoryContextFactory, encryptor, nifiProperties);
    processScheduler.setSchedulingAgent(SchedulingStrategy.TIMER_DRIVEN, timerDrivenSchedulingAgent);

    flowFileSwapManager = flowController.createSwapManager();
    resourceClaimManager = flowController.getResourceClaimManager();
}
 

protected final FlowFileRepository getFlowFileRepository() {
    return getRepositoryContext().getFlowFileRepository();
}
 

@Override
protected void injectExtensionTypes(final DirectInjectionExtensionManager extensionManager) {
    extensionManager.injectExtensionType(FlowFileRepository.class, OOMEWriteAheadFlowFileRepository.class);
}
 

@Test
public void testSwapFileUnknownToRepoNotSwappedIn() throws IOException {
    final FlowFileQueue flowFileQueue = Mockito.mock(FlowFileQueue.class);
    when(flowFileQueue.getIdentifier()).thenReturn("");

    final File targetDir = new File("target/swap");
    targetDir.mkdirs();

    final File targetFile = new File(targetDir, "444-old-swap-file.swap");
    final File originalSwapFile = new File("src/test/resources/swap/444-old-swap-file.swap");
    try (final OutputStream fos = new FileOutputStream(targetFile);
         final InputStream fis = new FileInputStream(originalSwapFile)) {
        StreamUtils.copy(fis, fos);
    }

    final FileSystemSwapManager swapManager = new FileSystemSwapManager(Paths.get("target"));
    final ResourceClaimManager resourceClaimManager = new NopResourceClaimManager();
    final FlowFileRepository flowFileRepo = Mockito.mock(FlowFileRepository.class);

    swapManager.initialize(new SwapManagerInitializationContext() {
        @Override
        public ResourceClaimManager getResourceClaimManager() {
            return resourceClaimManager;
        }

        @Override
        public FlowFileRepository getFlowFileRepository() {
            return flowFileRepo;
        }

        @Override
        public EventReporter getEventReporter() {
            return EventReporter.NO_OP;
        }
    });

    when(flowFileRepo.isValidSwapLocationSuffix(anyString())).thenReturn(false);
    final List<String> recoveredLocations = swapManager.recoverSwapLocations(flowFileQueue, null);
    assertEquals(1, recoveredLocations.size());

    final String firstLocation = recoveredLocations.get(0);
    final SwapContents emptyContents = swapManager.swapIn(firstLocation, flowFileQueue);
    assertEquals(0, emptyContents.getFlowFiles().size());

    when(flowFileRepo.isValidSwapLocationSuffix(anyString())).thenReturn(true);
    when(flowFileQueue.getIdentifier()).thenReturn("87bb99fe-412c-49f6-a441-d1b0af4e20b4");
    final SwapContents contents = swapManager.swapIn(firstLocation, flowFileQueue);
    assertEquals(10000, contents.getFlowFiles().size());
}
 

private FileSystemSwapManager createSwapManager() {
    final FlowFileRepository flowFileRepo = Mockito.mock(FlowFileRepository.class);
    return createSwapManager(flowFileRepo);
}
 

/**
 * Creates a connection between two Connectable objects.
 *
 * @param id required ID of the connection
 * @param name the name of the connection, or <code>null</code> to leave the
 * connection unnamed
 * @param source required source
 * @param destination required destination
 * @param relationshipNames required collection of relationship names
 * @return
 *
 * @throws NullPointerException if the ID, source, destination, or set of
 * relationships is null.
 * @throws IllegalArgumentException if <code>relationships</code> is an
 * empty collection
 */
public Connection createConnection(final String id, final String name, final Connectable source, final Connectable destination, final Collection<String> relationshipNames) {
    final StandardConnection.Builder builder = new StandardConnection.Builder(processScheduler);

    final List<Relationship> relationships = new ArrayList<>();
    for (final String relationshipName : requireNonNull(relationshipNames)) {
        relationships.add(new Relationship.Builder().name(relationshipName).build());
    }

    // Create and initialize a FlowFileSwapManager for this connection
    final FlowFileSwapManager swapManager = createSwapManager(nifiProperties);
    final EventReporter eventReporter = createEventReporter(getBulletinRepository());

    try (final NarCloseable narCloseable = NarCloseable.withNarLoader()) {
        final SwapManagerInitializationContext initializationContext = new SwapManagerInitializationContext() {
            @Override
            public ResourceClaimManager getResourceClaimManager() {
                return resourceClaimManager;
            }

            @Override
            public FlowFileRepository getFlowFileRepository() {
                return flowFileRepository;
            }

            @Override
            public EventReporter getEventReporter() {
                return eventReporter;
            }
        };

        swapManager.initialize(initializationContext);
    }

    return builder.id(requireNonNull(id).intern())
            .name(name == null ? null : name.intern())
            .relationships(relationships)
            .source(requireNonNull(source))
            .destination(destination)
            .swapManager(swapManager)
            .queueSwapThreshold(nifiProperties.getQueueSwapThreshold())
            .eventReporter(eventReporter)
            .resourceClaimManager(resourceClaimManager)
            .flowFileRepository(flowFileRepository)
            .provenanceRepository(provenanceRepository)
            .build();
}
 

public FlowFileRepository getFlowFileRepository() {
    return flowFileRepo;
}