Add interception for system shutdown to centrally manage shutdown hooks

joylive-core/joylive-core-api/src/main/java/com/jd/live/agent/core/util/Close.java

@@ -15,6 +15,8 @@
  */
 package com.jd.live.agent.core.util;
 
+import java.util.function.Consumer;
+
 /**
  * A utility class that provides methods to facilitate the closing of resources and threads.
  * This class is designed to be used as a singleton and supports method chaining for ease of use.
@@ -85,6 +87,28 @@ public Close close(AutoCloseable... resources) {
         return this;
     }
 
+    /**
+     * Attempts to close or clean up the provided resource using the specified {@code closer} action.
+     *
+     * @param value  the resource to be closed or cleaned up; if {@code null}, the {@code closer} is not invoked
+     * @param closer a {@link Consumer} that accepts the resource and performs the necessary action to close
+     *               or clean it up
+     * @return this instance, enabling method chaining
+     */
+    public <T> Close closeIfExists(T value, Consumer<T> closer) {
+        if (value != null) {
+            closer.accept(value);
+        }
+        return this;
+    }
+
+    public Close close(Runnable closer) {
+        if (closer != null) {
+            closer.run();
+        }
+        return this;
+    }
+
     /**
      * Interrupts the given Thread and returns this Close instance for method chaining.
      * If the thread is null, this method does nothing.

joylive-core/joylive-core-api/src/main/java/com/jd/live/agent/core/util/shutdown/Shutdown.java

@@ -0,0 +1,203 @@
+/*
+ * Copyright © ${year} ${owner} (${email})
+ *
+ * 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.jd.live.agent.core.util.shutdown;
+
+import com.jd.live.agent.bootstrap.logger.Logger;
+import com.jd.live.agent.bootstrap.logger.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Manages the orderly shutdown of application components, ensuring that all registered
+ * {@link ShutdownHook} instances are executed within a specified timeout period.
+ * <p>
+ * This class allows for the graceful shutdown of resources by registering {@code ShutdownHook}
+ * instances which can be executed when the application is shutting down. Hooks can be added
+ * at any time before the shutdown process begins. Once the shutdown process starts, no new hooks
+ * can be added.
+ * </p>
+ * </p>
+ */
+public class Shutdown {
+
+    private static final Logger logger = LoggerFactory.getLogger(Shutdown.class);
+
+    private static final int DEFAULT_TIMEOUT = 10 * 1000;
+
+    private final List<ShutdownHook> hooks = new CopyOnWriteArrayList<>();
+
+    private final AtomicBoolean register = new AtomicBoolean();
+
+    private final AtomicBoolean shutdown = new AtomicBoolean();
+
+    private final long timeout;
+
+    private final Thread shutdownTask;
+
+    public Shutdown() {
+        this(DEFAULT_TIMEOUT);
+    }
+
+    public Shutdown(long timeout) {
+        this.timeout = timeout;
+        this.shutdownTask = new Thread(() -> {
+            try {
+                doShutdown().get(this.timeout, TimeUnit.MILLISECONDS);
+            } catch (InterruptedException | ExecutionException | TimeoutException ignored) {
+            }
+        }, "LiveAgent-Shutdown");
+    }
+
+    /**
+     * Initiates the shutdown process, executing all registered {@code ShutdownHook} instances
+     * within the specified timeout period.
+     * <p>
+     * This method is synchronized to ensure that the shutdown process is only executed once.
+     * If the shutdown process has already started, this method returns immediately with a
+     * completed future.
+     * </p>
+     *
+     * @return a {@link CompletableFuture} representing the pending completion of the shutdown process
+     */
+    private synchronized CompletableFuture<Void> doShutdown() {
+        if (!shutdown.compareAndSet(false, true)) {
+            return CompletableFuture.completedFuture(null);
+        }
+        logger.info("LiveAgent shutdown....");
+        CompletableFuture<Void> result = null;
+        if (!hooks.isEmpty()) {
+            List<ShutdownHookGroup> groups = sortGroup();
+
+            // Sequentially execute hooks
+            for (ShutdownHookGroup group : groups) {
+                if (result == null) {
+                    result = group.run();
+                } else {
+                    result = result.handle((t, r) -> group.run().join());
+                }
+            }
+        }
+        result = result == null ? CompletableFuture.completedFuture(null) : result;
+        return result.whenComplete((r, t) -> {
+            if (t == null) {
+                logger.info("LiveAgent shutdown successfully.");
+            } else {
+                logger.error("LiveAgent shutdown failed.", t);
+            }
+        });
+    }
+
+    /**
+     * Groups and sorts the registered shutdown hooks by their priority, preparing them for execution.
+     *
+     * @return a list of {@link ShutdownHookGroup} instances, each containing hooks of the same priority
+     */
+    private List<ShutdownHookGroup> sortGroup() {
+        // Group hook by priority
+        List<ShutdownHookGroup> groups = new ArrayList<>();
+
+        // Sort hook by priority.
+        hooks.sort(Comparator.comparingInt(ShutdownHook::priority));
+        ShutdownHookGroup lastGroup = null;
+
+        for (ShutdownHook hook : hooks) {
+            if (lastGroup == null || lastGroup.priority != hook.priority()) {
+                lastGroup = new ShutdownHookGroup(hook.priority());
+                lastGroup.add(hook);
+                groups.add(lastGroup);
+            } else if (lastGroup.priority == hook.priority()) {
+                lastGroup.add(hook);
+            }
+        }
+        return groups;
+    }
+
+    /**
+     * Registers a JVM shutdown hook to trigger the application shutdown process.
+     */
+    public synchronized void register() {
+        if (register.compareAndSet(false, true)) {
+            Runtime.getRuntime().addShutdownHook(shutdownTask);
+        }
+    }
+
+    /**
+     * Unregisters the JVM shutdown hook associated with this {@code Shutdown} instance, preventing
+     * the automatic execution of registered shutdown hooks upon JVM shutdown.
+     */
+    public synchronized void unregister() {
+        if (register.compareAndSet(true, false)) {
+            Runtime.getRuntime().removeShutdownHook(shutdownTask);
+        }
+    }
+
+    /**
+     * Adds a {@link ShutdownHook} instance to be executed during the shutdown process.
+     *
+     * @param hook the {@code ShutdownHook} to add
+     */
+    public void addHook(ShutdownHook hook) {
+        if (hook != null) {
+            hooks.add(hook);
+        }
+    }
+
+    /**
+     * Adds a {@link Runnable} as a shutdown hook to be executed during the shutdown process.
+     *
+     * @param runnable the {@code Runnable} to add as a shutdown hook
+     */
+    public void addHook(Runnable runnable) {
+        if (runnable != null) {
+            addHook(new ShutdownHookAdapter(runnable));
+        }
+    }
+
+    /**
+     * Checks if the shutdown process has started.
+     *
+     * @return {@code true} if the shutdown process has started, {@code false} otherwise
+     */
+    public boolean isShutdown() {
+        return shutdown.get();
+    }
+
+    /**
+     * Initiates the shutdown process and returns a {@link CompletableFuture} representing its completion.
+     *
+     * @return a {@code CompletableFuture} representing the pending completion of the shutdown process
+     */
+    public CompletableFuture<Void> shutdown() {
+        return doShutdown();
+    }
+
+    /**
+     * Initiates the shutdown process and waits for its completion within the specified timeout period.
+     *
+     * @param timeout the timeout in milliseconds to wait for the shutdown process to complete
+     * @throws InterruptedException if the current thread is interrupted while waiting
+     * @throws ExecutionException   if the shutdown process throws an exception
+     * @throws TimeoutException     if the timeout expires before the shutdown process is complete
+     */
+    public void shutdown(long timeout) throws InterruptedException, ExecutionException, TimeoutException {
+        doShutdown().get(timeout, TimeUnit.MILLISECONDS);
+    }
+
+}

