Supporting graceful shutdown based on SIGTERM handler. (#749)

* Supporting graceful shutdown based on SIGTERM handler.

Author: mindaugasbarcauskas

src/main/java/com/uber/cadence/internal/sync/WorkflowClientInternal.java

@@ -408,7 +408,7 @@ public static <A1, A2, A3, A4, A5, A6, R> CompletableFuture<R> execute(
     return execute(() -> workflow.apply(arg1, arg2, arg3, arg4, arg5, arg6));
   }
 
-  private synchronized static void emitClientVersion(WorkflowClientOptions options) {
+  private static synchronized void emitClientVersion(WorkflowClientOptions options) {
     if (emittingClientVersion) {
       return;
     }

src/main/java/com/uber/cadence/internal/worker/WorkerShutDownHandler.java

@@ -0,0 +1,69 @@
+/*
+ *  Modifications Copyright (c) 2020-2022 Uber Technologies Inc.
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file 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.uber.cadence.internal.worker;
+
+import com.uber.cadence.internal.common.InternalUtils;
+import com.uber.cadence.worker.WorkerFactory;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+public class WorkerShutDownHandler {
+
+  private static final List<WorkerFactory> workerFactories = new ArrayList<>();
+  private static Thread registeredHandler;
+
+  public static void registerHandler() {
+    if (registeredHandler != null) {
+      return;
+    }
+
+    registeredHandler = new Thread("SHUTDOWN_WORKERS") {
+      @Override
+      public void run() {
+        for (WorkerFactory workerFactory : workerFactories) {
+          workerFactory.suspendPolling();
+        }
+
+        for (WorkerFactory workerFactory : workerFactories) {
+          workerFactory.shutdownNow();
+        }
+
+        long remainingTimeout = 10000;
+        for (WorkerFactory workerFactory : workerFactories) {
+          final long timeoutMillis = remainingTimeout;
+          remainingTimeout = InternalUtils.awaitTermination(timeoutMillis,
+                  () -> workerFactory.awaitTermination(timeoutMillis, TimeUnit.MILLISECONDS));
+        }
+      }
+    };
+
+    Runtime.getRuntime()
+        .addShutdownHook(registeredHandler);
+  }
+
+  public static synchronized void registerWorkerFactory(WorkerFactory workerFactory) {
+    if (workerFactory != null) {
+      workerFactories.add(workerFactory);
+    }
+  }
+
+  // Only for tests
+  protected static void execute() {
+    registeredHandler.run();
+  }
+}

src/main/java/com/uber/cadence/worker/WorkerFactory.java

@@ -30,6 +30,7 @@
 import com.uber.cadence.internal.worker.PollDecisionTaskDispatcher;
 import com.uber.cadence.internal.worker.Poller;
 import com.uber.cadence.internal.worker.PollerOptions;
+import com.uber.cadence.internal.worker.WorkerShutDownHandler;
 import com.uber.cadence.internal.worker.WorkflowPollTaskFactory;
 import com.uber.m3.tally.Scope;
 import com.uber.m3.util.ImmutableMap;
@@ -57,7 +58,10 @@ public static WorkerFactory newInstance(WorkflowClient workflowClient) {
 
   public static WorkerFactory newInstance(
       WorkflowClient workflowClient, WorkerFactoryOptions options) {
-    return new WorkerFactory(workflowClient, options);
+    WorkerShutDownHandler.registerHandler();
+    WorkerFactory workerFactory = new WorkerFactory(workflowClient, options);
+    WorkerShutDownHandler.registerWorkerFactory(workerFactory);
+    return workerFactory;
   }
 
   private final List<Worker> workers = new ArrayList<>();

src/test/java/com/uber/cadence/internal/worker/WorkerShutDownHandlerTest.java

@@ -0,0 +1,64 @@
+/*
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file 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.uber.cadence.internal.worker;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.when;
+
+import com.uber.cadence.client.WorkflowClient;
+import com.uber.cadence.client.WorkflowClientOptions;
+import com.uber.cadence.serviceclient.IWorkflowService;
+import com.uber.cadence.worker.WorkerFactory;
+import com.uber.m3.tally.NoopScope;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class WorkerShutDownHandlerTest {
+
+  @Mock private WorkflowClient mockClient;
+
+  @Mock private IWorkflowService mockService;
+
+  @Before
+  public void setup() {
+    WorkflowClientOptions clientOptions =
+        WorkflowClientOptions.newBuilder().setMetricsScope(new NoopScope()).build();
+    when(mockClient.getOptions()).thenReturn(clientOptions);
+    when(mockClient.getService()).thenReturn(mockService);
+  }
+
+  @Test
+  public void shutDownHookShutsDownFactories() {
+
+    WorkerFactory workerFactory = WorkerFactory.newInstance(mockClient);
+    workerFactory.newWorker("TL1");
+    workerFactory.newWorker("TL2");
+
+    WorkerFactory workerFactory2 = WorkerFactory.newInstance(mockClient);
+    workerFactory2.newWorker("TL3");
+
+    WorkerShutDownHandler.execute();
+
+    assertTrue(workerFactory.isShutdown());
+    assertTrue(workerFactory2.isShutdown());
+  }
+}