Merge pull request #506 from danthe1st/graceful-redeploy

graceful shutdown on redeploy

Author: danthe1st

src/main/java/net/discordjug/javabot/SpringConfig.java

@@ -5,8 +5,6 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 
-import javax.sql.DataSource;
-
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -19,7 +17,6 @@
 import net.discordjug.javabot.annotations.PreRegisteredListener;
 import net.discordjug.javabot.data.config.BotConfig;
 import net.discordjug.javabot.data.config.SystemsConfig;
-import net.discordjug.javabot.data.h2db.DbHelper;
 import net.discordjug.javabot.tasks.PresenceUpdater;
 import net.dv8tion.jda.api.JDA;
 import net.dv8tion.jda.api.JDABuilder;
@@ -36,16 +33,8 @@
 @RequiredArgsConstructor
 public class SpringConfig {
 	@Bean
-	PresenceUpdater standardActivityPresenceUpdater() {
-		return PresenceUpdater.standardActivities();
-	}
-
-	@Bean
-	DataSource dataSource(BotConfig config) {
-		if (config.getSystems().getJdaBotToken().isEmpty()) {
-			throw new RuntimeException("JDA Token not set. Stopping Bot...");
-		}
-		return DbHelper.initDataSource(config);
+	PresenceUpdater standardActivityPresenceUpdater(ScheduledExecutorService threadPool) {
+		return PresenceUpdater.standardActivities(threadPool);
 	}
 
 	@Bean
@@ -95,6 +84,10 @@ DIH4JDA initializeDIH4JDA(JDA jda) throws DIH4JDAException {
 
 	@Bean
 	BotConfig botConfig() {
-		return new BotConfig(Path.of("config"));
+		BotConfig botConfig = new BotConfig(Path.of("config"));
+		if (botConfig.getSystems().getJdaBotToken().isEmpty()) {
+			throw new RuntimeException("JDA Token not set. Stopping Bot...");
+		}
+		return botConfig;
 	}
 }

src/main/java/net/discordjug/javabot/data/h2db/DbHelper.java

@@ -2,8 +2,7 @@
 
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
-
-import lombok.Getter;
+import jakarta.annotation.PreDestroy;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import net.discordjug.javabot.data.config.BotConfig;
@@ -12,6 +11,7 @@
 
 import org.h2.tools.Server;
 import org.jetbrains.annotations.NotNull;
+import org.springframework.context.annotation.Bean;
 import org.springframework.stereotype.Service;
 
 import java.io.IOException;
@@ -32,12 +32,12 @@
 /**
  * Class that provides helper methods for dealing with the database.
  */
-@Getter
 @Slf4j
 @Service
 @RequiredArgsConstructor
 public class DbHelper {
-	private final DataSource dataSource;
+	private Server server;
+	private HikariDataSource ds;
 
 	/**
 	 * Initializes the data source that'll be used throughout the bot to access
@@ -48,12 +48,12 @@ public class DbHelper {
 	 * @throws IllegalStateException If an error occurs and we're unable to
 	 *                               start the database.
 	 */
-	public static @NotNull HikariDataSource initDataSource(@NotNull BotConfig config) {
+	@Bean
+	DataSource initDataSource(@NotNull BotConfig config) {
 		// Determine if we need to initialize the schema, before starting up the server.
 		boolean shouldInitSchema = shouldInitSchema(config.getSystems().getHikariConfig().getJdbcUrl());
 
 		// Now that we have remembered whether we need to initialize the schema, start up the server.
-		Server server;
 		try {
 			System.setProperty("h2.bindAddress", "127.0.0.1");
 			server = Server.createTcpServer("-tcpPort", "9122", "-ifNotExists").start();
@@ -66,12 +66,8 @@ public class DbHelper {
 		hikariConfig.setJdbcUrl(hikariConfigSource.getJdbcUrl());
 		hikariConfig.setMaximumPoolSize(hikariConfigSource.getMaximumPoolSize());
 		hikariConfig.setLeakDetectionThreshold(hikariConfigSource.getLeakDetectionThreshold());
-		HikariDataSource ds = new HikariDataSource(hikariConfig);
+		ds = new HikariDataSource(hikariConfig);
 		// Add a shutdown hook to close down the datasource and server when the JVM terminates.
-		Runtime.getRuntime().addShutdownHook(new Thread(() -> {
-			ds.close();
-			server.stop();
-		}));
 		if (shouldInitSchema) {
 			try {
 				initializeSchema(ds);
@@ -124,4 +120,17 @@ public static void initializeSchema(DataSource dataSource) throws IOException, S
 			log.info("Successfully initialized H2 database.");
 		}
 	}
+	
+	/**
+	 * Closes the database (server) when the application is shut down.
+	 */
+	@PreDestroy
+	public void stop() {
+		if (ds != null) {
+			ds.close();
+		}
+		if (server != null) {
+			server.stop();
+		}
+	}
 }

src/main/java/net/discordjug/javabot/systems/staff_commands/RedeployCommand.java

@@ -1,6 +1,10 @@
 package net.discordjug.javabot.systems.staff_commands;
 
 import xyz.dynxsty.dih4jda.interactions.commands.application.SlashCommand;
+
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
 import lombok.extern.slf4j.Slf4j;
 import net.discordjug.javabot.data.config.BotConfig;
 import net.discordjug.javabot.data.h2db.message_cache.MessageCache;
@@ -12,6 +16,8 @@
 import net.dv8tion.jda.api.interactions.commands.build.Commands;
 
 import org.jetbrains.annotations.NotNull;
+import org.springframework.boot.SpringApplication;
+import org.springframework.context.ConfigurableApplicationContext;
 
 /**
  * <h3>This class represents the /redeploy command.</h3>
@@ -26,15 +32,21 @@
 public class RedeployCommand extends SlashCommand {
 	private final MessageCache messageCache;
 	private final BotConfig botConfig;
+	private final ScheduledExecutorService asyncPool;
+	private final ConfigurableApplicationContext applicationContext;
 
 	/**
 	 * The constructor of this class, which sets the corresponding {@link net.dv8tion.jda.api.interactions.commands.build.SlashCommandData}.
 	 * @param messageCache A service managing recent messages
 	 * @param botConfig The main configuration of the bot
+	 * @param asyncPool The thread pool used for asynchronous operations
+	 * @param applicationContext The spring application context
 	 */
-	public RedeployCommand(MessageCache messageCache, BotConfig botConfig) {
+	public RedeployCommand(MessageCache messageCache, BotConfig botConfig, ScheduledExecutorService asyncPool, ConfigurableApplicationContext applicationContext) {
 		this.messageCache = messageCache;
 		this.botConfig=botConfig;
+		this.asyncPool = asyncPool;
+		this.applicationContext = applicationContext;
 		setCommandData(Commands.slash("redeploy", "(ADMIN-ONLY) Makes the bot redeploy.")
 				.setDefaultPermissions(DefaultMemberPermissions.DISABLED)
 				.setGuildOnly(true)
@@ -51,6 +63,14 @@ public void execute(@NotNull SlashCommandInteractionEvent event) {
 		log.warn("Redeploying... Requested by: " + UserUtils.getUserTag(event.getUser()));
 		event.reply("**Redeploying...** This may take some time.").queue();
 		messageCache.synchronize();
-		System.exit(0);
+		asyncPool.shutdownNow();
+		try {
+			asyncPool.awaitTermination(3, TimeUnit.SECONDS);
+			event.getJDA().shutdown();
+			event.getJDA().awaitShutdown(3, TimeUnit.SECONDS);
+		} catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+		}
+		SpringApplication.exit(applicationContext, () -> 0);
 	}
 }

src/main/java/net/discordjug/javabot/tasks/PresenceUpdater.java

@@ -12,7 +12,6 @@
 
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -32,7 +31,7 @@ public class PresenceUpdater extends ListenerAdapter {
 	 * The executor that is responsible for the scheduled updates of the bot's
 	 * presence data.
 	 */
-	private final ScheduledExecutorService threadPool = Executors.newSingleThreadScheduledExecutor();
+	private final ScheduledExecutorService threadPool;
 
 	/**
 	 * A list of functions that take a reference to the bot's JDA client, and
@@ -64,24 +63,27 @@ public class PresenceUpdater extends ListenerAdapter {
 	 * @param activities The list of activity-producing functions.
 	 * @param delay      The amount of time the updater should wait before updating the activity.
 	 * @param delayUnit  The unit of time that {@link PresenceUpdater#delay} is counted in.
+	 * @param threadPool The thread pool to use for updating presences
 	 */
-	public PresenceUpdater(List<Function<JDA, Activity>> activities, long delay, TimeUnit delayUnit) {
+	public PresenceUpdater(List<Function<JDA, Activity>> activities, long delay, TimeUnit delayUnit, ScheduledExecutorService threadPool) {
 		this.activities = new CopyOnWriteArrayList<>(activities);
 		this.delay = delay;
 		this.delayUnit = delayUnit;
+		this.threadPool = threadPool;
 	}
 
 	/**
 	 * A list of standard Activities.
 	 *
+	 * @param threadPool The thread pool to use for updating activities
 	 * @return A pre-built implementation of the {@link PresenceUpdater} that
 	 * has all the necessary properties defined to reasonable defaults.
 	 */
-	public static PresenceUpdater standardActivities() {
+	public static PresenceUpdater standardActivities(ScheduledExecutorService threadPool) {
 		return new PresenceUpdater(List.of(
 				jda -> Activity.watching(String.format("%s members", jda.getGuilds().stream().mapToLong(Guild::getMemberCount).sum())),
 				jda -> Activity.customStatus("Use /report, 'Report User' or 'Report Message' to report disruptive behaviour!")
-		), 35, TimeUnit.SECONDS);
+		), 35, TimeUnit.SECONDS, threadPool);
 	}
 
 	/**