Java Code Examples for org.elasticsearch.client.RestClientBuilder#setRequestConfigCallback()

/**
 * 异步httpclient连接延时配置
 */
private void setRequestConfig(RestClientBuilder builder, RestClientPoolProperties poolProperties){
    builder.setRequestConfigCallback(requestConfigBuilder -> {
        requestConfigBuilder.setConnectTimeout(poolProperties.getConnectTimeOut())
                .setSocketTimeout(poolProperties.getSocketTimeOut())
                .setConnectionRequestTimeout(poolProperties.getConnectionRequestTimeOut());
        return requestConfigBuilder;
    });
}
 

@Bean
public RestHighLevelClient elasticsearchRestConnection(ESConfiguration esConfiguration) {
    String auth = esConfiguration.getServiceElasticsearchUsername() + ":" + esConfiguration.getServiceElasticsearchPassword();
    String authB64;
    try {
        authB64 = Base64.getEncoder().encodeToString(auth.getBytes("utf-8"));
    } catch (UnsupportedEncodingException e) {
        throw new RuntimeException("Impossible encoding for user " + esConfiguration.getServiceElasticsearchUsername() + " and password " + esConfiguration.getServiceElasticsearchPassword() + " msg " + e);
    }
    RestClientBuilder builder = RestClient.builder(
            new HttpHost(esConfiguration.getHost(), Integer.valueOf(esConfiguration.getPort()), "http"));
    Header[] defaultHeaders = new Header[]{
            new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"),
            new BasicHeader("cluster.name", esConfiguration.getClusterName()),
            new BasicHeader(HttpHeaders.AUTHORIZATION, "Basic " + authB64)
    };
    builder.setDefaultHeaders(defaultHeaders);
    builder.setMaxRetryTimeoutMillis(esConfiguration.getSocketTimeout() * 1000);
    builder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
        @Override
        public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
            return requestConfigBuilder
                    .setConnectionRequestTimeout(connectionRequestTimeout)
                    .setConnectTimeout(esConfiguration.getConnectionTimeout() * 1000)
                    .setSocketTimeout(esConfiguration.getSocketTimeout() * 1000);
        }
    });
    return new RestHighLevelClient(builder);
}
 

private RestClient getRestClient() {
    if (restClient == null) {
    	String esHost = config.getElasticSearch().getDevIngestHost();
		int esPort = config.getElasticSearch().getDevIngestPort();
        RestClientBuilder builder = RestClient.builder(new HttpHost(esHost, esPort));
        builder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
            @Override
            public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
                return requestConfigBuilder.setConnectionRequestTimeout(0);
            }
        });
        restClient = builder.build();
    }
    return restClient;
}
 

private RestClient getRestClient() {
    if (restClient == null) {
    	String esHost = config.getElasticSearch().getDevIngestHost();
		int esPort = config.getElasticSearch().getDevIngestPort();
        RestClientBuilder builder = RestClient.builder(new HttpHost(esHost, esPort));
        builder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
            @Override
            public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
                return requestConfigBuilder.setConnectionRequestTimeout(0);
            }
        });
        restClient = builder.build();
    }
    return restClient;
}
 

private RestClient getRestClient() {
	if (restClient == null) {
		String esHost = config.getElasticSearch().getDevIngestHost();
		int esPort = config.getElasticSearch().getDevIngestPort();
		RestClientBuilder builder = RestClient.builder(new HttpHost(esHost, esPort));
		RequestConfigCallback requestConfigCallback = requestConfigBuilder -> requestConfigBuilder
				.setConnectionRequestTimeout(0);
		builder.setRequestConfigCallback(requestConfigCallback);
		restClient = builder.build();
	}
	return restClient;
}
 

private RestClient getRestClient() {
    if (restClient == null) {

        RestClientBuilder builder = RestClient.builder(new HttpHost(esHost, esPort));
        builder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
            @Override
            public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
                return requestConfigBuilder.setConnectionRequestTimeout(0);
            }
        });
        restClient = builder.build();
    }
    return restClient;

}
 

/**
 * get restHistLevelClient
 *
 * @param builder                 RestClientBuilder
 * @param elasticsearchProperties elasticsearch default properties
 * @return {@link org.elasticsearch.client.RestHighLevelClient}
 * @author fxbin
 */
