org.apache.nifi.provenance.lucene.IndexManager Java Examples

@Test(timeout = 5000)
public void testExpiration() throws InterruptedException, IOException {
    final RepositoryConfiguration repoConfig = createConfig(1);
    repoConfig.setDesiredIndexSize(1L);
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final List<ProvenanceEventRecord> events = new ArrayList<>();
    events.add(createEvent(500000L));
    events.add(createEvent());

    final EventStore eventStore = Mockito.mock(EventStore.class);
    Mockito.doAnswer(new Answer<List<ProvenanceEventRecord>>() {
        @Override
        public List<ProvenanceEventRecord> answer(final InvocationOnMock invocation) throws Throwable {
            final Long eventId = invocation.getArgumentAt(0, Long.class);
            assertEquals(0, eventId.longValue());
            assertEquals(1, invocation.getArgumentAt(1, Integer.class).intValue());
            return Collections.singletonList(events.get(0));
        }
    }).when(eventStore).getEvents(Mockito.anyLong(), Mockito.anyInt());

    index.initialize(eventStore);
    index.addEvent(events.get(0), createStorageSummary(events.get(0).getEventId()));

    // Add the first event to the index and wait for it to be indexed, since indexing is asynchronous.
    List<File> allDirectories = Collections.emptyList();
    while (allDirectories.isEmpty()) {
        allDirectories = index.getDirectoryManager().getDirectories(null, null);
    }

    events.remove(0); // Remove the first event from the store
    index.performMaintenance();
    assertEquals(1, index.getDirectoryManager().getDirectories(null, null).size());
}
 

public LuceneEventIndex(final RepositoryConfiguration config, final IndexManager indexManager, final int maxEventsPerCommit, final EventReporter eventReporter) {
    this.eventReporter = eventReporter;
    queryExecutor = Executors.newFixedThreadPool(config.getQueryThreadPoolSize(), new NamedThreadFactory("Provenance Query"));
    indexExecutor = Executors.newFixedThreadPool(config.getIndexThreadPoolSize(), new NamedThreadFactory("Index Provenance Events"));
    cacheWarmerExecutor = Executors.newScheduledThreadPool(config.getStorageDirectories().size(), new NamedThreadFactory("Warm Lucene Index", true));
    directoryManager = new IndexDirectoryManager(config);

    // Limit number of indexing threads to 100. When we restore the repository on restart,
    // we have to re-index up to MAX_THREADS * MAX_DOCUMENTS_PER_THREADS events prior to
    // the last event that the index holds. This is done because we could have that many
    // events 'in flight', waiting to be indexed when the last index writer was committed,
    // so even though the index says the largest event ID is 1,000,000 for instance, Event
    // with ID 999,999 may still not have been indexed because another thread was in the
    // process of writing the event to the index.
    final int configuredIndexPoolSize = config.getIndexThreadPoolSize();
    final int numIndexThreads;
    if (configuredIndexPoolSize > MAX_INDEX_THREADS) {
        logger.warn("The Provenance Repository is configured to perform indexing of events using {} threads. This number exceeds the maximum allowable number of threads, which is {}. "
            + "Will proceed using {} threads. This value is limited because the performance of indexing will decrease and startup times will increase when setting this value too high.",
            configuredIndexPoolSize, MAX_INDEX_THREADS, MAX_INDEX_THREADS);
        numIndexThreads = MAX_INDEX_THREADS;
    } else {
        numIndexThreads = configuredIndexPoolSize;
    }

    for (int i = 0; i < numIndexThreads; i++) {
        final EventIndexTask task = new EventIndexTask(documentQueue, indexManager, directoryManager, maxEventsPerCommit, eventReporter);
        indexTasks.add(task);
        indexExecutor.submit(task);
    }

    this.config = config;
    this.indexManager = indexManager;
    this.eventConverter = new ConvertEventToLuceneDocument(config.getSearchableFields(), config.getSearchableAttributes());
}
 

public MigrateDefunctIndex(final File indexDirectory, final IndexManager indexManager, final IndexDirectoryManager directoryManager, final long minTimestamp, final long maxTimestamp,
                           final EventStore eventStore, final EventReporter eventReporter, final ConvertEventToLuceneDocument eventConverter, final AtomicInteger rebuildCount,
                           final int totalCount) {
    this.indexDirectory = indexDirectory;
    this.indexManager = indexManager;
    this.directoryManager = directoryManager;
    this.eventStore = eventStore;
    this.eventReporter = eventReporter;
    this.minTimestamp = minTimestamp;
    this.maxTimestamp = maxTimestamp;
    this.eventConverter = eventConverter;
    this.rebuildCount = rebuildCount;
    this.totalCount = totalCount;
}
 

