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

Example 6 with AtomicArray

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

the class ClearScrollControllerTests method testClearScrollIds.

public void testClearScrollIds() throws IOException, InterruptedException {
    DiscoveryNode node1 = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node2 = new DiscoveryNode("node_2", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node3 = new DiscoveryNode("node_3", buildNewFakeTransportAddress(), Version.CURRENT);
    AtomicArray<SearchPhaseResult> array = new AtomicArray<>(3);
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult1 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 1), node1);
    testSearchPhaseResult1.setSearchShardTarget(new SearchShardTarget("node_1", new ShardId("idx", "uuid1", 2), null, null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult2 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 12), node2);
    testSearchPhaseResult2.setSearchShardTarget(new SearchShardTarget("node_2", new ShardId("idy", "uuid2", 42), null, null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult3 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), 42), node3);
    testSearchPhaseResult3.setSearchShardTarget(new SearchShardTarget("node_3", new ShardId("idy", "uuid2", 43), null, null));
    array.setOnce(0, testSearchPhaseResult1);
    array.setOnce(1, testSearchPhaseResult2);
    array.setOnce(2, testSearchPhaseResult3);
    AtomicInteger numFreed = new AtomicInteger(0);
    String scrollId = TransportSearchHelper.buildScrollId(array, VersionUtils.randomVersion(random()));
    DiscoveryNodes nodes = DiscoveryNodes.builder().add(node1).add(node2).add(node3).build();
    CountDownLatch latch = new CountDownLatch(1);
    ActionListener<ClearScrollResponse> listener = new LatchedActionListener<>(new ActionListener<ClearScrollResponse>() {

        @Override
        public void onResponse(ClearScrollResponse clearScrollResponse) {
            assertEquals(numFreed.get(), clearScrollResponse.getNumFreed());
            assertTrue(clearScrollResponse.isSucceeded());
        }

        @Override
        public void onFailure(Exception e) {
            throw new AssertionError(e);
        }
    }, latch);
    List<DiscoveryNode> nodesInvoked = new CopyOnWriteArrayList<>();
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendFreeContext(Transport.Connection connection, ShardSearchContextId contextId, ActionListener<SearchFreeContextResponse> listener) {
            nodesInvoked.add(connection.getNode());
            boolean freed = randomBoolean();
            if (freed) {
                numFreed.incrementAndGet();
            }
            Thread t = new Thread(() -> listener.onResponse(new SearchFreeContextResponse(freed)));
            t.start();
        }

        @Override
        public Transport.Connection getConnection(String clusterAlias, DiscoveryNode node) {
            return new SearchAsyncActionTests.MockConnection(node);
        }
    };
    ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
    clearScrollRequest.scrollIds(Arrays.asList(scrollId));
    ClearScrollController controller = new ClearScrollController(clearScrollRequest, listener, nodes, logger, searchTransportService);
    controller.run();
    latch.await();
    assertEquals(3, nodesInvoked.size());
    Collections.sort(nodesInvoked, Comparator.comparing(DiscoveryNode::getId));
    assertEquals(nodesInvoked, Arrays.asList(node1, node2, node3));
}

Example 7 with AtomicArray

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

the class CountedCollectorTests method testCollect.

