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

Example 1 with InternalScrollSearchRequest

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

the class SearchScrollAsyncAction method run.

private void run(BiFunction<String, String, DiscoveryNode> clusterNodeLookup, final SearchContextIdForNode[] context) {
    final CountDown counter = new CountDown(scrollId.getContext().length);
    for (int i = 0; i < context.length; i++) {
        SearchContextIdForNode target = context[i];
        final int shardIndex = i;
        final Transport.Connection connection;
        try {
            DiscoveryNode node = clusterNodeLookup.apply(target.getClusterAlias(), target.getNode());
            if (node == null) {
                throw new IllegalStateException("node [" + target.getNode() + "] is not available");
            }
            connection = getConnection(target.getClusterAlias(), node);
        } catch (Exception ex) {
            onShardFailure("query", counter, target.getSearchContextId(), ex, null, () -> SearchScrollAsyncAction.this.moveToNextPhase(clusterNodeLookup));
            continue;
        }
        final InternalScrollSearchRequest internalRequest = TransportSearchHelper.internalScrollSearchRequest(target.getSearchContextId(), request);
        // we can't create a SearchShardTarget here since we don't know the index and shard ID we are talking to
        // we only know the node and the search context ID. Yet, the response will contain the SearchShardTarget
        // from the target node instead...that's why we pass null here
        SearchActionListener<T> searchActionListener = new SearchActionListener<T>(null, shardIndex) {

            @Override
            protected void setSearchShardTarget(T response) {
                // don't do this - it's part of the response...
                assert response.getSearchShardTarget() != null : "search shard target must not be null";
                if (target.getClusterAlias() != null) {
                    // re-create the search target and add the cluster alias if there is any,
                    // we need this down the road for subseq. phases
                    SearchShardTarget searchShardTarget = response.getSearchShardTarget();
                    response.setSearchShardTarget(new SearchShardTarget(searchShardTarget.getNodeId(), searchShardTarget.getShardId(), target.getClusterAlias(), null));
                }
            }

            @Override
            protected void innerOnResponse(T result) {
                assert shardIndex == result.getShardIndex() : "shard index mismatch: " + shardIndex + " but got: " + result.getShardIndex();
                onFirstPhaseResult(shardIndex, result);
                if (counter.countDown()) {
                    SearchPhase phase = moveToNextPhase(clusterNodeLookup);
                    try {
                        phase.run();
                    } catch (Exception e) {
                        // we need to fail the entire request here - the entire phase just blew up
                        // don't call onShardFailure or onFailure here since otherwise we'd countDown the counter
                        // again which would result in an exception
                        listener.onFailure(new SearchPhaseExecutionException(phase.getName(), "Phase failed", e, ShardSearchFailure.EMPTY_ARRAY));
                    }
                }
            }

            @Override
            public void onFailure(Exception t) {
                onShardFailure("query", counter, target.getSearchContextId(), t, null, () -> SearchScrollAsyncAction.this.moveToNextPhase(clusterNodeLookup));
            }
        };
        executeInitialPhase(connection, internalRequest, searchActionListener);
    }
}

Example 2 with InternalScrollSearchRequest

use of org.opensearch.search.internal.InternalScrollSearchRequest 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 3 with InternalScrollSearchRequest

use of org.opensearch.search.internal.InternalScrollSearchRequest 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 4 with InternalScrollSearchRequest

use of org.opensearch.search.internal.InternalScrollSearchRequest 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());
    }
}

Example 5 with InternalScrollSearchRequest

use of org.opensearch.search.internal.InternalScrollSearchRequest 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));
    }
}