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

Example 16 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget 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 17 with SearchShardTarget

use of org.opensearch.search.SearchShardTarget 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 18 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 19 with SearchShardTarget

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

the class SearchAsyncActionTests method testFanOutAndCollect.

public void testFanOutAndCollect() throws InterruptedException {
    SearchRequest request = new SearchRequest();
    request.allowPartialSearchResults(true);
    request.setMaxConcurrentShardRequests(randomIntBetween(1, 100));
    AtomicReference<TestSearchResponse> response = new AtomicReference<>();
    ActionListener<SearchResponse> responseListener = ActionListener.wrap(searchResponse -> response.set((TestSearchResponse) searchResponse), (e) -> {
        throw new AssertionError("unexpected", e);
    });
    DiscoveryNode primaryNode = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode replicaNode = new DiscoveryNode("node_2", buildNewFakeTransportAddress(), Version.CURRENT);
    Map<DiscoveryNode, Set<ShardSearchContextId>> nodeToContextMap = newConcurrentMap();
    AtomicInteger contextIdGenerator = new AtomicInteger(0);
    int numShards = randomIntBetween(1, 10);
    GroupShardsIterator<SearchShardIterator> shardsIter = getShardsIter("idx", new OriginalIndices(new String[] { "idx" }, SearchRequest.DEFAULT_INDICES_OPTIONS), numShards, randomBoolean(), primaryNode, replicaNode);
    AtomicInteger numFreedContext = new AtomicInteger();
    SearchTransportService transportService = new SearchTransportService(null, null) {

        @Override
        public void sendFreeContext(Transport.Connection connection, ShardSearchContextId contextId, OriginalIndices originalIndices) {
            numFreedContext.incrementAndGet();
            assertTrue(nodeToContextMap.containsKey(connection.getNode()));
            assertTrue(nodeToContextMap.get(connection.getNode()).remove(contextId));
        }
    };
    Map<String, Transport.Connection> lookup = new HashMap<>();
    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));
    ExecutorService executor = Executors.newFixedThreadPool(randomIntBetween(1, Runtime.getRuntime().availableProcessors()));
    final CountDownLatch latch = new CountDownLatch(numShards);
    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(), executor, request, responseListener, shardsIter, new TransportSearchAction.SearchTimeProvider(0, 0, () -> 0), ClusterState.EMPTY_STATE, null, new ArraySearchPhaseResults<>(shardsIter.size()), request.getMaxConcurrentShardRequests(), SearchResponse.Clusters.EMPTY) {

        TestSearchResponse response = new TestSearchResponse();

        @Override
        protected void executePhaseOnShard(SearchShardIterator shardIt, SearchShardTarget shard, SearchActionListener<TestSearchPhaseResult> listener) {
            assertTrue("shard: " + shard.getShardId() + " has been queried twice", response.queried.add(shard.getShardId()));
            Transport.Connection connection = getConnection(null, shard.getNodeId());
            TestSearchPhaseResult testSearchPhaseResult = new TestSearchPhaseResult(new ShardSearchContextId(UUIDs.randomBase64UUID(), contextIdGenerator.incrementAndGet()), connection.getNode());
            Set<ShardSearchContextId> ids = nodeToContextMap.computeIfAbsent(connection.getNode(), (n) -> newConcurrentSet());
            ids.add(testSearchPhaseResult.getContextId());
            if (randomBoolean()) {
                listener.onResponse(testSearchPhaseResult);
            } else {
                new Thread(() -> listener.onResponse(testSearchPhaseResult)).start();
            }
        }

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

                @Override
                public void run() {
                    for (int i = 0; i < results.getNumShards(); i++) {
                        TestSearchPhaseResult result = results.getAtomicArray().get(i);
                        assertEquals(result.node.getId(), result.getSearchShardTarget().getNodeId());
                        sendReleaseSearchContext(result.getContextId(), new MockConnection(result.node), OriginalIndices.NONE);
                    }
                    responseListener.onResponse(response);
                }
            };
        }

        @Override
        protected void executeNext(Runnable runnable, Thread originalThread) {
            super.executeNext(runnable, originalThread);
            latch.countDown();
        }
    };
    asyncAction.start();
    latch.await();
    assertNotNull(response.get());
    assertFalse(nodeToContextMap.isEmpty());
    assertTrue(nodeToContextMap.toString(), nodeToContextMap.containsKey(primaryNode) || nodeToContextMap.containsKey(replicaNode));
    assertEquals(shardsIter.size(), numFreedContext.get());
    if (nodeToContextMap.containsKey(primaryNode)) {
        assertTrue(nodeToContextMap.get(primaryNode).toString(), nodeToContextMap.get(primaryNode).isEmpty());
    } else {
        assertTrue(nodeToContextMap.get(replicaNode).toString(), nodeToContextMap.get(replicaNode).isEmpty());
    }
    executor.shutdown();
}

Example 20 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());
}