Examples with TransportConnectionListener org.opensearch.transport.TransportConnectionListener used on opensource projects

Example 1 with TransportConnectionListener

use of org.opensearch.transport.TransportConnectionListener in project OpenSearch by opensearch-project.

the class TransportSearchActionTests method testCollectSearchShards.

public void testCollectSearchShards() throws Exception {
    int numClusters = randomIntBetween(2, 10);
    DiscoveryNode[] nodes = new DiscoveryNode[numClusters];
    Map<String, OriginalIndices> remoteIndicesByCluster = new HashMap<>();
    Settings.Builder builder = Settings.builder();
    MockTransportService[] mockTransportServices = startTransport(numClusters, nodes, remoteIndicesByCluster, builder);
    Settings settings = builder.build();
    try (MockTransportService service = MockTransportService.createNewService(settings, Version.CURRENT, threadPool, null)) {
        service.start();
        service.acceptIncomingRequests();
        RemoteClusterService remoteClusterService = service.getRemoteClusterService();
        {
            final CountDownLatch latch = new CountDownLatch(1);
            AtomicReference<Map<String, ClusterSearchShardsResponse>> response = new AtomicReference<>();
            AtomicInteger skippedClusters = new AtomicInteger();
            TransportSearchAction.collectSearchShards(IndicesOptions.lenientExpandOpen(), null, null, skippedClusters, remoteIndicesByCluster, remoteClusterService, threadPool, new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch));
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertEquals(0, skippedClusters.get());
            assertNotNull(response.get());
            Map<String, ClusterSearchShardsResponse> map = response.get();
            assertEquals(numClusters, map.size());
            for (int i = 0; i < numClusters; i++) {
                String clusterAlias = "remote" + i;
                assertTrue(map.containsKey(clusterAlias));
                ClusterSearchShardsResponse shardsResponse = map.get(clusterAlias);
                assertEquals(1, shardsResponse.getNodes().length);
            }
        }
        {
            final CountDownLatch latch = new CountDownLatch(1);
            AtomicReference<Exception> failure = new AtomicReference<>();
            AtomicInteger skippedClusters = new AtomicInteger(0);
            TransportSearchAction.collectSearchShards(IndicesOptions.lenientExpandOpen(), "index_not_found", null, skippedClusters, remoteIndicesByCluster, remoteClusterService, threadPool, new LatchedActionListener<>(ActionListener.wrap(r -> fail("no response expected"), failure::set), latch));
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertEquals(0, skippedClusters.get());
            assertNotNull(failure.get());
            assertThat(failure.get(), instanceOf(RemoteTransportException.class));
            RemoteTransportException remoteTransportException = (RemoteTransportException) failure.get();
            assertEquals(RestStatus.NOT_FOUND, remoteTransportException.status());
        }
        int numDisconnectedClusters = randomIntBetween(1, numClusters);
        Set<DiscoveryNode> disconnectedNodes = new HashSet<>(numDisconnectedClusters);
        Set<Integer> disconnectedNodesIndices = new HashSet<>(numDisconnectedClusters);
        while (disconnectedNodes.size() < numDisconnectedClusters) {
            int i = randomIntBetween(0, numClusters - 1);
            if (disconnectedNodes.add(nodes[i])) {
                assertTrue(disconnectedNodesIndices.add(i));
            }
        }
        CountDownLatch disconnectedLatch = new CountDownLatch(numDisconnectedClusters);
        RemoteClusterServiceTests.addConnectionListener(remoteClusterService, new TransportConnectionListener() {

            @Override
            public void onNodeDisconnected(DiscoveryNode node, Transport.Connection connection) {
                if (disconnectedNodes.remove(node)) {
                    disconnectedLatch.countDown();
                }
            }
        });
        for (DiscoveryNode disconnectedNode : disconnectedNodes) {
            service.addFailToSendNoConnectRule(disconnectedNode.getAddress());
        }
        {
            final CountDownLatch latch = new CountDownLatch(1);
            AtomicInteger skippedClusters = new AtomicInteger(0);
            AtomicReference<Exception> failure = new AtomicReference<>();
            TransportSearchAction.collectSearchShards(IndicesOptions.lenientExpandOpen(), null, null, skippedClusters, remoteIndicesByCluster, remoteClusterService, threadPool, new LatchedActionListener<>(ActionListener.wrap(r -> fail("no response expected"), failure::set), latch));
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertEquals(0, skippedClusters.get());
            assertNotNull(failure.get());
            assertThat(failure.get(), instanceOf(RemoteTransportException.class));
            assertThat(failure.get().getMessage(), containsString("error while communicating with remote cluster ["));
            assertThat(failure.get().getCause(), instanceOf(NodeDisconnectedException.class));
        }
        // setting skip_unavailable to true for all the disconnected clusters will make the request succeed again
        for (int i : disconnectedNodesIndices) {
            RemoteClusterServiceTests.updateSkipUnavailable(remoteClusterService, "remote" + i, true);
        }
        {
            final CountDownLatch latch = new CountDownLatch(1);
            AtomicInteger skippedClusters = new AtomicInteger(0);
            AtomicReference<Map<String, ClusterSearchShardsResponse>> response = new AtomicReference<>();
            TransportSearchAction.collectSearchShards(IndicesOptions.lenientExpandOpen(), null, null, skippedClusters, remoteIndicesByCluster, remoteClusterService, threadPool, new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch));
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertNotNull(response.get());
            Map<String, ClusterSearchShardsResponse> map = response.get();
            assertEquals(numClusters - disconnectedNodesIndices.size(), map.size());
            assertEquals(skippedClusters.get(), disconnectedNodesIndices.size());
            for (int i = 0; i < numClusters; i++) {
                String clusterAlias = "remote" + i;
                if (disconnectedNodesIndices.contains(i)) {
                    assertFalse(map.containsKey(clusterAlias));
                } else {
                    assertNotNull(map.get(clusterAlias));
                }
            }
        }
        // give transport service enough time to realize that the node is down, and to notify the connection listeners
        // so that RemoteClusterConnection is left with no connected nodes, hence it will retry connecting next
        assertTrue(disconnectedLatch.await(5, TimeUnit.SECONDS));
        service.clearAllRules();
        if (randomBoolean()) {
            for (int i : disconnectedNodesIndices) {
                if (randomBoolean()) {
                    RemoteClusterServiceTests.updateSkipUnavailable(remoteClusterService, "remote" + i, true);
                }
            }
        }
        {
            final CountDownLatch latch = new CountDownLatch(1);
            AtomicInteger skippedClusters = new AtomicInteger(0);
            AtomicReference<Map<String, ClusterSearchShardsResponse>> response = new AtomicReference<>();
            TransportSearchAction.collectSearchShards(IndicesOptions.lenientExpandOpen(), null, null, skippedClusters, remoteIndicesByCluster, remoteClusterService, threadPool, new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch));
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertEquals(0, skippedClusters.get());
            assertNotNull(response.get());
            Map<String, ClusterSearchShardsResponse> map = response.get();
            assertEquals(numClusters, map.size());
            for (int i = 0; i < numClusters; i++) {
                String clusterAlias = "remote" + i;
                assertTrue(map.containsKey(clusterAlias));
                assertNotNull(map.get(clusterAlias));
            }
        }
        assertEquals(0, service.getConnectionManager().size());
    } finally {
        for (MockTransportService mockTransportService : mockTransportServices) {
            mockTransportService.close();
        }
    }
}