public QueryTask(final Query query, final ProgressiveResult result, final int maxResults, final IndexManager indexManager,
    final File indexDir, final EventStore eventStore, final EventAuthorizer authorizer,
    final EventTransformer unauthorizedTransformer) {
    this.query = query;
    this.queryResult = result;
    this.maxResults = maxResults;
    this.indexManager = indexManager;
    this.indexDir = indexDir;
    this.eventStore = eventStore;
    this.authorizer = authorizer;
    this.transformer = unauthorizedTransformer;
}
 

public EventIndexTask(final BlockingQueue<StoredDocument> documentQueue, final IndexManager indexManager,
    final IndexDirectoryManager directoryManager, final int maxEventsPerCommit, final EventReporter eventReporter) {
    this.documentQueue = documentQueue;
    this.indexManager = indexManager;
    this.directoryManager = directoryManager;
    this.commitThreshold = maxEventsPerCommit;
    this.eventReporter = eventReporter;
}
 

synchronized void init(RecordWriterFactory recordWriterFactory, RecordReaderFactory recordReaderFactory,
                       final EventReporter eventReporter, final Authorizer authorizer,
                       final ProvenanceAuthorizableFactory resourceFactory) throws IOException {
    final EventFileManager fileManager = new EventFileManager();

    eventStore = new PartitionedWriteAheadEventStore(config, recordWriterFactory, recordReaderFactory, eventReporter, fileManager);

    final IndexManager indexManager = new StandardIndexManager(config);
    eventIndex = new LuceneEventIndex(config, indexManager, eventReporter);

    this.eventReporter = eventReporter;
    this.authorizer = authorizer;
    this.resourceFactory = resourceFactory;

    eventStore.initialize();
    eventIndex.initialize(eventStore);

    if (eventIndex.isReindexNecessary()) {
        try {
            eventStore.reindexLatestEvents(eventIndex);
        } catch (final Exception e) {
            logger.error("Failed to re-index some of the Provenance Events. It is possible that some of the latest "
                + "events will not be available from the Provenance Repository when a query is issued.", e);
        }
    } else {
        logger.info("Provenance Event Index indicates that no events should be re-indexed upon startup. Will not wait for re-indexing to occur.");
    }
}
 

@Test(timeout = 60000)
public void testUnauthorizedEventsGetPlaceholdersForLineage() throws InterruptedException {
    assumeFalse(isWindowsEnvironment());
    final RepositoryConfiguration repoConfig = createConfig(1);
    repoConfig.setDesiredIndexSize(1L);
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 3, EventReporter.NO_OP);
    index.initialize(eventStore);

    for (int i = 0; i < 3; i++) {
        final ProvenanceEventRecord event = createEvent("1234");
        final StorageResult storageResult = eventStore.addEvent(event);
        index.addEvents(storageResult.getStorageLocations());
    }

    final NiFiUser user = createUser();

    List<LineageNode> nodes = Collections.emptyList();
    while (nodes.size() < 3) {
        final ComputeLineageSubmission submission = index.submitLineageComputation(1L, user, EventAuthorizer.DENY_ALL);
        assertTrue(submission.getResult().awaitCompletion(15, TimeUnit.SECONDS));

        nodes = submission.getResult().getNodes();
        Thread.sleep(25L);
    }

    assertEquals(3, nodes.size());

    for (final LineageNode node : nodes) {
        assertEquals(LineageNodeType.PROVENANCE_EVENT_NODE, node.getNodeType());
        final ProvenanceEventLineageNode eventNode = (ProvenanceEventLineageNode) node;
        assertEquals(ProvenanceEventType.UNKNOWN, eventNode.getEventType());
    }
}
 

@Test(timeout = 60000)
public void testExpiration() throws IOException {
    final RepositoryConfiguration repoConfig = createConfig(1);
    repoConfig.setDesiredIndexSize(1L);
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final List<ProvenanceEventRecord> events = new ArrayList<>();
    events.add(createEvent(500000L));
    events.add(createEvent());

    final EventStore eventStore = Mockito.mock(EventStore.class);
    Mockito.doAnswer(new Answer<List<ProvenanceEventRecord>>() {
        @Override
        public List<ProvenanceEventRecord> answer(final InvocationOnMock invocation) {
            final Long eventId = invocation.getArgument(0);
            assertEquals(0, eventId.longValue());
            assertEquals(1, invocation.<Integer>getArgument(1).intValue());
            return Collections.singletonList(events.get(0));
        }
    }).when(eventStore).getEvents(Mockito.anyLong(), Mockito.anyInt());

    index.initialize(eventStore);
    index.addEvent(events.get(0), createStorageSummary(events.get(0).getEventId()));

    // Add the first event to the index and wait for it to be indexed, since indexing is asynchronous.
    List<File> allDirectories = Collections.emptyList();
    while (allDirectories.isEmpty()) {
        allDirectories = index.getDirectoryManager().getDirectories(null, null);
    }

    events.remove(0); // Remove the first event from the store
    index.performMaintenance();
    assertEquals(1, index.getDirectoryManager().getDirectories(null, null).size());
}
 

