Examples with SearchShardTarget org.opensearch.search.SearchShardTarget used on opensource projects

Example 26 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class SearchAsyncActionTests method testLimitConcurrentShardRequests.

public void testLimitConcurrentShardRequests() throws InterruptedException {
    SearchRequest request = new SearchRequest();
    request.allowPartialSearchResults(true);
    int numConcurrent = randomIntBetween(1, 5);
    request.setMaxConcurrentShardRequests(numConcurrent);
    boolean doReplicas = randomBoolean();
    int numShards = randomIntBetween(5, 10);
    int numShardAttempts = numShards;
    Boolean[] shardFailures = new Boolean[numShards];
    // at least one response otherwise the entire request fails
    shardFailures[randomIntBetween(0, shardFailures.length - 1)] = false;
    for (int i = 0; i < shardFailures.length; i++) {
        if (shardFailures[i] == null) {
            boolean failure = randomBoolean();
            shardFailures[i] = failure;
            if (failure && doReplicas) {
                numShardAttempts++;
            }
        }
    }
    CountDownLatch latch = new CountDownLatch(numShardAttempts);
    AtomicBoolean searchPhaseDidRun = new AtomicBoolean(false);
    ActionListener<SearchResponse> responseListener = ActionListener.wrap(response -> {
    }, (e) -> {
        throw new AssertionError("unexpected", e);
    });
    DiscoveryNode primaryNode = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    // for the sake of this test we place the replica on the same node. ie. this is not a mistake since we limit per node now
    DiscoveryNode replicaNode = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    AtomicInteger contextIdGenerator = new AtomicInteger(0);
    GroupShardsIterator<SearchShardIterator> shardsIter = getShardsIter("idx", new OriginalIndices(new String[] { "idx" }, SearchRequest.DEFAULT_INDICES_OPTIONS), numShards, doReplicas, primaryNode, replicaNode);
    SearchTransportService transportService = new SearchTransportService(null, null);
    Map<String, Transport.Connection> lookup = new HashMap<>();
    Map<ShardId, Boolean> seenShard = new ConcurrentHashMap<>();
    lookup.put(primaryNode.getId(), new MockConnection(primaryNode));
    lookup.put(replicaNode.getId(), new MockConnection(replicaNode));
    Map<String, AliasFilter> aliasFilters = Collections.singletonMap("_na_", new AliasFilter(null, Strings.EMPTY_ARRAY));
    CountDownLatch awaitInitialRequests = new CountDownLatch(1);
    AtomicInteger numRequests = new AtomicInteger(0);
    AbstractSearchAsyncAction<TestSearchPhaseResult> asyncAction = new AbstractSearchAsyncAction<TestSearchPhaseResult>("test", logger, transportService, (cluster, node) -> {
        assert cluster == null : "cluster was not null: " + cluster;
        return lookup.get(node);
    }, aliasFilters, Collections.emptyMap(), Collections.emptyMap(), null, request, responseListener, shardsIter, new TransportSearchAction.SearchTimeProvider(0, 0, () -> 0), ClusterState.EMPTY_STATE, null, new ArraySearchPhaseResults<>(shardsIter.size()), request.getMaxConcurrentShardRequests(), SearchResponse.Clusters.EMPTY) {

        @Override
        protected void executePhaseOnShard(SearchShardIterator shardIt, SearchShardTarget shard, SearchActionListener<TestSearchPhaseResult> listener) {
            seenShard.computeIfAbsent(shard.getShardId(), (i) -> {
                // only count this once per shard copy
                numRequests.incrementAndGet();
                return Boolean.TRUE;
            });
            new Thread(() -> {
                try {
                    awaitInitialRequests.await();
                } catch (InterruptedException e) {
                    throw new AssertionError(e);
                }
                Transport.Connection connection = getConnection(null, shard.getNodeId());
                TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()), connection.getNode());
                if (shardFailures[shard.getShardId().id()]) {
                    listener.onFailure(new RuntimeException());
                } else {
                    listener.onResponse(testSearchPhaseResult);
                }
            }).start();
        }

        @Override
        protected SearchPhase getNextPhase(SearchPhaseResults<TestSearchPhaseResult> results, SearchPhaseContext context) {
            return new SearchPhase("test") {

                @Override
                public void run() {
                    assertTrue(searchPhaseDidRun.compareAndSet(false, true));
                }
            };
        }

        @Override
        protected void executeNext(Runnable runnable, Thread originalThread) {
            super.executeNext(runnable, originalThread);
            latch.countDown();
        }
    };
    asyncAction.start();
    assertEquals(numConcurrent, numRequests.get());
    awaitInitialRequests.countDown();
    latch.await();
    assertTrue(searchPhaseDidRun.get());
    assertEquals(numShards, numRequests.get());
}

