Examples with ShardFetchRequest org.opensearch.search.fetch.ShardFetchRequest used on opensource projects

Example 1 with ShardFetchRequest

use of org.opensearch.search.fetch.ShardFetchRequest in project OpenSearch by opensearch-project.

the class SearchServiceTests method testSearchWhileIndexDeleted.

public void testSearchWhileIndexDeleted() throws InterruptedException {
    createIndex("index");
    client().prepareIndex("index").setId("1").setSource("field", "value").setRefreshPolicy(IMMEDIATE).get();
    SearchService service = getInstanceFromNode(SearchService.class);
    IndicesService indicesService = getInstanceFromNode(IndicesService.class);
    IndexService indexService = indicesService.indexServiceSafe(resolveIndex("index"));
    IndexShard indexShard = indexService.getShard(0);
    AtomicBoolean running = new AtomicBoolean(true);
    CountDownLatch startGun = new CountDownLatch(1);
    Semaphore semaphore = new Semaphore(Integer.MAX_VALUE);
    final Thread thread = new Thread() {

        @Override
        public void run() {
            startGun.countDown();
            while (running.get()) {
                service.afterIndexRemoved(indexService.index(), indexService.getIndexSettings(), DELETED);
                if (randomBoolean()) {
                    // context in a non-sane way.
                    try {
                        semaphore.acquire();
                    } catch (InterruptedException e) {
                        throw new AssertionError(e);
                    }
                    client().prepareIndex("index").setSource("field", "value").setRefreshPolicy(randomFrom(WriteRequest.RefreshPolicy.values())).execute(new ActionListener<IndexResponse>() {

                        @Override
                        public void onResponse(IndexResponse indexResponse) {
                            semaphore.release();
                        }

                        @Override
                        public void onFailure(Exception e) {
                            semaphore.release();
                        }
                    });
                }
            }
        }
    };
    thread.start();
    startGun.await();
    try {
        final int rounds = scaledRandomIntBetween(100, 1000);
        SearchRequest searchRequest = new SearchRequest().allowPartialSearchResults(true);
        SearchRequest scrollSearchRequest = new SearchRequest().allowPartialSearchResults(true).scroll(new Scroll(TimeValue.timeValueMinutes(1)));
        for (int i = 0; i < rounds; i++) {
            try {
                try {
                    PlainActionFuture<SearchPhaseResult> result = new PlainActionFuture<>();
                    final boolean useScroll = randomBoolean();
                    service.executeQueryPhase(new ShardSearchRequest(OriginalIndices.NONE, useScroll ? scrollSearchRequest : searchRequest, indexShard.shardId(), 1, new AliasFilter(null, Strings.EMPTY_ARRAY), 1.0f, -1, null, null), true, new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()), result);
                    SearchPhaseResult searchPhaseResult = result.get();
                    IntArrayList intCursors = new IntArrayList(1);
                    intCursors.add(0);
                    ShardFetchRequest req = new ShardFetchRequest(searchPhaseResult.getContextId(), intCursors, null);
                    PlainActionFuture<FetchSearchResult> listener = new PlainActionFuture<>();
                    service.executeFetchPhase(req, new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()), listener);
                    listener.get();
                    if (useScroll) {
                        // have to free context since this test does not remove the index from IndicesService.
                        service.freeReaderContext(searchPhaseResult.getContextId());
                    }
                } catch (ExecutionException ex) {
                    assertThat(ex.getCause(), instanceOf(RuntimeException.class));
                    throw ((RuntimeException) ex.getCause());
                }
            } catch (AlreadyClosedException ex) {
                throw ex;
            } catch (IllegalStateException ex) {
                assertEquals("reader_context is already closed can't increment refCount current count [0]", ex.getMessage());
            } catch (SearchContextMissingException ex) {
            // that's fine
            }
        }
    } finally {
        running.set(false);
        thread.join();
        semaphore.acquire(Integer.MAX_VALUE);
    }
    assertEquals(0, service.getActiveContexts());
    SearchStats.Stats totalStats = indexShard.searchStats().getTotal();
    assertEquals(0, totalStats.getQueryCurrent());
    assertEquals(0, totalStats.getScrollCurrent());
    assertEquals(0, totalStats.getFetchCurrent());
}

Example 2 with ShardFetchRequest

use of org.opensearch.search.fetch.ShardFetchRequest in project OpenSearch by opensearch-project.

the class SearchScrollQueryThenFetchAsyncAction method moveToNextPhase.

@Override
protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
    return new SearchPhase("fetch") {

        @Override
        public void run() {
            final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = searchPhaseController.reducedScrollQueryPhase(queryResults.asList());
            ScoreDoc[] scoreDocs = reducedQueryPhase.sortedTopDocs.scoreDocs;
            if (scoreDocs.length == 0) {
                sendResponse(reducedQueryPhase, fetchResults);
                return;
            }
            final IntArrayList[] docIdsToLoad = searchPhaseController.fillDocIdsToLoad(queryResults.length(), scoreDocs);
            final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(reducedQueryPhase, queryResults.length());
            final CountDown counter = new CountDown(docIdsToLoad.length);
            for (int i = 0; i < docIdsToLoad.length; i++) {
                final int index = i;
                final IntArrayList docIds = docIdsToLoad[index];
                if (docIds != null) {
                    final QuerySearchResult querySearchResult = queryResults.get(index);
                    ScoreDoc lastEmittedDoc = lastEmittedDocPerShard[index];
                    ShardFetchRequest shardFetchRequest = new ShardFetchRequest(querySearchResult.getContextId(), docIds, lastEmittedDoc);
                    SearchShardTarget searchShardTarget = querySearchResult.getSearchShardTarget();
                    DiscoveryNode node = clusterNodeLookup.apply(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId());
                    assert node != null : "target node is null in secondary phase";
                    Transport.Connection connection = getConnection(searchShardTarget.getClusterAlias(), node);
                    searchTransportService.sendExecuteFetchScroll(connection, shardFetchRequest, task, new SearchActionListener<FetchSearchResult>(querySearchResult.getSearchShardTarget(), index) {

                        @Override
                        protected void innerOnResponse(FetchSearchResult response) {
                            fetchResults.setOnce(response.getShardIndex(), response);
                            if (counter.countDown()) {
                                sendResponse(reducedQueryPhase, fetchResults);
                            }
                        }

                        @Override
                        public void onFailure(Exception t) {
                            onShardFailure(getName(), counter, querySearchResult.getContextId(), t, querySearchResult.getSearchShardTarget(), () -> sendResponsePhase(reducedQueryPhase, fetchResults));
                        }
                    });
                } else {
                    // which can have null values so we have to count them down too
                    if (counter.countDown()) {
                        sendResponse(reducedQueryPhase, fetchResults);
                    }
                }
            }
        }
    };
}