(fix) Mdodsworth/graceful shutdown (#236)

thomaszurkan-optimizely · web-flow · commit af3af22ebeb8 · 2018-11-20T11:36:06.000-08:00
* adding graceful shutdown to the async event dispatcher

* adding tests for AsyncEventHandler

* run through travis and compat tests
diff --git a/core-httpclient-impl/build.gradle b/core-httpclient-impl/build.gradle
@@ -1,5 +1,6 @@
 dependencies {
     compile project(':core-api')
+    testCompile project(':core-api').sourceSets.test.output
 
     compileOnly group: 'com.google.code.gson', name: 'gson', version: gsonVersion
 
diff --git a/core-httpclient-impl/src/main/java/com/optimizely/ab/event/AsyncEventHandler.java b/core-httpclient-impl/src/main/java/com/optimizely/ab/event/AsyncEventHandler.java
@@ -18,6 +18,7 @@
 
 import com.optimizely.ab.HttpClientUtils;
 import com.optimizely.ab.NamedThreadFactory;
+import com.optimizely.ab.annotations.VisibleForTesting;
 
 import org.apache.http.HttpResponse;
 import org.apache.http.client.ClientProtocolException;
@@ -33,23 +34,23 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.Closeable;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.net.URISyntaxException;
 import java.util.Map;
 import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 
 import javax.annotation.CheckForNull;
 
 /**
  * {@link EventHandler} implementation that queues events and has a separate pool of threads responsible
  * for the dispatch.
  */
-public class AsyncEventHandler implements EventHandler, Closeable {
+public class AsyncEventHandler implements EventHandler {
 
     // The following static values are public so that they can be tweaked if necessary.
     // These are the recommended settings for http protocol.  https://hc.apache.org/httpcomponents-client-ga/tutorial/html/connmgmt.html
@@ -65,7 +66,6 @@ public class AsyncEventHandler implements EventHandler, Closeable {
 
     private final CloseableHttpClient httpClient;
     private final ExecutorService workerExecutor;
-    private final BlockingQueue<LogEvent> logEventQueue;
 
     public AsyncEventHandler(int queueCapacity, int numWorkers) {
         this(queueCapacity, numWorkers, 200, 20, 5000);
@@ -80,21 +80,22 @@ public AsyncEventHandler(int queueCapacity, int numWorkers, int maxConnections,
         this.maxPerRoute = connectionsPerRoute;
         this.validateAfterInactivity = validateAfter;
 
-        this.logEventQueue = new ArrayBlockingQueue<LogEvent>(queueCapacity);
-        this.httpClient = HttpClients.custom()
+      this.httpClient = HttpClients.custom()
             .setDefaultRequestConfig(HttpClientUtils.DEFAULT_REQUEST_CONFIG)
                 .setConnectionManager(poolingHttpClientConnectionManager())
             .disableCookieManagement()
             .build();
 
-        this.workerExecutor = Executors.newFixedThreadPool(
-            numWorkers, new NamedThreadFactory("optimizely-event-dispatcher-thread-%s", true));
+        this.workerExecutor = new ThreadPoolExecutor(numWorkers, numWorkers,
+                                                     0L, TimeUnit.MILLISECONDS,
+                                                     new ArrayBlockingQueue<Runnable>(queueCapacity),
+                                                     new NamedThreadFactory("optimizely-event-dispatcher-thread-%s", true));
+    }
 
-        // create dispatch workers
-        for (int i = 0; i < numWorkers; i++) {
-            EventDispatchWorker worker = new EventDispatchWorker();
-            workerExecutor.submit(worker);
-        }
+    @VisibleForTesting
+    public AsyncEventHandler(CloseableHttpClient httpClient, ExecutorService workerExecutor) {
+        this.httpClient = httpClient;
+        this.workerExecutor = workerExecutor;
     }
 
     private PoolingHttpClientConnectionManager poolingHttpClientConnectionManager()
@@ -108,61 +109,87 @@ private PoolingHttpClientConnectionManager poolingHttpClientConnectionManager()
 
     @Override
     public void dispatchEvent(LogEvent logEvent) {
-        // attempt to enqueue the log event for processing
-        boolean submitted = logEventQueue.offer(logEvent);
-        if (!submitted) {
-            logger.error("unable to enqueue event because queue is full");
+        try {
+            // attempt to enqueue the log event for processing
+            workerExecutor.execute(new EventDispatcher(logEvent));
+        } catch (RejectedExecutionException e) {
+            logger.error("event dispatch rejected");
         }
     }
 
-    @Override
-    public void close() throws IOException {
-        logger.info("closing event dispatcher");
+    /**
+     * Attempts to gracefully terminate all event dispatch workers and close all resources.
+     * This method blocks, awaiting the completion of any queued or ongoing event dispatches.
+     *
+     * Note: termination of ongoing event dispatching is best-effort.
+     *
+     * @param timeout maximum time to wait for event dispatches to complete
+     * @param unit the time unit of the timeout argument
+     */
+    public void shutdownAndAwaitTermination(long timeout, TimeUnit unit) {
+
+        // Disable new tasks from being submitted
+        logger.info("event handler shutting down. Attempting to dispatch previously submitted events");
+        workerExecutor.shutdown();
 
-        // "close" all workers and the http client
         try {
-            httpClient.close();
-        } catch (IOException e) {
-            logger.error("unable to close the event handler httpclient cleanly", e);
-        } finally {
+            // Wait a while for existing tasks to terminate
+            if (!workerExecutor.awaitTermination(timeout, unit)) {
+                int unprocessedCount = workerExecutor.shutdownNow().size();
+                logger.warn("timed out waiting for previously submitted events to be dispatched. "
+                            + "{} events were dropped. "
+                            + "Interrupting dispatch worker(s)", unprocessedCount);
+                // Cancel currently executing tasks
+                // Wait a while for tasks to respond to being cancelled
+                if (!workerExecutor.awaitTermination(timeout, unit)) {
+                    logger.error("unable to gracefully shutdown event handler");
+                }
+            }
+        } catch (InterruptedException ie) {
+            // (Re-)Cancel if current thread also interrupted
             workerExecutor.shutdownNow();
+            // Preserve interrupt status
+            Thread.currentThread().interrupt();
+        } finally {
+            try {
+                httpClient.close();
+            } catch (IOException e) {
+                logger.error("unable to close event dispatcher http client", e);
+            }
         }
+
+        logger.info("event handler shutdown complete");
     }
 
     //======== Helper classes ========//
 
-    private class EventDispatchWorker implements Runnable {
+    /**
+     * Wrapper runnable for the actual event dispatch.
+     */
+    private class EventDispatcher implements Runnable {
+
+        private final LogEvent logEvent;
+
+        EventDispatcher(LogEvent logEvent) {
+            this.logEvent = logEvent;
+        }
 
         @Override
         public void run() {
-            boolean terminate = false;
-
-            logger.info("starting event dispatch worker");
-            // event loop that'll block waiting for events to appear in the queue
-            //noinspection InfiniteLoopStatement
-            while (!terminate) {
-                try {
-                    LogEvent event = logEventQueue.take();
-                    HttpRequestBase request;
-                    if (event.getRequestMethod() == LogEvent.RequestMethod.GET) {
-                        request = generateGetRequest(event);
-                    } else {
-                        request = generatePostRequest(event);
-                    }
-                    httpClient.execute(request, EVENT_RESPONSE_HANDLER);
-                } catch (InterruptedException e) {
-                    logger.info("terminating event dispatcher event loop");
-                    terminate = true;
-                } catch (Throwable t) {
-                    logger.error("event dispatcher threw exception but will continue", t);
-                }
+            try {
+                HttpGet request = generateRequest(logEvent);
+                httpClient.execute(request, EVENT_RESPONSE_HANDLER);
+            } catch (IOException e) {
+                logger.error("event dispatch failed", e);
+            } catch (URISyntaxException e) {
+                logger.error("unable to parse generated URI", e);
             }
         }
 
         /**
          * Helper method that generates the event request for the given {@link LogEvent}.
          */
-        private HttpGet generateGetRequest(LogEvent event) throws URISyntaxException {
+        private HttpGet generateRequest(LogEvent event) throws URISyntaxException {
 
             URIBuilder builder = new URIBuilder(event.getEndpointUrl());
             for (Map.Entry<String, String> param : event.getRequestParams().entrySet()) {
@@ -171,13 +198,6 @@ private HttpGet generateGetRequest(LogEvent event) throws URISyntaxException {
 
             return new HttpGet(builder.build());
         }
-
-        private HttpPost generatePostRequest(LogEvent event) throws UnsupportedEncodingException {
-            HttpPost post = new HttpPost(event.getEndpointUrl());
-            post.setEntity(new StringEntity(event.getBody()));
-            post.addHeader("Content-Type", "application/json");
-            return post;
-        }
     }
 
     /**
diff --git a/core-httpclient-impl/src/test/java/com/optimizely/ab/event/AsyncEventHandlerTest.java b/core-httpclient-impl/src/test/java/com/optimizely/ab/event/AsyncEventHandlerTest.java
@@ -0,0 +1,138 @@
+/**
+ *
+ *    Copyright 2016, Optimizely
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License");
+ *    you may not use this file except in compliance with the License.
+ *    You may obtain a copy of the License at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS,
+ *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *    See the License for the specific language governing permissions and
+ *    limitations under the License.
+ */
+package com.optimizely.ab.event;
+
+import com.google.common.util.concurrent.MoreExecutors;
+
+import com.optimizely.ab.event.internal.payload.EventBatch;
+import org.apache.http.client.ResponseHandler;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link AsyncEventHandler}.
+ */
+public class AsyncEventHandlerTest {
+
+    @Rule
+    @SuppressFBWarnings("URF_UNREAD_PUBLIC_OR_PROTECTED_FIELD")
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+
+    @Mock CloseableHttpClient mockHttpClient;
+    @Mock ExecutorService mockExecutorService;
+
+    @Test
+    public void testQueueCapacityPreconditionCheck() throws Exception {
+        thrown.expect(IllegalArgumentException.class);
+        new AsyncEventHandler(-1, 1);
+    }
+
+    @Test
+    public void testDispatch() throws Exception {
+        AsyncEventHandler eventHandler = new AsyncEventHandler(mockHttpClient, MoreExecutors.newDirectExecutorService());
+        eventHandler.dispatchEvent(createLogEvent());
+        verify(mockHttpClient).execute(any(HttpGet.class), any(ResponseHandler.class));
+    }
+
+    /**
+     * Verify that {@link RejectedExecutionException}s are caught, rather than being propagated.
+     */
+    @Test
+    public void testRejectedExecutionsAreHandled() throws Exception {
+        AsyncEventHandler eventHandler = new AsyncEventHandler(mockHttpClient, mockExecutorService);
+        doThrow(RejectedExecutionException.class).when(mockExecutorService).execute(any(Runnable.class));
+        eventHandler.dispatchEvent(createLogEvent());
+    }
+
+    /**
+     * Verify that {@link IOException}s are caught, rather than being propagated (which would cause a worker
+     * thread to die).
+     */
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testIOExceptionsCaughtInDispatch() throws Exception {
+        AsyncEventHandler eventHandler = new AsyncEventHandler(mockHttpClient, MoreExecutors.newDirectExecutorService());
+
+        // have the http client throw an IOException on execute
+        when(mockHttpClient.execute(any(HttpGet.class), any(ResponseHandler.class))).thenThrow(IOException.class);
+        eventHandler.dispatchEvent(createLogEvent());
+        verify(mockHttpClient).execute(any(HttpGet.class), any(ResponseHandler.class));
+    }
+
+    /**
+     * Verifies the case where all queued events could be processed before the timeout is exceeded.
+     */
+    @Test
+    public void testShutdownAndAwaitTermination() throws Exception {
+        AsyncEventHandler eventHandler = new AsyncEventHandler(mockHttpClient, mockExecutorService);
+        when(mockExecutorService.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(true);
+
+        eventHandler.shutdownAndAwaitTermination(1, TimeUnit.SECONDS);
+        verify(mockExecutorService).shutdown();
+        verify(mockExecutorService, never()).shutdownNow();
+
+        verify(mockHttpClient).close();
+    }
+
+    /**
+     * Verify the case where all queued events count NOT be processed before the timeout was exceeded.
+     * {@link ExecutorService#shutdownNow()} should be called to drop the queued events and attempt to interrupt
+     * ongoing tasks.
+     */
+    @Test
+    public void testShutdownAndForcedTermination() throws Exception {
+        AsyncEventHandler eventHandler = new AsyncEventHandler(mockHttpClient, mockExecutorService);
+        when(mockExecutorService.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(false);
+
+        eventHandler.shutdownAndAwaitTermination(1, TimeUnit.SECONDS);
+        verify(mockExecutorService).shutdown();
+        verify(mockExecutorService).shutdownNow();
+        verify(mockHttpClient).close();
+    }
+
+    //======== Helper methods ========//
+
+    private LogEvent createLogEvent() {Map<String, String> testParams = new HashMap<String, String>();
+        testParams.put("test", "params");
+        return new LogEvent(LogEvent.RequestMethod.GET, "test_url", testParams, new EventBatch());
+    }
+}