Example 27 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class SearchAsyncActionTests method testSkipSearchShards.

public void testSkipSearchShards() throws InterruptedException {
    SearchRequest request = new SearchRequest();
    request.allowPartialSearchResults(true);
    int numShards = 10;
    ActionListener<SearchResponse> responseListener = ActionListener.wrap(response -> {
    }, (e) -> {
        throw new AssertionError("unexpected", e);
    });
    DiscoveryNode primaryNode = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode replicaNode = new DiscoveryNode("node_2", buildNewFakeTransportAddress(), Version.CURRENT);
    AtomicInteger contextIdGenerator = new AtomicInteger(0);
    GroupShardsIterator<SearchShardIterator> shardsIter = getShardsIter("idx", new OriginalIndices(new String[] { "idx" }, SearchRequest.DEFAULT_INDICES_OPTIONS), numShards, randomBoolean(), primaryNode, replicaNode);
    int numSkipped = 0;
    for (SearchShardIterator iter : shardsIter) {
        if (iter.shardId().id() % 2 == 0) {
            iter.resetAndSkip();
            numSkipped++;
        }
    }
    CountDownLatch latch = new CountDownLatch(numShards - numSkipped);
    AtomicBoolean searchPhaseDidRun = new AtomicBoolean(false);
    SearchTransportService transportService = new SearchTransportService(null, null);
    Map<String, Transport.Connection> lookup = new HashMap<>();
    Map<ShardId, Boolean> seenShard = new ConcurrentHashMap<>();
    lookup.put(primaryNode.getId(), new MockConnection(primaryNode));
    lookup.put(replicaNode.getId(), new MockConnection(replicaNode));
    Map<String, AliasFilter> aliasFilters = Collections.singletonMap("_na_", new AliasFilter(null, Strings.EMPTY_ARRAY));
    AtomicInteger numRequests = new AtomicInteger(0);
    AbstractSearchAsyncAction<TestSearchPhaseResult> asyncAction = new AbstractSearchAsyncAction<TestSearchPhaseResult>("test", logger, transportService, (cluster, node) -> {
        assert cluster == null : "cluster was not null: " + cluster;
        return lookup.get(node);
    }, aliasFilters, Collections.emptyMap(), Collections.emptyMap(), null, request, responseListener, shardsIter, new TransportSearchAction.SearchTimeProvider(0, 0, () -> 0), ClusterState.EMPTY_STATE, null, new ArraySearchPhaseResults<>(shardsIter.size()), request.getMaxConcurrentShardRequests(), SearchResponse.Clusters.EMPTY) {

        @Override
        protected void executePhaseOnShard(SearchShardIterator shardIt, SearchShardTarget shard, SearchActionListener<TestSearchPhaseResult> listener) {
            seenShard.computeIfAbsent(shard.getShardId(), (i) -> {
                // only count this once per replica
                numRequests.incrementAndGet();
                return Boolean.TRUE;
            });
            new Thread(() -> {
                Transport.Connection connection = getConnection(null, shard.getNodeId());
                TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()), connection.getNode());
                listener.onResponse(testSearchPhaseResult);
            }).start();
        }

        @Override
        protected SearchPhase getNextPhase(SearchPhaseResults<TestSearchPhaseResult> results, SearchPhaseContext context) {
            return new SearchPhase("test") {

                @Override
                public void run() {
                    assertTrue(searchPhaseDidRun.compareAndSet(false, true));
                }
            };
        }

        @Override
        protected void executeNext(Runnable runnable, Thread originalThread) {
            super.executeNext(runnable, originalThread);
            latch.countDown();
        }
    };
    asyncAction.start();
    latch.await();
    assertTrue(searchPhaseDidRun.get());
    SearchResponse searchResponse = asyncAction.buildSearchResponse(null, asyncAction.buildShardFailures(), null, null);
    assertEquals(shardsIter.size() - numSkipped, numRequests.get());
    assertEquals(0, searchResponse.getFailedShards());
    assertEquals(numSkipped, searchResponse.getSkippedShards());
    assertEquals(shardsIter.size(), searchResponse.getSuccessfulShards());
}