Example 2 with TransportConnectionListener

use of org.opensearch.transport.TransportConnectionListener in project OpenSearch by opensearch-project.

the class TransportSearchActionTests method testCCSRemoteReduce.

public void testCCSRemoteReduce() throws Exception {
    int numClusters = randomIntBetween(1, 10);
    DiscoveryNode[] nodes = new DiscoveryNode[numClusters];
    Map<String, OriginalIndices> remoteIndicesByCluster = new HashMap<>();
    Settings.Builder builder = Settings.builder();
    MockTransportService[] mockTransportServices = startTransport(numClusters, nodes, remoteIndicesByCluster, builder);
    Settings settings = builder.build();
    boolean local = randomBoolean();
    OriginalIndices localIndices = local ? new OriginalIndices(new String[] { "index" }, SearchRequest.DEFAULT_INDICES_OPTIONS) : null;
    int totalClusters = numClusters + (local ? 1 : 0);
    TransportSearchAction.SearchTimeProvider timeProvider = new TransportSearchAction.SearchTimeProvider(0, 0, () -> 0);
    try (MockTransportService service = MockTransportService.createNewService(settings, Version.CURRENT, threadPool, null)) {
        service.start();
        service.acceptIncomingRequests();
        RemoteClusterService remoteClusterService = service.getRemoteClusterService();
        {
            SearchRequest searchRequest = new SearchRequest();
            final CountDownLatch latch = new CountDownLatch(1);
            SetOnce<Tuple<SearchRequest, ActionListener<SearchResponse>>> setOnce = new SetOnce<>();
            AtomicReference<SearchResponse> response = new AtomicReference<>();
            LatchedActionListener<SearchResponse> listener = new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch);
            TransportSearchAction.ccsRemoteReduce(searchRequest, localIndices, remoteIndicesByCluster, timeProvider, emptyReduceContextBuilder(), remoteClusterService, threadPool, listener, (r, l) -> setOnce.set(Tuple.tuple(r, l)));
            if (localIndices == null) {
                assertNull(setOnce.get());
            } else {
                Tuple<SearchRequest, ActionListener<SearchResponse>> tuple = setOnce.get();
                assertEquals("", tuple.v1().getLocalClusterAlias());
                assertThat(tuple.v2(), instanceOf(TransportSearchAction.CCSActionListener.class));
                tuple.v2().onResponse(emptySearchResponse());
            }
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            SearchResponse searchResponse = response.get();
            assertEquals(0, searchResponse.getClusters().getSkipped());
            assertEquals(totalClusters, searchResponse.getClusters().getTotal());
            assertEquals(totalClusters, searchResponse.getClusters().getSuccessful());
            assertEquals(totalClusters == 1 ? 1 : totalClusters + 1, searchResponse.getNumReducePhases());
        }
        {
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.preference("index_not_found");
            final CountDownLatch latch = new CountDownLatch(1);
            SetOnce<Tuple<SearchRequest, ActionListener<SearchResponse>>> setOnce = new SetOnce<>();
            AtomicReference<Exception> failure = new AtomicReference<>();
            LatchedActionListener<SearchResponse> listener = new LatchedActionListener<>(ActionListener.wrap(r -> fail("no response expected"), failure::set), latch);
            TransportSearchAction.ccsRemoteReduce(searchRequest, localIndices, remoteIndicesByCluster, timeProvider, emptyReduceContextBuilder(), remoteClusterService, threadPool, listener, (r, l) -> setOnce.set(Tuple.tuple(r, l)));
            if (localIndices == null) {
                assertNull(setOnce.get());
            } else {
                Tuple<SearchRequest, ActionListener<SearchResponse>> tuple = setOnce.get();
                assertEquals("", tuple.v1().getLocalClusterAlias());
                assertThat(tuple.v2(), instanceOf(TransportSearchAction.CCSActionListener.class));
                tuple.v2().onResponse(emptySearchResponse());
            }
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertNotNull(failure.get());
            assertThat(failure.get(), instanceOf(RemoteTransportException.class));
            RemoteTransportException remoteTransportException = (RemoteTransportException) failure.get();
            assertEquals(RestStatus.NOT_FOUND, remoteTransportException.status());
        }
        int numDisconnectedClusters = randomIntBetween(1, numClusters);
        Set<DiscoveryNode> disconnectedNodes = new HashSet<>(numDisconnectedClusters);
        Set<Integer> disconnectedNodesIndices = new HashSet<>(numDisconnectedClusters);
        while (disconnectedNodes.size() < numDisconnectedClusters) {
            int i = randomIntBetween(0, numClusters - 1);
            if (disconnectedNodes.add(nodes[i])) {
                assertTrue(disconnectedNodesIndices.add(i));
            }
        }
        CountDownLatch disconnectedLatch = new CountDownLatch(numDisconnectedClusters);
        RemoteClusterServiceTests.addConnectionListener(remoteClusterService, new TransportConnectionListener() {

            @Override
            public void onNodeDisconnected(DiscoveryNode node, Transport.Connection connection) {
                if (disconnectedNodes.remove(node)) {
                    disconnectedLatch.countDown();
                }
            }
        });
        for (DiscoveryNode disconnectedNode : disconnectedNodes) {
            service.addFailToSendNoConnectRule(disconnectedNode.getAddress());
        }
        {
            SearchRequest searchRequest = new SearchRequest();
            final CountDownLatch latch = new CountDownLatch(1);
            SetOnce<Tuple<SearchRequest, ActionListener<SearchResponse>>> setOnce = new SetOnce<>();
            AtomicReference<Exception> failure = new AtomicReference<>();
            LatchedActionListener<SearchResponse> listener = new LatchedActionListener<>(ActionListener.wrap(r -> fail("no response expected"), failure::set), latch);
            TransportSearchAction.ccsRemoteReduce(searchRequest, localIndices, remoteIndicesByCluster, timeProvider, emptyReduceContextBuilder(), remoteClusterService, threadPool, listener, (r, l) -> setOnce.set(Tuple.tuple(r, l)));
            if (localIndices == null) {
                assertNull(setOnce.get());
            } else {
                Tuple<SearchRequest, ActionListener<SearchResponse>> tuple = setOnce.get();
                assertEquals("", tuple.v1().getLocalClusterAlias());
                assertThat(tuple.v2(), instanceOf(TransportSearchAction.CCSActionListener.class));
                tuple.v2().onResponse(emptySearchResponse());
            }
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            assertNotNull(failure.get());
            assertThat(failure.get(), instanceOf(RemoteTransportException.class));
            assertThat(failure.get().getMessage(), containsString("error while communicating with remote cluster ["));
            assertThat(failure.get().getCause(), instanceOf(NodeDisconnectedException.class));
        }
        // setting skip_unavailable to true for all the disconnected clusters will make the request succeed again
        for (int i : disconnectedNodesIndices) {
            RemoteClusterServiceTests.updateSkipUnavailable(remoteClusterService, "remote" + i, true);
        }
        {
            SearchRequest searchRequest = new SearchRequest();
            final CountDownLatch latch = new CountDownLatch(1);
            SetOnce<Tuple<SearchRequest, ActionListener<SearchResponse>>> setOnce = new SetOnce<>();
            AtomicReference<SearchResponse> response = new AtomicReference<>();
            LatchedActionListener<SearchResponse> listener = new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch);
            TransportSearchAction.ccsRemoteReduce(searchRequest, localIndices, remoteIndicesByCluster, timeProvider, emptyReduceContextBuilder(), remoteClusterService, threadPool, listener, (r, l) -> setOnce.set(Tuple.tuple(r, l)));
            if (localIndices == null) {
                assertNull(setOnce.get());
            } else {
                Tuple<SearchRequest, ActionListener<SearchResponse>> tuple = setOnce.get();
                assertEquals("", tuple.v1().getLocalClusterAlias());
                assertThat(tuple.v2(), instanceOf(TransportSearchAction.CCSActionListener.class));
                tuple.v2().onResponse(emptySearchResponse());
            }
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            SearchResponse searchResponse = response.get();
            assertEquals(disconnectedNodesIndices.size(), searchResponse.getClusters().getSkipped());
            assertEquals(totalClusters, searchResponse.getClusters().getTotal());
            int successful = totalClusters - disconnectedNodesIndices.size();
            assertEquals(successful, searchResponse.getClusters().getSuccessful());
            assertEquals(successful == 0 ? 0 : successful + 1, searchResponse.getNumReducePhases());
        }
        // give transport service enough time to realize that the node is down, and to notify the connection listeners
        // so that RemoteClusterConnection is left with no connected nodes, hence it will retry connecting next
        assertTrue(disconnectedLatch.await(5, TimeUnit.SECONDS));
        service.clearAllRules();
        if (randomBoolean()) {
            for (int i : disconnectedNodesIndices) {
                if (randomBoolean()) {
                    RemoteClusterServiceTests.updateSkipUnavailable(remoteClusterService, "remote" + i, true);
                }
            }
        }
        {
            SearchRequest searchRequest = new SearchRequest();
            final CountDownLatch latch = new CountDownLatch(1);
            SetOnce<Tuple<SearchRequest, ActionListener<SearchResponse>>> setOnce = new SetOnce<>();
            AtomicReference<SearchResponse> response = new AtomicReference<>();
            LatchedActionListener<SearchResponse> listener = new LatchedActionListener<>(ActionListener.wrap(response::set, e -> fail("no failures expected")), latch);
            TransportSearchAction.ccsRemoteReduce(searchRequest, localIndices, remoteIndicesByCluster, timeProvider, emptyReduceContextBuilder(), remoteClusterService, threadPool, listener, (r, l) -> setOnce.set(Tuple.tuple(r, l)));
            if (localIndices == null) {
                assertNull(setOnce.get());
            } else {
                Tuple<SearchRequest, ActionListener<SearchResponse>> tuple = setOnce.get();
                assertEquals("", tuple.v1().getLocalClusterAlias());
                assertThat(tuple.v2(), instanceOf(TransportSearchAction.CCSActionListener.class));
                tuple.v2().onResponse(emptySearchResponse());
            }
            awaitLatch(latch, 5, TimeUnit.SECONDS);
            SearchResponse searchResponse = response.get();
            assertEquals(0, searchResponse.getClusters().getSkipped());
            assertEquals(totalClusters, searchResponse.getClusters().getTotal());
            assertEquals(totalClusters, searchResponse.getClusters().getSuccessful());
            assertEquals(totalClusters == 1 ? 1 : totalClusters + 1, searchResponse.getNumReducePhases());
        }
        assertEquals(0, service.getConnectionManager().size());
    } finally {
        for (MockTransportService mockTransportService : mockTransportServices) {
            mockTransportService.close();
        }
    }
}