Examples with ReaderContext org.opensearch.search.internal.ReaderContext used on opensource projects

Example 11 with ReaderContext

use of org.opensearch.search.internal.ReaderContext in project OpenSearch by opensearch-project.

the class SearchService method putReaderContext.

protected void putReaderContext(ReaderContext context) {
    final ReaderContext previous = activeReaders.put(context.id().getId(), context);
    assert previous == null;
    // ensure that if we race against afterIndexRemoved, we remove the context from the active list.
    // this is important to ensure store can be cleaned up, in particular if the search is a scroll with a long timeout.
    final Index index = context.indexShard().shardId().getIndex();
    if (indicesService.hasIndex(index) == false) {
        removeReaderContext(context.id().getId());
        throw new IndexNotFoundException(index);
    }
}

Example 12 with ReaderContext

use of org.opensearch.search.internal.ReaderContext in project OpenSearch by opensearch-project.

the class SearchOperationListenerTests method testListenersAreExecuted.

// this test also tests if calls are correct if one or more listeners throw exceptions
public void testListenersAreExecuted() {
    AtomicInteger preQuery = new AtomicInteger();
    AtomicInteger failedQuery = new AtomicInteger();
    AtomicInteger onQuery = new AtomicInteger();
    AtomicInteger onFetch = new AtomicInteger();
    AtomicInteger preFetch = new AtomicInteger();
    AtomicInteger failedFetch = new AtomicInteger();
    AtomicInteger newContext = new AtomicInteger();
    AtomicInteger freeContext = new AtomicInteger();
    AtomicInteger newScrollContext = new AtomicInteger();
    AtomicInteger freeScrollContext = new AtomicInteger();
    AtomicInteger validateSearchContext = new AtomicInteger();
    AtomicInteger timeInNanos = new AtomicInteger(randomIntBetween(0, 10));
    SearchOperationListener listener = new SearchOperationListener() {

        @Override
        public void onPreQueryPhase(SearchContext searchContext) {
            assertNotNull(searchContext);
            preQuery.incrementAndGet();
        }

        @Override
        public void onFailedQueryPhase(SearchContext searchContext) {
            assertNotNull(searchContext);
            failedQuery.incrementAndGet();
        }

        @Override
        public void onQueryPhase(SearchContext searchContext, long tookInNanos) {
            assertEquals(timeInNanos.get(), tookInNanos);
            assertNotNull(searchContext);
            onQuery.incrementAndGet();
        }

        @Override
        public void onPreFetchPhase(SearchContext searchContext) {
            assertNotNull(searchContext);
            preFetch.incrementAndGet();
        }

        @Override
        public void onFailedFetchPhase(SearchContext searchContext) {
            assertNotNull(searchContext);
            failedFetch.incrementAndGet();
        }

        @Override
        public void onFetchPhase(SearchContext searchContext, long tookInNanos) {
            assertEquals(timeInNanos.get(), tookInNanos);
            onFetch.incrementAndGet();
        }

        @Override
        public void onNewReaderContext(ReaderContext readerContext) {
            assertNotNull(readerContext);
            newContext.incrementAndGet();
        }

        @Override
        public void onFreeReaderContext(ReaderContext readerContext) {
            assertNotNull(readerContext);
            freeContext.incrementAndGet();
        }

        @Override
        public void onNewScrollContext(ReaderContext readerContext) {
            assertNotNull(readerContext);
            newScrollContext.incrementAndGet();
        }

        @Override
        public void onFreeScrollContext(ReaderContext readerContext) {
            assertNotNull(readerContext);
            freeScrollContext.incrementAndGet();
        }

        @Override
        public void validateReaderContext(ReaderContext readerContext, TransportRequest request) {
            assertNotNull(readerContext);
            validateSearchContext.incrementAndGet();
        }
    };
    SearchOperationListener throwingListener = (SearchOperationListener) Proxy.newProxyInstance(SearchOperationListener.class.getClassLoader(), new Class[] { SearchOperationListener.class }, (a, b, c) -> {
        throw new RuntimeException();
    });
    int throwingListeners = 0;
    final List<SearchOperationListener> indexingOperationListeners = new ArrayList<>(Arrays.asList(listener, listener));
    if (randomBoolean()) {
        indexingOperationListeners.add(throwingListener);
        throwingListeners++;
        if (randomBoolean()) {
            indexingOperationListeners.add(throwingListener);
            throwingListeners++;
        }
    }
    Collections.shuffle(indexingOperationListeners, random());
    SearchOperationListener.CompositeListener compositeListener = new SearchOperationListener.CompositeListener(indexingOperationListeners, logger);
    SearchContext ctx = new TestSearchContext(null);
    compositeListener.onQueryPhase(ctx, timeInNanos.get());
    assertEquals(0, preFetch.get());
    assertEquals(0, preQuery.get());
    assertEquals(0, failedFetch.get());
    assertEquals(0, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(0, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onFetchPhase(ctx, timeInNanos.get());
    assertEquals(0, preFetch.get());
    assertEquals(0, preQuery.get());
    assertEquals(0, failedFetch.get());
    assertEquals(0, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onPreQueryPhase(ctx);
    assertEquals(0, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(0, failedFetch.get());
    assertEquals(0, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onPreFetchPhase(ctx);
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(0, failedFetch.get());
    assertEquals(0, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onFailedFetchPhase(ctx);
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(0, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onFailedQueryPhase(ctx);
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(0, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onNewReaderContext(mock(ReaderContext.class));
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(2, newContext.get());
    assertEquals(0, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onNewScrollContext(mock(ReaderContext.class));
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(2, newContext.get());
    assertEquals(2, newScrollContext.get());
    assertEquals(0, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onFreeReaderContext(mock(ReaderContext.class));
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(2, newContext.get());
    assertEquals(2, newScrollContext.get());
    assertEquals(2, freeContext.get());
    assertEquals(0, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    compositeListener.onFreeScrollContext(mock(ReaderContext.class));
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(2, newContext.get());
    assertEquals(2, newScrollContext.get());
    assertEquals(2, freeContext.get());
    assertEquals(2, freeScrollContext.get());
    assertEquals(0, validateSearchContext.get());
    if (throwingListeners == 0) {
        compositeListener.validateReaderContext(mock(ReaderContext.class), Empty.INSTANCE);
    } else {
        RuntimeException expected = expectThrows(RuntimeException.class, () -> compositeListener.validateReaderContext(mock(ReaderContext.class), Empty.INSTANCE));
        assertNull(expected.getMessage());
        assertEquals(throwingListeners - 1, expected.getSuppressed().length);
        if (throwingListeners > 1) {
            assertThat(expected.getSuppressed()[0], not(sameInstance(expected)));
        }
    }
    assertEquals(2, preFetch.get());
    assertEquals(2, preQuery.get());
    assertEquals(2, failedFetch.get());
    assertEquals(2, failedQuery.get());
    assertEquals(2, onQuery.get());
    assertEquals(2, onFetch.get());
    assertEquals(2, newContext.get());
    assertEquals(2, newScrollContext.get());
    assertEquals(2, freeContext.get());
    assertEquals(2, freeScrollContext.get());
    assertEquals(2, validateSearchContext.get());
}

Example 13 with ReaderContext

use of org.opensearch.search.internal.ReaderContext in project OpenSearch by opensearch-project.

the class SearchServiceTests method testTimeout.

public void testTimeout() throws IOException {
    createIndex("index");
    final SearchService service = getInstanceFromNode(SearchService.class);
    final IndicesService indicesService = getInstanceFromNode(IndicesService.class);
    final IndexService indexService = indicesService.indexServiceSafe(resolveIndex("index"));
    final IndexShard indexShard = indexService.getShard(0);
    SearchRequest searchRequest = new SearchRequest().allowPartialSearchResults(true);
    final ShardSearchRequest requestWithDefaultTimeout = new ShardSearchRequest(OriginalIndices.NONE, searchRequest, indexShard.shardId(), 1, new AliasFilter(null, Strings.EMPTY_ARRAY), 1.0f, -1, null, null);
    try (ReaderContext reader = createReaderContext(indexService, indexShard);
        SearchContext contextWithDefaultTimeout = service.createContext(reader, requestWithDefaultTimeout, null, randomBoolean())) {
        // the search context should inherit the default timeout
        assertThat(contextWithDefaultTimeout.timeout(), equalTo(TimeValue.timeValueSeconds(5)));
    }
    final long seconds = randomIntBetween(6, 10);
    searchRequest.source(new SearchSourceBuilder().timeout(TimeValue.timeValueSeconds(seconds)));
    final ShardSearchRequest requestWithCustomTimeout = new ShardSearchRequest(OriginalIndices.NONE, searchRequest, indexShard.shardId(), 1, new AliasFilter(null, Strings.EMPTY_ARRAY), 1.0f, -1, null, null);
    try (ReaderContext reader = createReaderContext(indexService, indexShard);
        SearchContext context = service.createContext(reader, requestWithCustomTimeout, null, randomBoolean())) {
        // the search context should inherit the query timeout
        assertThat(context.timeout(), equalTo(TimeValue.timeValueSeconds(seconds)));
    }
}

Example 14 with ReaderContext

use of org.opensearch.search.internal.ReaderContext in project OpenSearch by opensearch-project.

the class SearchServiceTests method testMaxScriptFieldsSearch.

/**
 * test that getting more than the allowed number of script_fields throws an exception
 */
public void testMaxScriptFieldsSearch() throws IOException {
    createIndex("index");
    final SearchService service = getInstanceFromNode(SearchService.class);
    final IndicesService indicesService = getInstanceFromNode(IndicesService.class);
    final IndexService indexService = indicesService.indexServiceSafe(resolveIndex("index"));
    final IndexShard indexShard = indexService.getShard(0);
    SearchRequest searchRequest = new SearchRequest().allowPartialSearchResults(true);
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchRequest.source(searchSourceBuilder);
    // adding the maximum allowed number of script_fields to retrieve
    int maxScriptFields = indexService.getIndexSettings().getMaxScriptFields();
    for (int i = 0; i < maxScriptFields; i++) {
        searchSourceBuilder.scriptField("field" + i, new Script(ScriptType.INLINE, MockScriptEngine.NAME, CustomScriptPlugin.DUMMY_SCRIPT, Collections.emptyMap()));
    }
    final ShardSearchRequest request = new ShardSearchRequest(OriginalIndices.NONE, searchRequest, indexShard.shardId(), 1, new AliasFilter(null, Strings.EMPTY_ARRAY), 1.0f, -1, null, null);
    try (ReaderContext reader = createReaderContext(indexService, indexShard)) {
        try (SearchContext context = service.createContext(reader, request, null, randomBoolean())) {
            assertNotNull(context);
        }
        searchSourceBuilder.scriptField("anotherScriptField", new Script(ScriptType.INLINE, MockScriptEngine.NAME, CustomScriptPlugin.DUMMY_SCRIPT, Collections.emptyMap()));
        IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> service.createContext(reader, request, null, randomBoolean()));
        assertEquals("Trying to retrieve too many script_fields. Must be less than or equal to: [" + maxScriptFields + "] but was [" + (maxScriptFields + 1) + "]. This limit can be set by changing the [index.max_script_fields] index level setting.", ex.getMessage());
    }
}

Example 15 with ReaderContext

use of org.opensearch.search.internal.ReaderContext in project OpenSearch by opensearch-project.

the class SearchServiceTests method testIgnoreScriptfieldIfSizeZero.

public void testIgnoreScriptfieldIfSizeZero() throws IOException {
    createIndex("index");
    final SearchService service = getInstanceFromNode(SearchService.class);
    final IndicesService indicesService = getInstanceFromNode(IndicesService.class);
    final IndexService indexService = indicesService.indexServiceSafe(resolveIndex("index"));
    final IndexShard indexShard = indexService.getShard(0);
    SearchRequest searchRequest = new SearchRequest().allowPartialSearchResults(true);
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchRequest.source(searchSourceBuilder);
    searchSourceBuilder.scriptField("field" + 0, new Script(ScriptType.INLINE, MockScriptEngine.NAME, CustomScriptPlugin.DUMMY_SCRIPT, Collections.emptyMap()));
    searchSourceBuilder.size(0);
    final ShardSearchRequest request = new ShardSearchRequest(OriginalIndices.NONE, searchRequest, indexShard.shardId(), 1, new AliasFilter(null, Strings.EMPTY_ARRAY), 1.0f, -1, null, null);
    try (ReaderContext reader = createReaderContext(indexService, indexShard);
        SearchContext context = service.createContext(reader, request, null, randomBoolean())) {
        assertEquals(0, context.scriptFields().fields().size());
    }
}