private static RestHighLevelClient getRestHighLevelClient(RestClientBuilder builder, ElasticsearchProperties elasticsearchProperties) {

    // Callback used the default {@link RequestConfig} being set to the {@link CloseableHttpClient}
    builder.setRequestConfigCallback(requestConfigBuilder -> {
        requestConfigBuilder.setConnectTimeout(elasticsearchProperties.getConnectTimeout());
        requestConfigBuilder.setSocketTimeout(elasticsearchProperties.getSocketTimeout());
        requestConfigBuilder.setConnectionRequestTimeout(elasticsearchProperties.getConnectionRequestTimeout());
        return requestConfigBuilder;
    });

    // Callback used to customize the {@link CloseableHttpClient} instance used by a {@link RestClient} instance.
    builder.setHttpClientConfigCallback(httpClientBuilder -> {
        httpClientBuilder.setMaxConnTotal(elasticsearchProperties.getMaxConnectTotal());
        httpClientBuilder.setMaxConnPerRoute(elasticsearchProperties.getMaxConnectPerRoute());
        return httpClientBuilder;
    });

    // Callback used the basic credential auth
    ElasticsearchProperties.Account account = elasticsearchProperties.getAccount();
    if (!StringUtils.isEmpty(account.getUsername()) && !StringUtils.isEmpty(account.getUsername())) {
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();

        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(account.getUsername(), account.getPassword()));
    }
    return new RestHighLevelClient(builder);
}
 

public void initialize() {
    RestClientBuilder builder = RestClient.builder(hosts);
    builder.setRequestConfigCallback(config -> config.setSocketTimeout((int) timeout.toMillis())
                                                     .setConnectionRequestTimeout((int) timeout.toMillis())); // timeout of requesting connection from connection pool
    builder.setHttpClientConfigCallback(config -> config.setMaxConnTotal(100)
                                                        .setMaxConnPerRoute(100)
                                                        .setKeepAliveStrategy((response, context) -> Duration.ofSeconds(30).toMillis()));
    client = new RestHighLevelClient(builder);
}
 

/**
 * Creates an Elasticsearch client from settings provided via the global config.
 *
 * @return new client
 */
public static ElasticsearchClient create(Map<String, Object> globalConfig) {
  ElasticsearchClientConfig esClientConfig = new ElasticsearchClientConfig(
      getEsSettings(globalConfig));
  HttpHost[] httpHosts = getHttpHosts(globalConfig, esClientConfig.getConnectionScheme());
  RestClientBuilder builder = RestClient.builder(httpHosts);

  builder.setRequestConfigCallback(reqConfigBuilder -> {
    // Modifies request config builder with connection and socket timeouts.
    // https://www.elastic.co/guide/en/elasticsearch/client/java-rest/5.6/_timeouts.html
    reqConfigBuilder.setConnectTimeout(esClientConfig.getConnectTimeoutMillis());
    reqConfigBuilder.setSocketTimeout(esClientConfig.getSocketTimeoutMillis());
    return reqConfigBuilder;
  });
  builder.setMaxRetryTimeoutMillis(esClientConfig.getMaxRetryTimeoutMillis());

  builder.setHttpClientConfigCallback(clientBuilder -> {
    clientBuilder.setDefaultIOReactorConfig(getIOReactorConfig(esClientConfig));
    clientBuilder.setDefaultCredentialsProvider(getCredentialsProvider(esClientConfig));
    clientBuilder.setSSLContext(getSSLContext(esClientConfig));
    return clientBuilder;
  });

  RestClient lowLevelClient = builder.build();
  RestHighLevelClient client = new RestHighLevelClient(lowLevelClient);
  return new ElasticsearchClient(lowLevelClient, client);
}
 

public void activate(Dictionary<String, Object> configuration) {
    this.configuration = configuration;
    String addressesString = (configuration.get("addresses") != null) ? configuration.get("addresses").toString() : "http://localhost:9200";
    String username = (configuration.get("username") != null) ? configuration.get("username").toString() : null;
    String password = (configuration.get("password") != null) ? configuration.get("password").toString() : null;

    Set<String> addresses = new HashSet<>(Arrays.asList(addressesString.split(",")));

    HttpHost[] hosts = new HttpHost[addresses.size()];
    int i = 0;
    for (String address : addresses) {
        try {
            URL url = new URL(address);
            hosts[i] = new HttpHost(url.getHost(), url.getPort(), url.getProtocol());
            i++;
        } catch (Exception e) {
            LOGGER.warn("Bad elasticsearch address {}", address, e);
        }
    }
    RestClientBuilder restClientBuilder = RestClient.builder(hosts);

    restClientBuilder.setRequestConfigCallback(requestConfigBuilder -> requestConfigBuilder.setConnectTimeout(1000)
            .setSocketTimeout(10000));

    if (username != null) {
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(username, password));
        restClientBuilder.setHttpClientConfigCallback(
                new RestClientBuilder.HttpClientConfigCallback() {
                    @Override
                    public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpAsyncClientBuilder) {
                        return httpAsyncClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                    }
                }
        );
    }

    restClient = new RestHighLevelClient(restClientBuilder);
}
 