public void testCollect() throws InterruptedException {
    ArraySearchPhaseResults<SearchPhaseResult> consumer = new ArraySearchPhaseResults<>(randomIntBetween(1, 100));
    List<Integer> state = new ArrayList<>();
    int numResultsExpected = randomIntBetween(1, consumer.getAtomicArray().length());
    MockSearchPhaseContext context = new MockSearchPhaseContext(consumer.getAtomicArray().length());
    CountDownLatch latch = new CountDownLatch(1);
    boolean maybeFork = randomBoolean();
    Executor executor = (runnable) -> {
        if (randomBoolean() && maybeFork) {
            new Thread(runnable).start();
        } else {
            runnable.run();
        }
    };
    CountedCollector<SearchPhaseResult> collector = new CountedCollector<>(consumer, numResultsExpected, latch::countDown, context);
    for (int i = 0; i < numResultsExpected; i++) {
        int shardID = i;
        switch(randomIntBetween(0, 2)) {
            case 0:
                state.add(0);
                executor.execute(() -> collector.countDown());
                break;
            case 1:
                state.add(1);
                executor.execute(() -> {
                    DfsSearchResult dfsSearchResult = new DfsSearchResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), shardID), null, null);
                    dfsSearchResult.setShardIndex(shardID);
                    dfsSearchResult.setSearchShardTarget(new SearchShardTarget("foo", new ShardId("bar", "baz", shardID), null, OriginalIndices.NONE));
                    collector.onResult(dfsSearchResult);
                });
                break;
            case 2:
                state.add(2);
                executor.execute(() -> collector.onFailure(shardID, new SearchShardTarget("foo", new ShardId("bar", "baz", shardID), null, OriginalIndices.NONE), new RuntimeException("boom")));
                break;
            default:
                fail("unknown state");
        }
    }
    latch.await();
    assertEquals(numResultsExpected, state.size());
    AtomicArray<SearchPhaseResult> results = consumer.getAtomicArray();
    for (int i = 0; i < numResultsExpected; i++) {
        switch(state.get(i)) {
            case 0:
                assertNull(results.get(i));
                break;
            case 1:
                assertNotNull(results.get(i));
                assertEquals(i, results.get(i).getContextId().getId());
                break;
            case 2:
                final int shardId = i;
                assertEquals(1, context.failures.stream().filter(f -> f.shardId() == shardId).count());
                break;
            default:
                fail("unknown state");
        }
    }
    for (int i = numResultsExpected; i < results.length(); i++) {
        assertNull("index: " + i, results.get(i));
    }
}

Example 8 with AtomicArray

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

the class DfsQueryPhaseTests method testFailPhaseOnException.

public void testFailPhaseOnException() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<SearchPhaseResult>> responseRef = new AtomicReference<>();
    results.set(0, newSearchResult(0, new ShardSearchContextId("", 1), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.set(1, newSearchResult(1, new ShardSearchContextId("", 2), new SearchShardTarget("node2", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, SearchActionListener<QuerySearchResult> listener) {
            if (request.contextId().getId() == 1) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE), null);
                queryResult.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(42, 1.0F) }), 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else if (request.contextId().getId() == 2) {
                throw new UncheckedIOException(new MockDirectoryWrapper.FakeIOException());
            } else {
                fail("no such request ID: " + request.contextId());
            }
        }
    };
    SearchPhaseController searchPhaseController = searchPhaseController();
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.searchRequest, results.length(), exc -> {
    });
    DfsQueryPhase phase = new DfsQueryPhase(results.asList(), null, consumer, (response) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(response.results);
        }
    }, mockSearchPhaseContext);
    assertEquals("dfs_query", phase.getName());
    expectThrows(UncheckedIOException.class, phase::run);
    // phase execution will clean up on the contexts
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}

Example 9 with AtomicArray

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

the class SearchScrollAsyncActionTests method testAllShardsFailed.

public void testAllShardsFailed() 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);
    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 on shard", ex.getCause().getMessage());
                assertEquals("query", ex.getPhaseName());
                assertEquals("all shards 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(() -> searchActionListener.onFailure(new IllegalArgumentException("BOOM on shard"))).start();
        }

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

        @Override
        protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
            fail("don't move all shards failed");
            return null;
        }

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

Example 10 with AtomicArray

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

the class SearchScrollAsyncActionTests method testSendRequestsToNodes.

public void testSendRequestsToNodes() throws InterruptedException {
    ParsedScrollId scrollId = getParsedScrollId(new SearchContextIdForNode(null, "node1", new ShardSearchContextId(UUIDs.randomBase64UUID(), 1)), new SearchContextIdForNode(null, "node2", new ShardSearchContextId(UUIDs.randomBase64UUID(), 2)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId(UUIDs.randomBase64UUID(), 17)), new SearchContextIdForNode(null, "node1", new ShardSearchContextId(UUIDs.randomBase64UUID(), 0)), new SearchContextIdForNode(null, "node3", new ShardSearchContextId(UUIDs.randomBase64UUID(), 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(() -> {
                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(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());
    }
}