Examples with SeedHostsProvider org.opensearch.discovery.SeedHostsProvider used on opensource projects

Example 1 with SeedHostsProvider

use of org.opensearch.discovery.SeedHostsProvider in project OpenSearch by opensearch-project.

the class EC2RetriesTests method testEC2DiscoveryRetriesOnRateLimiting.

public void testEC2DiscoveryRetriesOnRateLimiting() throws IOException {
    final String accessKey = "ec2_access";
    final List<String> hosts = Collections.singletonList("127.0.0.1:9300");
    final Map<String, Integer> failedRequests = new ConcurrentHashMap<>();
    // retry the same request 5 times at most
    final int maxRetries = randomIntBetween(1, 5);
    httpServer.createContext("/", exchange -> {
        if (exchange.getRequestMethod().equals(HttpMethodName.POST.name())) {
            final String request = Streams.readFully(exchange.getRequestBody()).utf8ToString();
            final String userAgent = exchange.getRequestHeaders().getFirst("User-Agent");
            if (userAgent != null && userAgent.startsWith("aws-sdk-java")) {
                final String auth = exchange.getRequestHeaders().getFirst("Authorization");
                if (auth == null || auth.contains(accessKey) == false) {
                    throw new IllegalArgumentException("wrong access key: " + auth);
                }
                if (failedRequests.compute(exchange.getRequestHeaders().getFirst("Amz-sdk-invocation-id"), (requestId, count) -> (count == null ? 0 : count) + 1) < maxRetries) {
                    exchange.sendResponseHeaders(HttpStatus.SC_SERVICE_UNAVAILABLE, -1);
                    return;
                }
                // Simulate an EC2 DescribeInstancesResponse
                byte[] responseBody = null;
                for (NameValuePair parse : URLEncodedUtils.parse(request, UTF_8)) {
                    if ("Action".equals(parse.getName())) {
                        responseBody = generateDescribeInstancesResponse(hosts.stream().map(address -> new Instance().withPublicIpAddress(address)).collect(Collectors.toList()));
                        break;
                    }
                }
                responseBody = responseBody == null ? new byte[0] : responseBody;
                exchange.getResponseHeaders().set("Content-Type", "text/xml; charset=UTF-8");
                exchange.sendResponseHeaders(HttpStatus.SC_OK, responseBody.length);
                exchange.getResponseBody().write(responseBody);
                return;
            }
        }
        fail("did not send response");
    });
    try (Ec2DiscoveryPlugin plugin = new Ec2DiscoveryPlugin(buildSettings(accessKey))) {
        final SeedHostsProvider seedHostsProvider = plugin.getSeedHostProviders(transportService, networkService).get("ec2").get();
        final SeedHostsResolver resolver = new SeedHostsResolver("test", Settings.EMPTY, transportService, seedHostsProvider);
        resolver.start();
        final List<TransportAddress> addressList = seedHostsProvider.getSeedAddresses(null);
        assertThat(addressList, Matchers.hasSize(1));
        assertThat(addressList.get(0).toString(), is(hosts.get(0)));
        assertThat(failedRequests, aMapWithSize(1));
        assertThat(failedRequests.values().iterator().next(), is(maxRetries));
    }
}