@VisibleForTesting
RestClient createClient() throws IOException {
  HttpHost[] hosts = new HttpHost[getAddresses().size()];
  int i = 0;
  for (String address : getAddresses()) {
    URL url = new URL(address);
    hosts[i] = new HttpHost(url.getHost(), url.getPort(), url.getProtocol());
    i++;
  }
  RestClientBuilder restClientBuilder = RestClient.builder(hosts);
  if (getUsername() != null) {
    final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
    credentialsProvider.setCredentials(
        AuthScope.ANY, new UsernamePasswordCredentials(getUsername(), getPassword()));
    restClientBuilder.setHttpClientConfigCallback(
        httpAsyncClientBuilder ->
            httpAsyncClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
  }
  if (getKeystorePath() != null && !getKeystorePath().isEmpty()) {
    try {
      KeyStore keyStore = KeyStore.getInstance("jks");
      try (InputStream is = new FileInputStream(new File(getKeystorePath()))) {
        String keystorePassword = getKeystorePassword();
        keyStore.load(is, (keystorePassword == null) ? null : keystorePassword.toCharArray());
      }
      final TrustStrategy trustStrategy =
          isTrustSelfSignedCerts() ? new TrustSelfSignedStrategy() : null;
      final SSLContext sslContext =
          SSLContexts.custom().loadTrustMaterial(keyStore, trustStrategy).build();
      final SSLIOSessionStrategy sessionStrategy = new SSLIOSessionStrategy(sslContext);
      restClientBuilder.setHttpClientConfigCallback(
          httpClientBuilder ->
              httpClientBuilder.setSSLContext(sslContext).setSSLStrategy(sessionStrategy));
    } catch (Exception e) {
      throw new IOException("Can't load the client certificate from the keystore", e);
    }
  }
  restClientBuilder.setRequestConfigCallback(
      new RestClientBuilder.RequestConfigCallback() {
        @Override
        public RequestConfig.Builder customizeRequestConfig(
            RequestConfig.Builder requestConfigBuilder) {
          if (getConnectTimeout() != null) {
            requestConfigBuilder.setConnectTimeout(getConnectTimeout());
          }
          if (getSocketTimeout() != null) {
            requestConfigBuilder.setSocketTimeout(getSocketTimeout());
          }
          return requestConfigBuilder;
        }
      });
  return restClientBuilder.build();
}
 

private RestClient createRestClient(Map<String, Serializable> params, String user, String password) throws IOException {
    final String hostName = getValue(HOSTNAME, params);
    final String[] hosts = hostName.split(",");
    final Integer defaultPort = getValue(HOSTPORT, params);
    final Boolean sslRejectUnauthorized = getValue(SSL_REJECT_UNAUTHORIZED, params);
    final String adminUser = getValue(USER, params);
    final String type = user == null || adminUser == null || user.equals(adminUser) ? "ADMIN" : "PROXY_USER";

    final Pattern pattern = Pattern.compile("(?<scheme>https?)?(://)?(?<host>[^:]+):?(?<port>\\d+)?");
    final HttpHost[] httpHosts = new HttpHost[hosts.length];
    final AuthScope[] auths = new AuthScope[hosts.length];
    for (int index=0; index < hosts.length; index++) {
        final Matcher matcher = pattern.matcher(hosts[index].trim());
        if (matcher.find()) {
            final String scheme = matcher.group("scheme") != null ? matcher.group("scheme") : "http";
            final String host = matcher.group("host");
            final Integer port = matcher.group("port") != null ? Integer.valueOf(matcher.group("port")) : defaultPort;
            httpHosts[index] = new HttpHost(host, port, scheme);
            auths[index] = new AuthScope(host, port);
        } else {
            throw new IOException("Unable to parse host");
        }
    }

    final RestClientBuilder builder = createClientBuilder(httpHosts);

    if (user != null) {
        builder.setRequestConfigCallback((b) -> {
            LOGGER.finest(String.format("Calling %s setRequestConfigCallback", type));
            return b.setAuthenticationEnabled(true);
        });
    }

    builder.setHttpClientConfigCallback((httpClientBuilder) -> {
        LOGGER.finest(String.format("Calling %s customizeHttpClient", type));

        httpClientBuilder.setThreadFactory((run) -> {
            final Thread thread = new Thread(run);
            thread.setDaemon(true);
            thread.setName(String.format("esrest-asynchttp-%s-%d", type, httpThreads.getAndIncrement()));
            return thread;
        });

        httpClientBuilder.useSystemProperties();

        if (!sslRejectUnauthorized) {
            httpClientBuilder.setSSLHostnameVerifier((host,session) -> true);
            try {
                httpClientBuilder.setSSLContext(SSLContextBuilder.create().loadTrustMaterial((chain,authType) ->true).build());
            } catch (KeyManagementException | NoSuchAlgorithmException | KeyStoreException e) {
                throw new UncheckedIOException(new IOException("Unable to create SSLContext", e));
            }
        }

        if (user != null) {
            final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
            final Credentials credentials = new org.apache.http.auth.UsernamePasswordCredentials(user, password);
            for (AuthScope scope : auths) {
                credentialsProvider.setCredentials(scope, credentials);
            }

            httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
        }

        return httpClientBuilder;
    });

    LOGGER.fine(String.format("Building a %s RestClient for %s @ %s:%d", type, user, hostName, defaultPort));
    return builder.build();
}
 

public void open(Dictionary<String, Object> config) {
    this.config = config;

    String addressesString = getValue(config, ADDRESSES_PROPERTY, ADDRESSES_DEFAULT);
    String username = getValue(config, USERNAME_PROPERTY, USERNAME_DEFAULT);
    String password = getValue(config, PASSWORD_PROPERTY, PASSWORD_DEFAULT);

    Set<String> addresses = new HashSet<String>(Arrays.asList(addressesString.split(",")));

    HttpHost[] hosts = new HttpHost[addresses.size()];
    int i = 0;
    for (String address : addresses) {
        try {
            URL url = new URL(address);
            hosts[i] = new HttpHost(url.getHost(), url.getPort(), url.getProtocol());
            i++;
        } catch (Exception e) {
            LOGGER.warn("Bad elasticsearch address {}", address, e);
        }
    }
    RestClientBuilder restClientBuilder = RestClient.builder(hosts);

    restClientBuilder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
        @Override
        public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
            return requestConfigBuilder.setConnectTimeout(1000)
                    .setSocketTimeout(10000);
        }
    });

    if (username != null) {
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(username, password));
        restClientBuilder.setHttpClientConfigCallback(
                new RestClientBuilder.HttpClientConfigCallback() {
                    @Override
                    public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpAsyncClientBuilder) {
                        return httpAsyncClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                    }
                }
        );
    }

    client = restClientBuilder.build();

    TimeZone tz = TimeZone.getTimeZone( "UTC" );
    tsFormat.setTimeZone(tz);
    indexDateFormat.setTimeZone(tz);
}
 

