Adjust the Greengrass Discovery client to use its own executor service (

#427) * Adjust the Greengrass Discovery client to use its own executor service --------- Co-authored-by: Bret Ambrose <[email protected]> Co-authored-by: Steve Kim <[email protected]>
aws · Jun 15, 2023 · 242e97d · 242e97d
1 parent 1b184d5
commit 242e97d
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 4 deletions.
diff --git a/sdk/src/main/java/software/amazon/awssdk/iot/discovery/DiscoveryClient.java b/sdk/src/main/java/software/amazon/awssdk/iot/discovery/DiscoveryClient.java
@@ -13,6 +13,9 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Class for performing network-based discovery of the connectivity properties of registered greengrass cores
@@ -35,12 +38,27 @@ public class DiscoveryClient implements AutoCloseable {
 
     private final HttpClientConnectionManager httpClientConnectionManager;
 
+    /**
+     * We need to use a code defined executor to avoid leaving threads alive when the discovery client is closed
+     * It also fixes issues with the SecurityManager as well - as created ExecutorServices created via code
+     * inherit the permissions of the application, unlike the default common thread pool that is used by default
+     * with supplyAsync otherwise.
+     */
+    private ExecutorService executorService = null;
+    private boolean cleanExecutor = false;
+
     /**
      *
      * @param config Greengrass discovery client configuration
      */
     public DiscoveryClient(final DiscoveryClientConfig config) {
-        this.httpClientConnectionManager = HttpClientConnectionManager.create(
+        executorService = config.getDiscoveryExecutor();
+        if (executorService == null) {
+            // If an executor is not set, then create one and make sure to clean it when finished.
+            executorService = Executors.newFixedThreadPool(1);
+            cleanExecutor = true;
+        }
+        httpClientConnectionManager = HttpClientConnectionManager.create(
                 new HttpClientConnectionManagerOptions()
                     .withClientBootstrap(config.getBootstrap())
                     .withProxyOptions(config.getProxyOptions())
@@ -102,11 +120,11 @@ public void onResponseComplete(HttpStream httpStream, int errorCode) {
             catch(InterruptedException | ExecutionException e) {
                 throw new RuntimeException(e.getMessage(), e);
             }
-        });
+        }, executorService);
     }
 
     private static String getHostname(final DiscoveryClientConfig config) {
-        //allow greengrass server endpoint to be manualy set for unique endpoints
+        //allow greengrass server endpoint to be manually set for unique endpoints
         if (config.getGGServerName().equals("")) {
             return String.format("greengrass-ats.iot.%s.%s",
                 config.getRegion(), AWS_DOMAIN_SUFFIX_MAP.getOrDefault(config.getRegion(), AWS_DOMAIN_DEFAULT));
@@ -120,5 +138,20 @@ public void close() {
         if(httpClientConnectionManager != null) {
             httpClientConnectionManager.close();
         }
+        if (cleanExecutor == true) {
+            executorService.shutdown();
+            try{
+                // Give the executorService 30 seconds to finish existing tasks. If it takes longer, force it to shutdown
+                if(!executorService.awaitTermination(30,TimeUnit.SECONDS)){
+                    executorService.shutdownNow();
+                }
+            } catch (InterruptedException ie){
+                // If current thread is interrupted, force executorService shutdown
+                executorService.shutdownNow();
+                // Preserve interrupt status
+                Thread.currentThread().interrupt();
+            }
+        }
+        executorService = null;
     }
 }
diff --git a/sdk/src/main/java/software/amazon/awssdk/iot/discovery/DiscoveryClientConfig.java b/sdk/src/main/java/software/amazon/awssdk/iot/discovery/DiscoveryClientConfig.java
@@ -1,5 +1,7 @@
 package software.amazon.awssdk.iot.discovery;
 
+import java.util.concurrent.ExecutorService;
+
 import software.amazon.awssdk.crt.http.HttpProxyOptions;
 import software.amazon.awssdk.crt.io.ClientBootstrap;
 import software.amazon.awssdk.crt.io.SocketOptions;
@@ -17,6 +19,7 @@ public class DiscoveryClientConfig implements AutoCloseable {
     final private int maxConnections;
     final private HttpProxyOptions proxyOptions;
     final private String ggServerName;
+    private ExecutorService discoveryExecutor;
 
     /**
      * Constructor for DiscoveryClientConfig creates the correct endpoint if not in a special region
@@ -44,6 +47,7 @@ public DiscoveryClientConfig(
         this.maxConnections = maxConnections;
         this.proxyOptions = proxyOptions;
         this.ggServerName = "";
+        this.discoveryExecutor = null;
     }
 
     /**
@@ -70,7 +74,7 @@ public DiscoveryClientConfig(
 
     /**
      * Default Constructor for DiscoveryClientConfig that allows the specification of a specific ggServerName to use in special regions
-     * 
+     *
      * @param bootstrap client bootstrap to use to establish network connections
      * @param tlsContextOptions tls configuration for client network connections.  For greengrass discovery, the
      *                          tls context must be initialized with the certificate and private key of the
@@ -94,6 +98,7 @@ public DiscoveryClientConfig(
         this.maxConnections = maxConnections;
         this.proxyOptions = proxyOptions;
         this.ggServerName = ggServerName;
+        this.discoveryExecutor = null;
     }
 
     /**
@@ -164,6 +169,27 @@ public String getGGServerName() {
         return ggServerName;
     }
 
+    /**
+     * @return the ExecutorService set for this discover client config. Returns null if one hasn't been set.
+     */
+    public ExecutorService getDiscoveryExecutor() {
+        return this.discoveryExecutor;
+    }
+
+    /**
+     * Sets the executor that is used when calling discover().
+     * If set using this function, you are expected to close/shutdown the executor when finished with it.
+     *
+     * If it is not set, the discovery client will internally manage its own executor.
+     *
+     * @param override The executor to use
+     * @return The client config
+     */
+    public DiscoveryClientConfig setDiscoveryExecutor(ExecutorService override) {
+        this.discoveryExecutor = override;
+        return this;
+    }
+
     @Override
     public void close() {
         if(tlsContext != null) {