Gracefully shutdown worker on WorkerTerminate request

Author: shreyas-gopalakrishna

src/main/java/com/microsoft/azure/functions/worker/Application.java

@@ -3,6 +3,7 @@
 import java.util.logging.*;
 import javax.annotation.*;
 
+import com.microsoft.azure.functions.worker.handler.GracefulTerminationException;
 import org.apache.commons.cli.*;
 import org.apache.commons.lang3.exception.ExceptionUtils;
 

src/main/java/com/microsoft/azure/functions/worker/JavaWorkerClient.java

@@ -42,6 +42,7 @@ private void addHandlers() {
         this.handlerSuppliers.put(StreamingMessage.ContentCase.FUNCTION_LOAD_REQUEST, () -> new FunctionLoadRequestHandler(broker));
         this.handlerSuppliers.put(StreamingMessage.ContentCase.INVOCATION_REQUEST, () -> new InvocationRequestHandler(broker));
         this.handlerSuppliers.put(StreamingMessage.ContentCase.WORKER_STATUS_REQUEST, WorkerStatusRequestHandler::new);
+        this.handlerSuppliers.put(StreamingMessage.ContentCase.WORKER_TERMINATE, () -> new WorkerTerminateRequestHandler(this));
     }
 
     public Future<Void> listen(String workerId, String requestId) {

src/main/java/com/microsoft/azure/functions/worker/WorkerLogManager.java

@@ -13,7 +13,8 @@ public class WorkerLogManager {
     public static Logger getInvocationLogger(String invocationId) { return INSTANCE.getInvocationLoggerImpl(invocationId); }
 
     static void initialize(JavaWorkerClient client, boolean logToConsole) { INSTANCE.initializeImpl(client, logToConsole); }
-    static void deinitialize() { INSTANCE.deinitializempl(); }
+    static void deinitialize() { INSTANCE.deinitializImpl(); }
+    public static void flushLogs() { INSTANCE.flushLogsImpl(); }
 
     private WorkerLogManager() {
         this.client = null;
@@ -28,20 +29,33 @@ private void initializeImpl(JavaWorkerClient client, boolean logToConsole) {
         assert this.client == null && client != null;
         this.client = client;
         this.logToConsole = logToConsole;
-        addHostClientHandlers(this.hostLogger, null);
+        addHandlers(this.hostLogger, null);
     }
 
-    private void deinitializempl() {
+    private void deinitializImpl() {
         assert this.client != null;
         clearHandlers(this.hostLogger);
         this.logToConsole = false;
         this.client = null;
     }
 
+    private void flushLogsImpl() {
+        flush(emptyLogger);
+        flush(systemLogger);
+        flush(hostLogger);
+    }
+
+    private static void flush(Logger logger) {
+        Handler[] handlers = logger.getHandlers();
+        for (Handler handler : handlers) {
+            handler.flush();
+        }
+    }
+
     private Logger getInvocationLoggerImpl(String invocationId) {
         Logger logger = Logger.getAnonymousLogger();
         logger.setLevel(Level.ALL);
-        addHostClientHandlers(logger, invocationId);
+        addHandlers(logger, invocationId);
         return logger;
     }
 
@@ -58,7 +72,7 @@ private static void addSystemHandler(Logger logger) {
         logger.addHandler(new SystemLoggerListener());
     }
 
-    private void addHostClientHandlers(Logger logger, String invocationId) {
+    private void addHandlers(Logger logger, String invocationId) {
         assert this.client != null;
         clearHandlers(logger);
         logger.addHandler(new HostLoggerListener(this.client, invocationId));

src/main/java/com/microsoft/azure/functions/worker/handler/GracefulTerminationException.java

@@ -0,0 +1,23 @@
+package com.microsoft.azure.functions.worker.handler;
+
+public class GracefulTerminationException extends RuntimeException {
+
+    public GracefulTerminationException() {
+    }
+
+    public GracefulTerminationException(String message) {
+        super(message);
+    }
+
+    public GracefulTerminationException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public GracefulTerminationException(Throwable cause) {
+        super(cause);
+    }
+
+    public GracefulTerminationException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+}

src/main/java/com/microsoft/azure/functions/worker/handler/MessageHandler.java

@@ -1,6 +1,5 @@
 package com.microsoft.azure.functions.worker.handler;
 
-import java.util.concurrent.*;
 import java.util.function.*;
 import java.util.logging.*;
 

src/main/java/com/microsoft/azure/functions/worker/handler/WorkerInitRequestHandler.java

@@ -24,6 +24,7 @@ String execute(WorkerInitRequest request, WorkerInitResponse.Builder response) {
         response.putCapabilities("WorkerStatus", "WorkerStatus");
         response.putCapabilities("RpcHttpBodyOnly", "RpcHttpBodyOnly");
         response.putCapabilities("RpcHttpTriggerMetadataRemoved", "RpcHttpTriggerMetadataRemoved");
+        response.putCapabilities("HandlesWorkerTerminateMessage", "HandlesWorkerTerminateMessage");
         response.setWorkerMetadata(composeWorkerMetadata());
         return "Worker initialized";
     }

src/main/java/com/microsoft/azure/functions/worker/handler/WorkerTerminateRequestHandler.java

@@ -0,0 +1,32 @@
+package com.microsoft.azure.functions.worker.handler;
+
+import com.google.protobuf.Message;
+import com.microsoft.azure.functions.rpc.messages.StreamingMessage;
+import com.microsoft.azure.functions.rpc.messages.WorkerTerminate;
+import com.microsoft.azure.functions.worker.JavaWorkerClient;
+import com.microsoft.azure.functions.worker.WorkerLogManager;
+import org.apache.commons.lang3.exception.ExceptionUtils;
+
+import java.util.logging.Level;
+
+public class WorkerTerminateRequestHandler extends MessageHandler<WorkerTerminate, Message.Builder> {
+
+    public WorkerTerminateRequestHandler(JavaWorkerClient client) {
+        super(StreamingMessage::getWorkerTerminate,
+                () -> null,
+                null,
+                null);
+        this.client = client;
+    }
+
+    @Override
+    String execute(WorkerTerminate workerTerminate, Message.Builder builder) {
+        WorkerLogManager.getSystemLogger().log(Level.INFO, "Worker terminate request received. Gracefully shutting down the worker.");
+        // Flushing the logs. Keeping the grpc client connection open here in case it is used by user's shutdown hooks
+        WorkerLogManager.flushLogs();
+        System.exit(0);
+        return null;
+    }
+
+    private final JavaWorkerClient client;
+}