@Test(timeout = 5000)
public void addThenQueryWithEmptyQuery() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final ProvenanceEventRecord event = createEvent();

    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));

    final Query query = new Query(UUID.randomUUID().toString());

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(100, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

@Test(timeout = 5000)
public void testUnauthorizedEventsGetPlaceholdersForLineage() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig(1);
    repoConfig.setDesiredIndexSize(1L);
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 3, EventReporter.NO_OP);
    index.initialize(eventStore);

    for (int i = 0; i < 3; i++) {
        final ProvenanceEventRecord event = createEvent("1234");
        final StorageResult storageResult = eventStore.addEvent(event);
        index.addEvents(storageResult.getStorageLocations());
    }

    final NiFiUser user = createUser();

    List<LineageNode> nodes = Collections.emptyList();
    while (nodes.size() < 3) {
        final ComputeLineageSubmission submission = index.submitLineageComputation(1L, user, EventAuthorizer.DENY_ALL);
        assertTrue(submission.getResult().awaitCompletion(5, TimeUnit.SECONDS));

        nodes = submission.getResult().getNodes();
        Thread.sleep(25L);
    }

    assertEquals(3, nodes.size());

    for (final LineageNode node : nodes) {
        assertEquals(LineageNodeType.PROVENANCE_EVENT_NODE, node.getNodeType());
        final ProvenanceEventLineageNode eventNode = (ProvenanceEventLineageNode) node;
        assertEquals(ProvenanceEventType.UNKNOWN, eventNode.getEventType());
    }
}
 

@Override
public synchronized void initialize(final EventReporter eventReporter, final Authorizer authorizer, final ProvenanceAuthorizableFactory resourceFactory,
    final IdentifierLookup idLookup) throws IOException {
    final RecordWriterFactory recordWriterFactory = (file, idGenerator, compressed, createToc) -> {
        final TocWriter tocWriter = createToc ? new StandardTocWriter(TocUtil.getTocFile(file), false, false) : null;
        return new EventIdFirstSchemaRecordWriter(file, idGenerator, tocWriter, compressed, BLOCK_SIZE, idLookup);
    };

    final EventFileManager fileManager = new EventFileManager();
    final RecordReaderFactory recordReaderFactory = (file, logs, maxChars) -> {
        fileManager.obtainReadLock(file);
        try {
            return RecordReaders.newRecordReader(file, logs, maxChars);
        } finally {
            fileManager.releaseReadLock(file);
        }
    };

    eventStore = new PartitionedWriteAheadEventStore(config, recordWriterFactory, recordReaderFactory, eventReporter, fileManager);

    final IndexManager indexManager = new SimpleIndexManager(config);
    eventIndex = new LuceneEventIndex(config, indexManager, eventReporter);

    this.eventReporter = eventReporter;
    this.authorizer = authorizer;
    this.resourceFactory = resourceFactory;

    eventStore.initialize();
    eventIndex.initialize(eventStore);

    try {
        eventStore.reindexLatestEvents(eventIndex);
    } catch (final Exception e) {
        logger.error("Failed to re-index some of the Provenance Events. It is possible that some of the latest "
            + "events will not be available from the Provenance Repository when a query is issued.", e);
    }
}
 