private static RestClient buildRestClient(List<InetSocketTransportAddress> hosts, int threadCount, boolean sslEnabled,
    String keyStoreType, String keyStoreFilePassword, String identityFilepath, String trustStoreType,
    String trustStoreFilePassword, String cacertsFilepath) throws Exception {


  HttpHost[] httpHosts = new HttpHost[hosts.size()];
  String scheme = sslEnabled?"https":"http";
  for (int h = 0; h < httpHosts.length; h++) {
    InetSocketTransportAddress host = hosts.get(h);
    httpHosts[h] = new HttpHost(host.getAddress(), host.getPort(), scheme);
  }

  RestClientBuilder builder = RestClient.builder(httpHosts);

  if (sslEnabled) {
    log.info("ssl configuration: trustStoreType = {}, cacertsFilePath = {}", trustStoreType, cacertsFilepath);
    KeyStore truststore = KeyStore.getInstance(trustStoreType);
    FileInputStream trustInputStream = new FileInputStream(cacertsFilepath);
    try {
      truststore.load(trustInputStream, trustStoreFilePassword.toCharArray());
    }
    finally {
      trustInputStream.close();
    }
    SSLContextBuilder sslBuilder = SSLContexts.custom().loadTrustMaterial(truststore, null);

    log.info("ssl key configuration: keyStoreType = {}, keyFilePath = {}", keyStoreType, identityFilepath);

    KeyStore keystore = KeyStore.getInstance(keyStoreType);
    FileInputStream keyInputStream = new FileInputStream(identityFilepath);
    try {
      keystore.load(keyInputStream, keyStoreFilePassword.toCharArray());
    }
    finally {
      keyInputStream.close();
    }
    sslBuilder.loadKeyMaterial(keystore, keyStoreFilePassword.toCharArray());

    final SSLContext sslContext = sslBuilder.build();
    builder = builder.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
        // Set ssl context
        .setSSLContext(sslContext).setSSLHostnameVerifier(new NoopHostnameVerifier())
        // Configure number of threads for clients
        .setDefaultIOReactorConfig(IOReactorConfig.custom().setIoThreadCount(threadCount).build()));
  } else {
    builder = builder.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
        // Configure number of threads for clients
        .setDefaultIOReactorConfig(IOReactorConfig.custom().setIoThreadCount(threadCount).build()));
  }

  // Configure timeouts
  builder.setRequestConfigCallback(requestConfigBuilder -> requestConfigBuilder
      .setConnectionRequestTimeout(0)); // Important, otherwise the client has spurious timeouts

  return builder.build();
}