Examples with SearchOperationListener org.opensearch.index.shard.SearchOperationListener used on opensource projects

Example 1 with SearchOperationListener

use of org.opensearch.index.shard.SearchOperationListener in project OpenSearch by opensearch-project.

the class SearchService method createAndPutReaderContext.

final ReaderContext createAndPutReaderContext(ShardSearchRequest request, IndexService indexService, IndexShard shard, Engine.SearcherSupplier reader, boolean keepStatesInContext) {
    assert request.readerId() == null;
    assert request.keepAlive() == null;
    ReaderContext readerContext = null;
    Releasable decreaseScrollContexts = null;
    try {
        if (request.scroll() != null) {
            decreaseScrollContexts = openScrollContexts::decrementAndGet;
            if (openScrollContexts.incrementAndGet() > maxOpenScrollContext) {
                throw new OpenSearchRejectedExecutionException("Trying to create too many scroll contexts. Must be less than or equal to: [" + maxOpenScrollContext + "]. " + "This limit can be set by changing the [" + MAX_OPEN_SCROLL_CONTEXT.getKey() + "] setting.");
            }
        }
        final long keepAlive = getKeepAlive(request);
        final ShardSearchContextId id = new ShardSearchContextId(sessionId, idGenerator.incrementAndGet());
        if (keepStatesInContext || request.scroll() != null) {
            readerContext = new LegacyReaderContext(id, indexService, shard, reader, request, keepAlive);
            if (request.scroll() != null) {
                readerContext.addOnClose(decreaseScrollContexts);
                decreaseScrollContexts = null;
            }
        } else {
            readerContext = new ReaderContext(id, indexService, shard, reader, keepAlive, request.keepAlive() == null);
        }
        reader = null;
        final ReaderContext finalReaderContext = readerContext;
        final SearchOperationListener searchOperationListener = shard.getSearchOperationListener();
        searchOperationListener.onNewReaderContext(finalReaderContext);
        if (finalReaderContext.scrollContext() != null) {
            searchOperationListener.onNewScrollContext(finalReaderContext);
        }
        readerContext.addOnClose(() -> {
            try {
                if (finalReaderContext.scrollContext() != null) {
                    searchOperationListener.onFreeScrollContext(finalReaderContext);
                }
            } finally {
                searchOperationListener.onFreeReaderContext(finalReaderContext);
            }
        });
        putReaderContext(finalReaderContext);
        readerContext = null;
        return finalReaderContext;
    } finally {
        Releasables.close(reader, readerContext, decreaseScrollContexts);
    }
}

Example 2 with SearchOperationListener

use of org.opensearch.index.shard.SearchOperationListener in project OpenSearch by opensearch-project.

the class SearchService method openReaderContext.

/**
 * Opens the reader context for given shardId. The newly opened reader context will be keep
 * until the {@code keepAlive} elapsed unless it is manually released.
 */
public void openReaderContext(ShardId shardId, TimeValue keepAlive, ActionListener<ShardSearchContextId> listener) {
    checkKeepAliveLimit(keepAlive.millis());
    final IndexService indexService = indicesService.indexServiceSafe(shardId.getIndex());
    final IndexShard shard = indexService.getShard(shardId.id());
    final SearchOperationListener searchOperationListener = shard.getSearchOperationListener();
    shard.awaitShardSearchActive(ignored -> {
        Engine.SearcherSupplier searcherSupplier = null;
        ReaderContext readerContext = null;
        try {
            searcherSupplier = shard.acquireSearcherSupplier();
            final ShardSearchContextId id = new ShardSearchContextId(sessionId, idGenerator.incrementAndGet());
            readerContext = new ReaderContext(id, indexService, shard, searcherSupplier, keepAlive.millis(), false);
            final ReaderContext finalReaderContext = readerContext;
            // transfer ownership to reader context
            searcherSupplier = null;
            searchOperationListener.onNewReaderContext(readerContext);
            readerContext.addOnClose(() -> searchOperationListener.onFreeReaderContext(finalReaderContext));
            putReaderContext(readerContext);
            readerContext = null;
            listener.onResponse(finalReaderContext.id());
        } catch (Exception exc) {
            Releasables.closeWhileHandlingException(searcherSupplier, readerContext);
            listener.onFailure(exc);
        }
    });
}

Example 3 with SearchOperationListener

use of org.opensearch.index.shard.SearchOperationListener in project OpenSearch by opensearch-project.

the class IndexModuleTests method testAddSearchOperationListener.

public void testAddSearchOperationListener() throws IOException {
    final IndexSettings indexSettings = IndexSettingsModule.newIndexSettings(index, settings);
    IndexModule module = createIndexModule(indexSettings, emptyAnalysisRegistry);
    AtomicBoolean executed = new AtomicBoolean(false);
    SearchOperationListener listener = new SearchOperationListener() {

        @Override
        public void onNewReaderContext(ReaderContext readerContext) {
            executed.set(true);
        }
    };
    module.addSearchOperationListener(listener);
    expectThrows(IllegalArgumentException.class, () -> module.addSearchOperationListener(listener));
    expectThrows(IllegalArgumentException.class, () -> module.addSearchOperationListener(null));
    IndexService indexService = newIndexService(module);
    assertEquals(2, indexService.getSearchOperationListener().size());
    assertEquals(SearchSlowLog.class, indexService.getSearchOperationListener().get(0).getClass());
    assertSame(listener, indexService.getSearchOperationListener().get(1));
    for (SearchOperationListener l : indexService.getSearchOperationListener()) {
        l.onNewReaderContext(mock(ReaderContext.class));
    }
    assertTrue(executed.get());
    indexService.close("simon says", false);
}