Re-use RestClient across DataStores ... only need 1 per host:port:user

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<String, RestClient> clients = new ConcurrentHashMap<>();
+
     @Override
     public String getDisplayName() {
         return DISPLAY_NAME;
@@ -185,8 +195,8 @@ public DataStore createDataStore(Map<String, Serializable> 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<String, Serializable> params) throws IOException {
-        return createRestClient(params, null, null);
+    RestClient getRestClient(Map<String, Serializable> params) throws IOException {
+        return getRestClient(params, null, null);
     }
 
-    private RestClient createRestClient(Map<String, Serializable> params, String user, String password) throws IOException {
+    private RestClient getRestClient(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 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("(?<scheme>https?)?(://)?(?<host>[^:]+):?(?<port>\\d+)?");
         final HttpHost[] httpHosts = new HttpHost[hosts.length];
@@ -282,8 +322,7 @@ private RestClient createRestClient(Map<String, Serializable> 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

gt-elasticsearch/src/test/java/mil/nga/giat/data/elasticsearch/ElasticDataStoreFinderIT.java

@@ -132,7 +132,7 @@ private List<HttpHost> getHosts(String hosts) throws IOException {
         Map<String,Serializable> 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());
     }
 
 }

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();

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<SimpleFeature> 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<SimpleFeature> features = readFeatures(featureSource.getFeatures(q).features());
+        assertEquals(5, features.size());
+    }
+
+
     private void assertCovered(SimpleFeatureCollection features, Integer... ids) {
         assertEquals(ids.length, features.size());