diff --git a/gt-elasticsearch/src/main/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFactory.java b/gt-elasticsearch/src/main/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFactory.java index 5da3cce..1c2c641 100644 --- a/gt-elasticsearch/src/main/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFactory.java +++ b/gt-elasticsearch/src/main/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFactory.java @@ -28,6 +28,7 @@ import java.security.NoSuchAlgorithmException; import java.util.Collections; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Logger; import java.util.regex.Matcher; @@ -136,6 +137,15 @@ public class ElasticDataStoreFactory implements DataStoreFactorySpi { GRID_THRESHOLD }; + /** + * The ES RestClient instances created by this factory. We don't want to create + * more than necessary, a single client can work for all stores at the same + * user:host:port. + * + * TODO: Close the clients when GeoServer is shutting down + */ + private Map clients = new ConcurrentHashMap<>(); + @Override public String getDisplayName() { return DISPLAY_NAME; @@ -185,8 +195,8 @@ public DataStore createDataStore(Map params) throws IOExce final String proxyUser = getValue(PROXY_USER, params); final String proxyPasswd = getValue(PROXY_PASSWD, params); - final RestClient client = createRestClient(params, user, passwd); - final RestClient proxyClient = proxyUser != null ? createRestClient(params, proxyUser, proxyPasswd) : null; + final RestClient client = getRestClient(params, user, passwd); + final RestClient proxyClient = proxyUser != null ? getRestClient(params, proxyUser, proxyPasswd) : null; return createDataStore(client, proxyClient, params); } @@ -211,17 +221,47 @@ public DataStore createDataStore(RestClient client, RestClient proxyClient, Map< return dataStore; } - public RestClient createRestClient(Map params) throws IOException { - return createRestClient(params, null, null); + RestClient getRestClient(Map params) throws IOException { + return getRestClient(params, null, null); } - private RestClient createRestClient(Map params, String user, String password) throws IOException { + private RestClient getRestClient(Map 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 int defaultPort = getValue(HOSTPORT, params); final String adminUser = getValue(USER, params); + final boolean sslRejectUnauthorized = getValue(SSL_REJECT_UNAUTHORIZED, params); + final String type = user == null || adminUser == null || user.equals(adminUser) ? "ADMIN" : "PROXY_USER"; + final String clientKey = String.format("%s @ %s:%d", user, hostName, defaultPort); + + /* if we already have the client ... just return it so we don't needlessly create the builder */ + if(clients.containsKey(clientKey)) + return clients.get(clientKey); + + /* creating the builder can throw IOException so we can't do it in the following functional call */ + RestClientBuilder rcb = createClientBuilder(user, password, type, hostName.split(","), defaultPort, sslRejectUnauthorized); + + /* note: ConcurrentHashMap performs this atomically ... only once per key */ + return this.clients.computeIfAbsent(clientKey, (key) -> { + LOGGER.info(String.format("Building a %s RestClient for", type, key)); + return rcb.build(); + }); + } + + /** + * Create a {@link RestClientBuilder}. + * + * @param user String + * @param password String + * @param type String + * @param hosts String[] + * @param defaultPort int + * @param sslRejectUnauthorized boolean + * @return RestClientBuilder + * @throws IOException when the hosts can not be parsed. + */ + private RestClientBuilder createClientBuilder(String user, String password, String type, String[] hosts, + int defaultPort, boolean sslRejectUnauthorized) throws IOException { final Pattern pattern = Pattern.compile("(?https?)?(://)?(?[^:]+):?(?\\d+)?"); final HttpHost[] httpHosts = new HttpHost[hosts.length]; @@ -282,8 +322,7 @@ private RestClient createRestClient(Map params, String use return httpClientBuilder; }); - LOGGER.fine(String.format("Building a %s RestClient for %s @ %s:%d", type, user, hostName, defaultPort)); - return builder.build(); + return builder; } @Override diff --git a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFinderIT.java b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFinderIT.java index 54df9e1..1e45a1c 100644 --- a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFinderIT.java +++ b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFinderIT.java @@ -132,7 +132,7 @@ private List getHosts(String hosts) throws IOException { Map params = createConnectionParams(); params.put(ElasticDataStoreFactory.HOSTNAME.key, hosts); ElasticDataStoreFactory factory = new ElasticDataStoreFactory(); - return factory.createRestClient(params).getNodes().stream().map(Node::getHost).collect(Collectors.toList()); + return factory.getRestClient(params).getNodes().stream().map(Node::getHost).collect(Collectors.toList()); } } diff --git a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDatastoreFactoryTest.java b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDatastoreFactoryTest.java index d7a1d53..e78accd 100644 --- a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDatastoreFactoryTest.java +++ b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDatastoreFactoryTest.java @@ -192,6 +192,47 @@ public void testBuildClientWithMultipleHosts() throws IOException { assertNotNull(credentialsProviderCaptor.getValue().getCredentials(new AuthScope("localhost2", 9201))); } + @Test + public void testGetRestClientSameClientKey() throws IOException { + /* multiple calls with same host:port:user should call builder once */ + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", null))); + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", null))); + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", null))); + verify(clientBuilder, times(1)).build(); + + } + + @Test + public void testGetRestClientWithProxy() throws IOException { + /* single call with proxy user should call builder twice */ + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", "proxy"))); + verify(clientBuilder, times(2)).build(); + } + + @Test + public void testGetRestClientDifferentHost() throws IOException { + /* multiple calls with different host should call builder once per host */ + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost1", 9200, "admin", null))); + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost2", 9200, "admin", null))); + verify(clientBuilder, times(2)).build(); + } + + @Test + public void testGetRestClientDifferentPort() throws IOException { + /* multiple calls with different port should call builder once per port */ + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", null))); + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9201, "admin", null))); + verify(clientBuilder, times(2)).build(); + } + + @Test + public void testGetRestClientDifferentUser() throws IOException { + /* multiple calls with different user should call builder once per user */ + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "admin", null))); + assertNotNull(dataStoreFactory.createDataStore(getParams("localhost", 9200, "proxy", null))); + verify(clientBuilder, times(2)).build(); + } + @Test public void testCreateClientbuilder() { ElasticDataStoreFactory factory = new ElasticDataStoreFactory(); diff --git a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticFeatureFilterIT.java b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticFeatureFilterIT.java index ad253eb..a8b9b3b 100644 --- a/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticFeatureFilterIT.java +++ b/gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticFeatureFilterIT.java @@ -714,6 +714,32 @@ public void testDefaultMaxFeatures() throws Exception { assertEquals(2, features.size()); } + @Test + public void testDefaultMaxWithRequestMaxFeatures() throws Exception { + init(); + dataStore.setDefaultMaxFeatures(2); + Query q = new Query(); + + /* default max should not override explicit query max */ + q.setMaxFeatures(5); + List features = readFeatures(featureSource.getFeatures(q).features()); + assertEquals(5, features.size()); + } + + @Test + public void testDefaultMaxWithOffsetLimit() throws Exception { + init(); + dataStore.setDefaultMaxFeatures(2); + Query q = new Query(); + + /* default max should not override offset-limit */ + q.setStartIndex(0); + q.setMaxFeatures(5); + List features = readFeatures(featureSource.getFeatures(q).features()); + assertEquals(5, features.size()); + } + + private void assertCovered(SimpleFeatureCollection features, Integer... ids) { assertEquals(ids.length, features.size());