Examples with AtomicArray org.opensearch.common.util.concurrent.AtomicArray used on opensource projects

Example 11 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class SearchScrollAsyncActionTests method testFailNextPhase.

public void testFailNextPhase() throws InterruptedException {
    ParsedScrollId scrollId = getParsedScrollId(new SearchContextIdForNode(null, "node1", new ShardSearchContextId("", 1)), new SearchContextIdForNode(null, "node2", new ShardSearchContextId("a", 2)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("b", 17)), new SearchContextIdForNode(null, "node1", new ShardSearchContextId("c", 0)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("d", 0)));
    DiscoveryNodes discoveryNodes = DiscoveryNodes.builder().add(new DiscoveryNode("node1", buildNewFakeTransportAddress(), Version.CURRENT)).add(new DiscoveryNode("node2", buildNewFakeTransportAddress(), Version.CURRENT)).add(new DiscoveryNode("node3", buildNewFakeTransportAddress(), Version.CURRENT)).build();
    AtomicArray<SearchAsyncActionTests.TestSearchPhaseResult> results = new AtomicArray<>(scrollId.getContext().length);
    SearchScrollRequest request = new SearchScrollRequest();
    request.scroll(new Scroll(TimeValue.timeValueMinutes(1)));
    CountDownLatch latch = new CountDownLatch(1);
    AtomicInteger movedCounter = new AtomicInteger(0);
    ActionListener<SearchResponse> listener = new ActionListener<SearchResponse>() {

        @Override
        public void onResponse(SearchResponse o) {
            try {
                fail("got a result");
            } finally {
                latch.countDown();
            }
        }

        @Override
        public void onFailure(Exception e) {
            try {
                assertTrue(e instanceof SearchPhaseExecutionException);
                SearchPhaseExecutionException ex = (SearchPhaseExecutionException) e;
                assertEquals("BOOM", ex.getCause().getMessage());
                assertEquals("TEST_PHASE", ex.getPhaseName());
                assertEquals("Phase failed", ex.getMessage());
            } finally {
                latch.countDown();
            }
        }
    };
    SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult> action = new SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult>(scrollId, logger, discoveryNodes, listener, null, request, null) {

        @Override
        protected void executeInitialPhase(Transport.Connection connection, InternalScrollSearchRequest internalRequest, SearchActionListener<SearchAsyncActionTests.TestSearchPhaseResult> searchActionListener) {
            new Thread(() -> {
                SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult = new SearchAsyncActionTests.TestSearchPhaseResult(internalRequest.contextId(), connection.getNode());
                testSearchPhaseResult.setSearchShardTarget(new SearchShardTarget(connection.getNode().getId(), new ShardId("test", "_na_", 1), null, OriginalIndices.NONE));
                searchActionListener.onResponse(testSearchPhaseResult);
            }).start();
        }

        @Override
        protected Transport.Connection getConnection(String clusterAlias, DiscoveryNode node) {
            return new SearchAsyncActionTests.MockConnection(node);
        }

        @Override
        protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
            assertEquals(1, movedCounter.incrementAndGet());
            return new SearchPhase("TEST_PHASE") {

                @Override
                public void run() throws IOException {
                    throw new IllegalArgumentException("BOOM");
                }
            };
        }

        @Override
        protected void onFirstPhaseResult(int shardId, SearchAsyncActionTests.TestSearchPhaseResult result) {
            results.setOnce(shardId, result);
        }
    };
    action.run();
    latch.await();
    ShardSearchFailure[] shardSearchFailures = action.buildShardFailures();
    assertEquals(0, shardSearchFailures.length);
    SearchContextIdForNode[] context = scrollId.getContext();
    for (int i = 0; i < results.length(); i++) {
        assertNotNull(results.get(i));
        assertEquals(context[i].getSearchContextId(), results.get(i).getContextId());
        assertEquals(context[i].getNode(), results.get(i).node.getId());
    }
}

Example 12 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class SearchScrollAsyncActionTests method testNodeNotAvailable.