joylive-core/joylive-core-api/src/main/java/com/jd/live/agent/core/util/shutdown/ShutdownHook.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright © ${year} ${owner} (${email})
+ *
+ * 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.jd.live.agent.core.util.shutdown;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Functional interface representing a shutdown hook.
+ * <p>
+ * A shutdown hook is a thread that is executed when the program is shutting down.
+ * This interface defines a method to run during the shutdown process with an optional
+ * priority that determines the order of execution if multiple hooks are present.
+ * </p>
+ */
+@FunctionalInterface
+public interface ShutdownHook {
+
+    /**
+     * The default priority that is used if no other priority is specified.
+     */
+    int DEFAULT_PRIORITY = 100;
+
+    /**
+     * Executes the shutdown hook logic asynchronously.
+     * <p>
+     * This method should contain the logic to be performed during the shutdown
+     * process. It returns a {@link CompletableFuture} which can be used to
+     * track the completion of the hook execution.
+     * </p>
+     *
+     * @return a CompletableFuture representing the pending completion of the hook
+     */
+    CompletableFuture<Void> run();
+
+    /**
+     * Returns the priority of this shutdown hook.
+     * <p>
+     * The priority is used to determine the order in which shutdown hooks are called.
+     * A lower number indicates a higher priority. If not overridden, the default
+     * priority ({@value #DEFAULT_PRIORITY}) is returned.
+     * </p>
+     *
+     * @return the priority of the hook
+     */
+    default int priority() {
+        return DEFAULT_PRIORITY;
+    }
+}
+