@Test(timeout = 60000)
public void addThenQueryWithEmptyQuery() throws InterruptedException {
    assumeFalse(isWindowsEnvironment());
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final ProvenanceEventRecord event = createEvent();

    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));

    final Query query = new Query(UUID.randomUUID().toString());

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(4000, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

public LuceneEventIndex(final RepositoryConfiguration config, final IndexManager indexManager, final int maxEventsPerCommit, final EventReporter eventReporter) {
    this.eventReporter = eventReporter;
    queryExecutor = Executors.newFixedThreadPool(config.getQueryThreadPoolSize(), new NamedThreadFactory("Provenance Query"));
    indexExecutor = Executors.newFixedThreadPool(config.getIndexThreadPoolSize(), new NamedThreadFactory("Index Provenance Events"));
    cacheWarmerExecutor = Executors.newScheduledThreadPool(config.getStorageDirectories().size(), new NamedThreadFactory("Warm Lucene Index", true));
    directoryManager = new IndexDirectoryManager(config);

    // Limit number of indexing threads to 100. When we restore the repository on restart,
    // we have to re-index up to MAX_THREADS * MAX_DOCUMENTS_PER_THREADS events prior to
    // the last event that the index holds. This is done because we could have that many
    // events 'in flight', waiting to be indexed when the last index writer was committed,
    // so even though the index says the largest event ID is 1,000,000 for instance, Event
    // with ID 999,999 may still not have been indexed because another thread was in the
    // process of writing the event to the index.
    final int configuredIndexPoolSize = config.getIndexThreadPoolSize();
    final int numIndexThreads;
    if (configuredIndexPoolSize > MAX_INDEX_THREADS) {
        logger.warn("The Provenance Repository is configured to perform indexing of events using {} threads. This number exceeds the maximum allowable number of threads, which is {}. "
            + "Will proceed using {} threads. This value is limited because the performance of indexing will decrease and startup times will increase when setting this value too high.",
            configuredIndexPoolSize, MAX_INDEX_THREADS, MAX_INDEX_THREADS);
        numIndexThreads = MAX_INDEX_THREADS;
    } else {
        numIndexThreads = configuredIndexPoolSize;
    }

    for (int i = 0; i < numIndexThreads; i++) {
        final EventIndexTask task = new EventIndexTask(documentQueue, config, indexManager, directoryManager, maxEventsPerCommit, eventReporter);
        indexTasks.add(task);
        indexExecutor.submit(task);
    }
    this.config = config;
    this.indexManager = indexManager;
    this.eventConverter = new ConvertEventToLuceneDocument(config.getSearchableFields(), config.getSearchableAttributes());
}
 

public QueryTask(final Query query, final ProgressiveResult result, final int maxResults, final IndexManager indexManager,
    final File indexDir, final EventStore eventStore, final EventAuthorizer authorizer,
    final EventTransformer unauthorizedTransformer) {
    this.query = query;
    this.queryResult = result;
    this.maxResults = maxResults;
    this.indexManager = indexManager;
    this.indexDir = indexDir;
    this.eventStore = eventStore;
    this.authorizer = authorizer;
    this.transformer = unauthorizedTransformer;
}
 

public EventIndexTask(final BlockingQueue<StoredDocument> documentQueue, final RepositoryConfiguration repoConfig, final IndexManager indexManager,
    final IndexDirectoryManager directoryManager, final int maxEventsPerCommit, final EventReporter eventReporter) {
    this.documentQueue = documentQueue;
    this.indexManager = indexManager;
    this.directoryManager = directoryManager;
    this.commitThreshold = maxEventsPerCommit;
    this.eventReporter = eventReporter;
}
 

@Test(timeout = 50000)
public void testQuerySpecificField() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 2, EventReporter.NO_OP);

    // add 2 events, one of which we will query for.
    final ProvenanceEventRecord event = createEvent();
    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));
    index.addEvent(createEvent(), new StorageSummary(2L, "1.prov", "1", 1, 2L, 2L));

    // Create a query that searches for the event with the FlowFile UUID equal to the first event's.
    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.FlowFileUUID, event.getFlowFileUuid()));

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(4000, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

protected IndexManager getIndexManager() {
    return indexManager;
}
 

public LuceneCacheWarmer(final File storageDir, final IndexManager indexManager) {
    this.storageDir = storageDir;
    this.indexManager = indexManager;
}
 

public LuceneEventIndex(final RepositoryConfiguration config, final IndexManager indexManager, final EventReporter eventReporter) {
    this(config, indexManager, EventIndexTask.DEFAULT_MAX_EVENTS_PER_COMMIT, eventReporter);
}
 

protected IndexManager getIndexManager() {
    return indexManager;
}
 

protected IndexManager getIndexManager() {
    return indexManager;
}
 

@Test(timeout = 50000)
public void testQuerySpecificField() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 2, EventReporter.NO_OP);

    // add 2 events, one of which we will query for.
    final ProvenanceEventRecord event = createEvent();
    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));
    index.addEvent(createEvent(), new StorageSummary(2L, "1.prov", "1", 1, 2L, 2L));

    // Create a query that searches for the event with the FlowFile UUID equal to the first event's.
    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.FlowFileUUID, event.getFlowFileUuid()));

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(100, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

public LuceneCacheWarmer(final File storageDir, final IndexManager indexManager) {
    this.storageDir = storageDir;
    this.indexManager = indexManager;
}
 

public LuceneEventIndex(final RepositoryConfiguration config, final IndexManager indexManager, final EventReporter eventReporter) {
    this(config, indexManager, EventIndexTask.DEFAULT_MAX_EVENTS_PER_COMMIT, eventReporter);
}