public void testNodeNotAvailable() throws InterruptedException {
    ParsedScrollId scrollId = getParsedScrollId(new SearchContextIdForNode(null, "node1", new ShardSearchContextId("", 1)), new SearchContextIdForNode(null, "node2", new ShardSearchContextId("", 2)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("", 17)), new SearchContextIdForNode(null, "node1", new ShardSearchContextId("", 0)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("", 0)));
    // node2 is not available
    DiscoveryNodes discoveryNodes = DiscoveryNodes.builder().add(new DiscoveryNode("node1", buildNewFakeTransportAddress(), Version.CURRENT)).add(new DiscoveryNode("node3", buildNewFakeTransportAddress(), Version.CURRENT)).build();
    AtomicArray<SearchAsyncActionTests.TestSearchPhaseResult> results = new AtomicArray<>(scrollId.getContext().length);
    SearchScrollRequest request = new SearchScrollRequest();
    request.scroll(new Scroll(TimeValue.timeValueMinutes(1)));
    CountDownLatch latch = new CountDownLatch(1);
    AtomicInteger movedCounter = new AtomicInteger(0);
    SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult> action = new SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult>(scrollId, logger, discoveryNodes, dummyListener(), null, request, null) {

        @Override
        protected void executeInitialPhase(Transport.Connection connection, InternalScrollSearchRequest internalRequest, SearchActionListener<SearchAsyncActionTests.TestSearchPhaseResult> searchActionListener) {
            try {
                assertNotEquals("node2 is not available", "node2", connection.getNode().getId());
            } catch (NullPointerException e) {
                logger.warn(e);
            }
            new Thread(() -> {
                SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult = new SearchAsyncActionTests.TestSearchPhaseResult(internalRequest.contextId(), connection.getNode());
                testSearchPhaseResult.setSearchShardTarget(new SearchShardTarget(connection.getNode().getId(), new ShardId("test", "_na_", 1), null, OriginalIndices.NONE));
                searchActionListener.onResponse(testSearchPhaseResult);
            }).start();
        }

        @Override
        protected Transport.Connection getConnection(String clusterAlias, DiscoveryNode node) {
            return new SearchAsyncActionTests.MockConnection(node);
        }

        @Override
        protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
            assertEquals(1, movedCounter.incrementAndGet());
            return new SearchPhase("test") {

                @Override
                public void run() throws IOException {
                    latch.countDown();
                }
            };
        }

        @Override
        protected void onFirstPhaseResult(int shardId, SearchAsyncActionTests.TestSearchPhaseResult result) {
            results.setOnce(shardId, result);
        }
    };
    action.run();
    latch.await();
    ShardSearchFailure[] shardSearchFailures = action.buildShardFailures();
    assertEquals(1, shardSearchFailures.length);
    assertEquals("IllegalStateException[node [node2] is not available]", shardSearchFailures[0].reason());
    SearchContextIdForNode[] context = scrollId.getContext();
    for (int i = 0; i < results.length(); i++) {
        if (context[i].getNode().equals("node2")) {
            assertNull(results.get(i));
        } else {
            assertNotNull(results.get(i));
            assertEquals(context[i].getSearchContextId(), results.get(i).getContextId());
            assertEquals(context[i].getNode(), results.get(i).node.getId());
        }
    }
}

Example 13 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method generateQueryResults.

/**
 * Generate random query results received from the provided number of shards, including the provided
 * number of search hits and randomly generated completion suggestions based on the name and size of the provided ones.
 * Note that <code>shardIndex</code> is already set to the generated completion suggestions to simulate what
 * {@link SearchPhaseController#reducedQueryPhase} does,
 * meaning that the returned query results can be fed directly to {@link SearchPhaseController#sortDocs}
 */