Example 28 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class TransportSearchActionTests method testProcessRemoteShards.

public void testProcessRemoteShards() {
    try (TransportService transportService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, null)) {
        RemoteClusterService service = transportService.getRemoteClusterService();
        assertFalse(service.isCrossClusterSearchEnabled());
        Map<String, ClusterSearchShardsResponse> searchShardsResponseMap = new HashMap<>();
        DiscoveryNode[] nodes = new DiscoveryNode[] { new DiscoveryNode("node1", buildNewFakeTransportAddress(), Version.CURRENT), new DiscoveryNode("node2", buildNewFakeTransportAddress(), Version.CURRENT) };
        Map<String, AliasFilter> indicesAndAliases = new HashMap<>();
        indicesAndAliases.put("foo", new AliasFilter(new TermsQueryBuilder("foo", "bar"), "some_alias_for_foo", "some_other_foo_alias"));
        indicesAndAliases.put("bar", new AliasFilter(new MatchAllQueryBuilder(), Strings.EMPTY_ARRAY));
        ClusterSearchShardsGroup[] groups = new ClusterSearchShardsGroup[] { new ClusterSearchShardsGroup(new ShardId("foo", "foo_id", 0), new ShardRouting[] { TestShardRouting.newShardRouting("foo", 0, "node1", true, ShardRoutingState.STARTED), TestShardRouting.newShardRouting("foo", 0, "node2", false, ShardRoutingState.STARTED) }), new ClusterSearchShardsGroup(new ShardId("foo", "foo_id", 1), new ShardRouting[] { TestShardRouting.newShardRouting("foo", 0, "node1", true, ShardRoutingState.STARTED), TestShardRouting.newShardRouting("foo", 1, "node2", false, ShardRoutingState.STARTED) }), new ClusterSearchShardsGroup(new ShardId("bar", "bar_id", 0), new ShardRouting[] { TestShardRouting.newShardRouting("bar", 0, "node2", true, ShardRoutingState.STARTED), TestShardRouting.newShardRouting("bar", 0, "node1", false, ShardRoutingState.STARTED) }) };
        searchShardsResponseMap.put("test_cluster_1", new ClusterSearchShardsResponse(groups, nodes, indicesAndAliases));
        DiscoveryNode[] nodes2 = new DiscoveryNode[] { new DiscoveryNode("node3", buildNewFakeTransportAddress(), Version.CURRENT) };
        ClusterSearchShardsGroup[] groups2 = new ClusterSearchShardsGroup[] { new ClusterSearchShardsGroup(new ShardId("xyz", "xyz_id", 0), new ShardRouting[] { TestShardRouting.newShardRouting("xyz", 0, "node3", true, ShardRoutingState.STARTED) }) };
        Map<String, AliasFilter> filter = new HashMap<>();
        filter.put("xyz", new AliasFilter(null, "some_alias_for_xyz"));
        searchShardsResponseMap.put("test_cluster_2", new ClusterSearchShardsResponse(groups2, nodes2, filter));
        Map<String, OriginalIndices> remoteIndicesByCluster = new HashMap<>();
        remoteIndicesByCluster.put("test_cluster_1", new OriginalIndices(new String[] { "fo*", "ba*" }, SearchRequest.DEFAULT_INDICES_OPTIONS));
        remoteIndicesByCluster.put("test_cluster_2", new OriginalIndices(new String[] { "x*" }, SearchRequest.DEFAULT_INDICES_OPTIONS));
        Map<String, AliasFilter> remoteAliases = TransportSearchAction.getRemoteAliasFilters(searchShardsResponseMap);
        List<SearchShardIterator> iteratorList = TransportSearchAction.getRemoteShardsIterator(searchShardsResponseMap, remoteIndicesByCluster, remoteAliases);
        assertEquals(4, iteratorList.size());
        for (SearchShardIterator iterator : iteratorList) {
            if (iterator.shardId().getIndexName().endsWith("foo")) {
                assertArrayEquals(new String[] { "some_alias_for_foo", "some_other_foo_alias" }, iterator.getOriginalIndices().indices());
                assertTrue(iterator.shardId().getId() == 0 || iterator.shardId().getId() == 1);
                assertEquals("test_cluster_1", iterator.getClusterAlias());
                assertEquals("foo", iterator.shardId().getIndexName());
                SearchShardTarget shard = iterator.nextOrNull();
                assertNotNull(shard);
                assertEquals(shard.getShardId().getIndexName(), "foo");
                shard = iterator.nextOrNull();
                assertNotNull(shard);
                assertEquals(shard.getShardId().getIndexName(), "foo");
                assertNull(iterator.nextOrNull());
            } else if (iterator.shardId().getIndexName().endsWith("bar")) {
                assertArrayEquals(new String[] { "bar" }, iterator.getOriginalIndices().indices());
                assertEquals(0, iterator.shardId().getId());
                assertEquals("test_cluster_1", iterator.getClusterAlias());
                assertEquals("bar", iterator.shardId().getIndexName());
                SearchShardTarget shard = iterator.nextOrNull();
                assertNotNull(shard);
                assertEquals(shard.getShardId().getIndexName(), "bar");
                shard = iterator.nextOrNull();
                assertNotNull(shard);
                assertEquals(shard.getShardId().getIndexName(), "bar");
                assertNull(iterator.nextOrNull());
            } else if (iterator.shardId().getIndexName().endsWith("xyz")) {
                assertArrayEquals(new String[] { "some_alias_for_xyz" }, iterator.getOriginalIndices().indices());
                assertEquals(0, iterator.shardId().getId());
                assertEquals("xyz", iterator.shardId().getIndexName());
                assertEquals("test_cluster_2", iterator.getClusterAlias());
                SearchShardTarget shard = iterator.nextOrNull();
                assertNotNull(shard);
                assertEquals(shard.getShardId().getIndexName(), "xyz");
                assertNull(iterator.nextOrNull());
            }
        }
        assertEquals(3, remoteAliases.size());
        assertTrue(remoteAliases.toString(), remoteAliases.containsKey("foo_id"));
        assertTrue(remoteAliases.toString(), remoteAliases.containsKey("bar_id"));
        assertTrue(remoteAliases.toString(), remoteAliases.containsKey("xyz_id"));
        assertEquals(new TermsQueryBuilder("foo", "bar"), remoteAliases.get("foo_id").getQueryBuilder());
        assertEquals(new MatchAllQueryBuilder(), remoteAliases.get("bar_id").getQueryBuilder());
        assertNull(remoteAliases.get("xyz_id").getQueryBuilder());
    }
}

