use exeucutor service instead

Author: jeremyber-aws

java/AsyncIO/src/main/java/com/amazonaws/services/msf/ProcessingFunction.java

@@ -13,20 +13,17 @@
 
 import java.time.Duration;
 import java.util.Collections;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 
 public class ProcessingFunction extends RichAsyncFunction<IncomingEvent, ProcessedEvent> {
     private static final Logger LOG = LoggerFactory.getLogger(ProcessingFunction.class);
 
     private final String apiUrl;
     private final String apiKey;
-    private final long shutdownWaitTS;
-    private final int threadPoolSize;
+    private static ExecutorService executorService;
+
 
     private transient AsyncHttpClient client;
-    private transient ExecutorService executorService;
 
     public ProcessingFunction(String apiUrl, String apiKey) {
         Preconditions.checkNotNull(apiUrl, "API URL must not be null");
@@ -36,70 +33,62 @@ public ProcessingFunction(String apiUrl, String apiKey) {
 
         this.apiUrl = apiUrl;
         this.apiKey = apiKey;
-        this.shutdownWaitTS = 20000;     // Max time (ms) to wait for ExecutorService shutdown before forcing termination
-        this.threadPoolSize = 30;       // Number of threads in the ExecutorService pool for async operations
     }
+    /**
+     * Instantiate the connection to an async client here to use within asyncInvoke
+     */
 
-
-     // Instantiate the connection to an async client here to use within asyncInvoke
     @Override
     public void open(Configuration parameters) throws Exception {
         DefaultAsyncHttpClientConfig.Builder clientBuilder = Dsl.config().setConnectTimeout(Duration.ofSeconds(10));
         client = Dsl.asyncHttpClient(clientBuilder);
-        executorService = Executors.newFixedThreadPool(threadPoolSize);
-        LOG.info("Initialized ExecutorService with {} threads", threadPoolSize);
+
+        int numCores = Runtime.getRuntime().availableProcessors(); // get num cores on node for thread count
+        executorService = Executors.newFixedThreadPool(numCores);
+
     }
 
-    // close Async Client, Executor Service on shutdown
     @Override
-    public void close() throws Exception {
-        if (client != null) {
-            client.close();
-        }
-        if (executorService != null) {
-            executorService.shutdown();
-            try {
-                if (!executorService.awaitTermination(shutdownWaitTS, TimeUnit.MILLISECONDS)) {
-                    executorService.shutdownNow();
-                }
-            } catch (InterruptedException e) {
-                executorService.shutdownNow();
-            }
-        }
+    public void close() throws Exception
+    {
+        client.close();
+        executorService.shutdown();
     }
 
     @Override
-    public void asyncInvoke(final IncomingEvent incomingEvent, final ResultFuture<ProcessedEvent> resultFuture) {
+    public void asyncInvoke(IncomingEvent incomingEvent, ResultFuture<ProcessedEvent> resultFuture) {
+
         // Create a new ProcessedEvent instance
         ProcessedEvent processedEvent = new ProcessedEvent(incomingEvent.getMessage());
         LOG.debug("New request: {}", incomingEvent);
 
-        executorService.submit(new Runnable() {
-            @Override
-            public void run() {
+        // Note: The Async Client used must return a Future object or equivalent
+        Future<Response> future = client.prepareGet(apiUrl)
+                .setHeader("x-api-key", apiKey)
+                .execute();
+
+        // Process the request via a Completable Future, in order to not block request synchronously
+        // Notice we are passing executor service for thread management
+        CompletableFuture.supplyAsync(() ->
+            {
                 try {
                     LOG.debug("Trying to get response for {}", incomingEvent.getId());
-                    Response response = client.prepareGet(apiUrl)
-                            .setHeader("x-api-key", apiKey)
-                            .execute()
-                            .get();
-
-                    int statusCode = response.getStatusCode();
-
-                    if (statusCode == 200) {
-                        LOG.debug("Success! {}", incomingEvent.getId());
-                        resultFuture.complete(Collections.singleton(processedEvent));
-                    } else if (statusCode == 500) { // Retryable error
-                        LOG.error("Status code 500, retrying shortly...");
-                        resultFuture.completeExceptionally(new Throwable(String.valueOf(statusCode)));
-                    } else {
-                        LOG.error("Unexpected status code: {}", statusCode);
-                        resultFuture.completeExceptionally(new Throwable(String.valueOf(statusCode)));
-                    }
-                } catch (Exception e) {
+                    Response response = future.get();
+                    return response.getStatusCode();
+                } catch (InterruptedException | ExecutionException e) {
                     LOG.error("Error during async HTTP call: {}", e.getMessage());
-                    resultFuture.completeExceptionally(e);
+                    return -1;
                 }
+            }, executorService).thenAccept(statusCode -> {
+            if (statusCode == 200) {
+                LOG.debug("Success! {}", incomingEvent.getId());
+                resultFuture.complete(Collections.singleton(processedEvent));
+            } else if (statusCode == 500) { // Retryable error
+                LOG.error("Status code 500, retrying shortly...");
+                resultFuture.completeExceptionally(new Throwable(statusCode.toString()));
+            } else {
+                LOG.error("Unexpected status code: {}", statusCode);
+                resultFuture.completeExceptionally(new Throwable(statusCode.toString()));
             }
         });
     }