private static AtomicArray<SearchPhaseResult> generateQueryResults(int nShards, List<CompletionSuggestion> suggestions, int searchHitsSize, boolean useConstantScore) {
    AtomicArray<SearchPhaseResult> queryResults = new AtomicArray<>(nShards);
    for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
        String clusterAlias = randomBoolean() ? null : "remote";
        SearchShardTarget searchShardTarget = new SearchShardTarget("", new ShardId("", "", shardIndex), clusterAlias, OriginalIndices.NONE);
        QuerySearchResult querySearchResult = new QuerySearchResult(new ShardSearchContextId("", shardIndex), searchShardTarget, null);
        final TopDocs topDocs;
        float maxScore = 0;
        if (searchHitsSize == 0) {
            topDocs = Lucene.EMPTY_TOP_DOCS;
        } else {
            int nDocs = randomIntBetween(0, searchHitsSize);
            ScoreDoc[] scoreDocs = new ScoreDoc[nDocs];
            for (int i = 0; i < nDocs; i++) {
                float score = useConstantScore ? 1.0F : Math.abs(randomFloat());
                scoreDocs[i] = new ScoreDoc(i, score);
                maxScore = Math.max(score, maxScore);
            }
            topDocs = new TopDocs(new TotalHits(scoreDocs.length, TotalHits.Relation.EQUAL_TO), scoreDocs);
        }
        List<CompletionSuggestion> shardSuggestion = new ArrayList<>();
        for (CompletionSuggestion completionSuggestion : suggestions) {
            CompletionSuggestion suggestion = new CompletionSuggestion(completionSuggestion.getName(), completionSuggestion.getSize(), false);
            final CompletionSuggestion.Entry completionEntry = new CompletionSuggestion.Entry(new Text(""), 0, 5);
            suggestion.addTerm(completionEntry);
            int optionSize = randomIntBetween(1, suggestion.getSize());
            float maxScoreValue = randomIntBetween(suggestion.getSize(), (int) Float.MAX_VALUE);
            for (int i = 0; i < optionSize; i++) {
                completionEntry.addOption(new CompletionSuggestion.Entry.Option(i, new Text(""), maxScoreValue, Collections.emptyMap()));
                float dec = randomIntBetween(0, optionSize);
                if (dec <= maxScoreValue) {
                    maxScoreValue -= dec;
                }
            }
            suggestion.setShardIndex(shardIndex);
            shardSuggestion.add(suggestion);
        }
        querySearchResult.topDocs(new TopDocsAndMaxScore(topDocs, maxScore), null);
        querySearchResult.size(searchHitsSize);
        querySearchResult.suggest(new Suggest(new ArrayList<>(shardSuggestion)));
        querySearchResult.setShardIndex(shardIndex);
        queryResults.set(shardIndex, querySearchResult);
    }
    return queryResults;
}

Example 14 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method generateFetchResults.

private static AtomicArray<SearchPhaseResult> generateFetchResults(int nShards, ScoreDoc[] mergedSearchDocs, Suggest mergedSuggest) {
    AtomicArray<SearchPhaseResult> fetchResults = new AtomicArray<>(nShards);
    for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
        float maxScore = -1F;
        String clusterAlias = randomBoolean() ? null : "remote";
        SearchShardTarget shardTarget = new SearchShardTarget("", new ShardId("", "", shardIndex), clusterAlias, OriginalIndices.NONE);
        FetchSearchResult fetchSearchResult = new FetchSearchResult(new ShardSearchContextId("", shardIndex), shardTarget);
        List<SearchHit> searchHits = new ArrayList<>();
        for (ScoreDoc scoreDoc : mergedSearchDocs) {
            if (scoreDoc.shardIndex == shardIndex) {
                searchHits.add(new SearchHit(scoreDoc.doc, "", Collections.emptyMap(), Collections.emptyMap()));
                if (scoreDoc.score > maxScore) {
                    maxScore = scoreDoc.score;
                }
            }
        }
        for (Suggest.Suggestion<?> suggestion : mergedSuggest) {
            if (suggestion instanceof CompletionSuggestion) {
                for (CompletionSuggestion.Entry.Option option : ((CompletionSuggestion) suggestion).getOptions()) {
                    ScoreDoc doc = option.getDoc();
                    if (doc.shardIndex == shardIndex) {
                        searchHits.add(new SearchHit(doc.doc, "", Collections.emptyMap(), Collections.emptyMap()));
                        if (doc.score > maxScore) {
                            maxScore = doc.score;
                        }
                    }
                }
            }
        }
        SearchHit[] hits = searchHits.toArray(new SearchHit[0]);
        fetchSearchResult.hits(new SearchHits(hits, new TotalHits(hits.length, Relation.EQUAL_TO), maxScore));
        fetchResults.set(shardIndex, fetchSearchResult);
    }
    return fetchResults;
}