Example 29 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class ShardSearchFailureTests method testToXContentWithClusterAlias.

public void testToXContentWithClusterAlias() throws IOException {
    ShardSearchFailure failure = new ShardSearchFailure(new ParsingException(0, 0, "some message", null), new SearchShardTarget("nodeId", new ShardId(new Index("indexName", "indexUuid"), 123), "cluster1", OriginalIndices.NONE));
    BytesReference xContent = toXContent(failure, XContentType.JSON, randomBoolean());
    assertEquals("{\"shard\":123," + "\"index\":\"cluster1:indexName\"," + "\"node\":\"nodeId\"," + "\"reason\":{" + "\"type\":\"parsing_exception\"," + "\"reason\":\"some message\"," + "\"line\":0," + "\"col\":0" + "}" + "}", xContent.utf8ToString());
}

Example 30 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget in project OpenSearch by opensearch-project.

the class ShardSearchFailureTests method testToXContent.

public void testToXContent() throws IOException {
    ShardSearchFailure failure = new ShardSearchFailure(new ParsingException(0, 0, "some message", null), new SearchShardTarget("nodeId", new ShardId(new Index("indexName", "indexUuid"), 123), null, OriginalIndices.NONE));
    BytesReference xContent = toXContent(failure, XContentType.JSON, randomBoolean());
    assertEquals("{\"shard\":123," + "\"index\":\"indexName\"," + "\"node\":\"nodeId\"," + "\"reason\":{" + "\"type\":\"parsing_exception\"," + "\"reason\":\"some message\"," + "\"line\":0," + "\"col\":0" + "}" + "}", xContent.utf8ToString());
}