joylive-core/joylive-core-api/src/main/java/com/jd/live/agent/core/util/shutdown/ShutdownHookAdapter.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright © ${year} ${owner} (${email})
+ *
+ * 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.jd.live.agent.core.util.shutdown;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * An adapter class for {@link ShutdownHook} that allows the wrapping of {@link Runnable} instances
+ * or other {@code ShutdownHook} instances with custom priority and execution mode settings.
+ */
+public class ShutdownHookAdapter implements ShutdownHook {
+
+    /**
+     * The wrapped {@link ShutdownHook} instance.
+     */
+    protected ShutdownHook hook;
+
+    /**
+     * The priority of the shutdown hook.
+     */
+    protected int priority;
+
+    /**
+     * Constructs a {@code ShutdownHookAdapter} with a {@link Runnable} to be executed
+     * at the default priority and in a blocking manner.
+     *
+     * @param runnable the runnable to be executed during shutdown
+     */
+    public ShutdownHookAdapter(Runnable runnable) {
+        this(runnable, DEFAULT_PRIORITY, false);
+    }
+
+    /**
+     * Constructs a {@code ShutdownHookAdapter} with a {@link Runnable} to be executed
+     * at the specified priority and in a blocking manner.
+     *
+     * @param runnable the runnable to be executed during shutdown
+     * @param priority the priority of the shutdown hook
+     */
+    public ShutdownHookAdapter(Runnable runnable, int priority) {
+        this(runnable, priority, false);
+    }
+
+    /**
+     * Constructs a {@code ShutdownHookAdapter} with a {@link Runnable} to be executed
+     * at the specified priority, with an option for asynchronous execution.
+     *
+     * @param runnable the runnable to be executed during shutdown
+     * @param priority the priority of the shutdown hook
+     * @param async    if {@code true}, the runnable will be executed asynchronously
+     */
+    public ShutdownHookAdapter(Runnable runnable, int priority, boolean async) {
+        this.hook = !async ? () -> {
+            runnable.run();
+            return CompletableFuture.completedFuture(null);
+        } : () -> CompletableFuture.runAsync(runnable);
+        this.priority = priority;
+    }
+
+    /**
+     * Executes the wrapped shutdown logic.
+     *
+     * @return a CompletableFuture representing the pending completion of the hook
+     */
+    @Override
+    public CompletableFuture<Void> run() {
+        return hook.run();
+    }
+
+    /**
+     * Returns the priority of this shutdown hook.
+     *
+     * @return the priority of the hook
+     */
+    @Override
+    public int priority() {
+        return priority;
+    }
+}
+