Example 15 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class SearchScrollAsyncActionTests method testShardFailures.

public void testShardFailures() throws InterruptedException {
    ParsedScrollId scrollId = getParsedScrollId(new SearchContextIdForNode(null, "node1", new ShardSearchContextId("", 1)), new SearchContextIdForNode(null, "node2", new ShardSearchContextId("", 2)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("", 17)), new SearchContextIdForNode(null, "node1", new ShardSearchContextId("", 0)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId("", 0)));
    DiscoveryNodes discoveryNodes = DiscoveryNodes.builder().add(new DiscoveryNode("node1", buildNewFakeTransportAddress(), Version.CURRENT)).add(new DiscoveryNode("node2", buildNewFakeTransportAddress(), Version.CURRENT)).add(new DiscoveryNode("node3", buildNewFakeTransportAddress(), Version.CURRENT)).build();
    AtomicArray<SearchAsyncActionTests.TestSearchPhaseResult> results = new AtomicArray<>(scrollId.getContext().length);
    SearchScrollRequest request = new SearchScrollRequest();
    request.scroll(new Scroll(TimeValue.timeValueMinutes(1)));
    CountDownLatch latch = new CountDownLatch(1);
    AtomicInteger movedCounter = new AtomicInteger(0);
    SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult> action = new SearchScrollAsyncAction<SearchAsyncActionTests.TestSearchPhaseResult>(scrollId, logger, discoveryNodes, dummyListener(), null, request, null) {

        @Override
        protected void executeInitialPhase(Transport.Connection connection, InternalScrollSearchRequest internalRequest, SearchActionListener<SearchAsyncActionTests.TestSearchPhaseResult> searchActionListener) {
            new Thread(() -> {
                if (internalRequest.contextId().getId() == 17) {
                    searchActionListener.onFailure(new IllegalArgumentException("BOOM on shard"));
                } else {
                    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult = new SearchAsyncActionTests.TestSearchPhaseResult(internalRequest.contextId(), connection.getNode());
                    testSearchPhaseResult.setSearchShardTarget(new SearchShardTarget(connection.getNode().getId(), new ShardId("test", "_na_", 1), null, OriginalIndices.NONE));
                    searchActionListener.onResponse(testSearchPhaseResult);
                }
            }).start();
        }

        @Override
        protected Transport.Connection getConnection(String clusterAlias, DiscoveryNode node) {
            return new SearchAsyncActionTests.MockConnection(node);
        }

        @Override
        protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
            assertEquals(1, movedCounter.incrementAndGet());
            return new SearchPhase("test") {

                @Override
                public void run() throws IOException {
                    latch.countDown();
                }
            };
        }

        @Override
        protected void onFirstPhaseResult(int shardId, SearchAsyncActionTests.TestSearchPhaseResult result) {
            results.setOnce(shardId, result);
        }
    };
    action.run();
    latch.await();
    ShardSearchFailure[] shardSearchFailures = action.buildShardFailures();
    assertEquals(1, shardSearchFailures.length);
    assertEquals("IllegalArgumentException[BOOM on shard]", shardSearchFailures[0].reason());
    SearchContextIdForNode[] context = scrollId.getContext();
    for (int i = 0; i < results.length(); i++) {
        if (context[i].getSearchContextId().getId() == 17) {
            assertNull(results.get(i));
        } else {
            assertNotNull(results.get(i));
            assertEquals(context[i].getSearchContextId(), results.get(i).getContextId());
            assertEquals(context[i].getNode(), results.get(i).node.getId());
        }
    }
}