Use in-memory session tracking to fix slow shutdown

applenick · applenick · commit 58b6baac456c · 2026-01-27T00:35:24.000-08:00
Signed-off-by: applenick &lt;applenick@users.noreply.github.com&gt;
diff --git a/core/src/main/java/dev/pgm/community/sessions/VanishedSessionListener.java b/core/src/main/java/dev/pgm/community/sessions/VanishedSessionListener.java
@@ -19,12 +19,10 @@ public void onVanish(PlayerVanishEvent event) {
     Player player = event.getPlayer().getBukkit();
     if (sessions.isPlayerJoining(player)) return;
 
-    sessions
-        .getLatestSession(player.getUniqueId(), false)
-        .thenAcceptAsync(
-            session -> {
-              sessions.endSession(session);
-              sessions.startSession(player);
-            });
+    Session session = sessions.getActiveSession(player.getUniqueId());
+    if (session != null) {
+      sessions.endSession(session);
+    }
+    sessions.startSession(player);
   }
 }
diff --git a/core/src/main/java/dev/pgm/community/sessions/feature/SessionFeature.java b/core/src/main/java/dev/pgm/community/sessions/feature/SessionFeature.java
@@ -10,6 +10,8 @@ public interface SessionFeature extends Feature {
 
   CompletableFuture<Session> getLatestSession(UUID playerId, boolean ignoreDisguised);
 
+  Session getActiveSession(UUID playerId);
+
   Session startSession(Player player);
 
   void endSession(Session session);
@@ -18,7 +20,5 @@ public interface SessionFeature extends Feature {
 
   void endOngoingSessions();
 
-  default void endOngoingSessionsSync() {
-    endOngoingSessions();
-  }
+  void endOngoingSessionsSync();
 }
diff --git a/core/src/main/java/dev/pgm/community/sessions/feature/types/SessionFeatureCore.java b/core/src/main/java/dev/pgm/community/sessions/feature/types/SessionFeatureCore.java
@@ -11,10 +11,10 @@
 import dev.pgm.community.utils.PGMUtils;
 import dev.pgm.community.utils.VisibilityUtils;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Logger;
 import org.bukkit.Bukkit;
 import org.bukkit.entity.Player;
@@ -29,16 +29,16 @@
 public class SessionFeatureCore extends FeatureBase implements SessionFeature {
 
   private final SessionStore store;
-  private List<UUID> joiningPlayers;
+  private final Map<UUID, Session> activeSessions;
   private VanishedSessionListener vanishedSessionListener;
   private boolean serverRestarting;
 
   public SessionFeatureCore(UsersFeature users, Logger logger, SessionStore store) {
     super(users.getConfig(), logger, "Sessions");
     this.store = store;
+    this.activeSessions = new ConcurrentHashMap<>();
 
     if (getConfig().isEnabled()) {
-      this.joiningPlayers = new ArrayList<>();
       enable();
 
       if (PGMUtils.isPGMEnabled()) {
@@ -50,6 +50,7 @@ public SessionFeatureCore(UsersFeature users, Logger logger, SessionStore store)
 
   @Override
   public void disable() {
+    super.disable();
     if (vanishedSessionListener != null) HandlerList.unregisterAll(vanishedSessionListener);
     endOngoingSessionsSync();
   }
@@ -59,9 +60,15 @@ public CompletableFuture<Session> getLatestSession(UUID playerId, boolean ignore
     return store.query(new SessionQuery(playerId, ignoreDisguised));
   }
 
+  @Override
+  public Session getActiveSession(UUID playerId) {
+    return activeSessions.get(playerId);
+  }
+
   @Override
   public Session startSession(Player player) {
     Session session = new Session(player.getUniqueId(), VisibilityUtils.isDisguised(player));
+    activeSessions.put(player.getUniqueId(), session);
     store.save(session);
 
     return session;
@@ -75,34 +82,42 @@ public void endSession(Session session) {
 
   @Override
   public void endOngoingSessions() {
-    store.endOngoingSessions();
+    Instant now = Instant.now();
+    for (Session session : activeSessions.values()) {
+      session.setEndDate(now);
+      store.updateSessionEndTime(session);
+    }
+    activeSessions.clear();
   }
 
   @Override
   public void endOngoingSessionsSync() {
-    store.endOngoingSessionsSync();
-  }
-
-  @EventHandler(priority = EventPriority.LOWEST)
-  public void onJoinLowest(PlayerJoinEvent event) {
-    joiningPlayers.add(event.getPlayer().getUniqueId());
+    Instant now = Instant.now();
+    for (Session session : activeSessions.values()) {
+      session.setEndDate(now);
+      store.updateSessionEndTimeSync(session);
+    }
+    activeSessions.clear();
   }
 
   @EventHandler(priority = EventPriority.HIGHEST)
   public void onJoinHighest(PlayerJoinEvent event) {
     if (!serverRestarting) startSession(event.getPlayer());
-    joiningPlayers.remove(event.getPlayer().getUniqueId());
   }
 
   @EventHandler(priority = EventPriority.HIGHEST)
   public void onQuitHighest(PlayerQuitEvent event) {
-    if (!serverRestarting)
-      getLatestSession(event.getPlayer().getUniqueId(), false).thenAcceptAsync(this::endSession);
+    if (serverRestarting) return;
+
+    Session session = activeSessions.remove(event.getPlayer().getUniqueId());
+    if (session != null) {
+      endSession(session);
+    }
   }
 
   @Override
   public boolean isPlayerJoining(Player player) {
-    return joiningPlayers.contains(player.getUniqueId());
+    return !activeSessions.containsKey(player.getUniqueId());
   }
 
   @EventHandler
diff --git a/core/src/main/java/dev/pgm/community/sessions/services/SQLSessionService.java b/core/src/main/java/dev/pgm/community/sessions/services/SQLSessionService.java
@@ -90,30 +90,20 @@ public void updateSessionEndTime(Session session) {
         session.getSessionId().toString());
   }
 
-  public void endOngoingSessions() {
-    DatabaseExecutor.executeUpdateAsync(
-        UPDATE_ONGOING_SESSION_ENDING_QUERY,
-        Instant.now().toEpochMilli(),
-        Community.get().getServerId());
-    DatabaseExecutor.executeUpdateAsync(
-        UPDATE_LATEST_ONGOING_SESSION_ENDING_QUERY,
-        Instant.now().toEpochMilli(),
-        Community.get().getServerId());
-  }
-
-  public void endOngoingSessionsSync() {
-    long now = Instant.now().toEpochMilli();
+  public void updateSessionEndTimeSync(Session session) {
+    Long endTime = session.getEndDate() == null ? null : session.getEndDate().toEpochMilli();
     try {
       DatabaseExecutor.executeUpdateAsync(
-              UPDATE_ONGOING_SESSION_ENDING_QUERY, now, Community.get().getServerId())
+              UPDATE_SESSION_ENDTIME_QUERY, endTime, session.getSessionId().toString())
           .join();
       DatabaseExecutor.executeUpdateAsync(
-              UPDATE_LATEST_ONGOING_SESSION_ENDING_QUERY, now, Community.get().getServerId())
+              UPDATE_LATEST_ENDTIME_QUERY, endTime, session.getSessionId().toString())
           .join();
     } catch (Exception exception) {
       Community.get()
           .getLogger()
-          .warning("Failed to end ongoing sessions during shutdown: " + exception.getMessage());
+          .warning(
+              "Failed to end session " + session.getSessionId() + ": " + exception.getMessage());
     }
   }
 
diff --git a/core/src/main/java/dev/pgm/community/sessions/services/SessionDataQuery.java b/core/src/main/java/dev/pgm/community/sessions/services/SessionDataQuery.java
@@ -12,9 +12,6 @@ public interface SessionDataQuery {
   static final String UPDATE_SESSION_ENDTIME_QUERY =
       "UPDATE " + TABLE_NAME + " SET end_time = ? WHERE id = ?";
 
-  static final String UPDATE_ONGOING_SESSION_ENDING_QUERY =
-      "UPDATE " + TABLE_NAME + " SET end_time = ? WHERE server = ? AND end_time IS NULL";
-
   static final String LATEST_TABLE_NAME = "latest_sessions";
   static final String LATEST_TABLE_FIELDS =
       "(player VARCHAR(36), ignore_disguised BOOL, session_id VARCHAR(36), disguised BOOL, server VARCHAR(32), start_time BIGINT, end_time BIGINT, PRIMARY KEY (player, ignore_disguised))";
@@ -24,7 +21,4 @@ public interface SessionDataQuery {
 
   static final String UPDATE_LATEST_ENDTIME_QUERY =
       "UPDATE " + LATEST_TABLE_NAME + " SET end_time = ? WHERE session_id = ?";
-
-  static final String UPDATE_LATEST_ONGOING_SESSION_ENDING_QUERY =
-      "UPDATE " + LATEST_TABLE_NAME + " SET end_time = ? WHERE server = ? AND end_time IS NULL";
 }
diff --git a/core/src/main/java/dev/pgm/community/sessions/store/SQLSessionStore.java b/core/src/main/java/dev/pgm/community/sessions/store/SQLSessionStore.java
@@ -29,12 +29,7 @@ public void updateSessionEndTime(Session session) {
   }
 
   @Override
-  public void endOngoingSessions() {
-    service.endOngoingSessions();
-  }
-
-  @Override
-  public void endOngoingSessionsSync() {
-    service.endOngoingSessionsSync();
+  public void updateSessionEndTimeSync(Session session) {
+    service.updateSessionEndTimeSync(session);
   }
 }
diff --git a/core/src/main/java/dev/pgm/community/sessions/store/SessionStore.java b/core/src/main/java/dev/pgm/community/sessions/store/SessionStore.java
@@ -12,7 +12,5 @@ public interface SessionStore {
 
   void updateSessionEndTime(Session session);
 
-  void endOngoingSessions();
-
-  void endOngoingSessionsSync();
+  void updateSessionEndTimeSync(Session session);
 }

Original file line number	Diff line number	Diff line change
`@@ -29,12 +29,7 @@ public void updateSessionEndTime(Session session) {`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`@Override`
`32`		`- public void endOngoingSessions() {`
`33`		`- service.endOngoingSessions();`
`34`		`- }`
`35`		`-`
`36`		`- @Override`
`37`		`- public void endOngoingSessionsSync() {`
`38`		`- service.endOngoingSessionsSync();`
	`32`	`+ public void updateSessionEndTimeSync(Session session) {`
	`33`	`+ service.updateSessionEndTimeSync(session);`
`39`	`34`	`}`
`40`	`35`	`}`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,5 @@ public interface SessionStore {`
`12`	`12`
`13`	`13`	`void updateSessionEndTime(Session session);`
`14`	`14`
`15`		`- void endOngoingSessions();`
`16`		`-`
`17`		`- void endOngoingSessionsSync();`
	`15`	`+ void updateSessionEndTimeSync(Session session);`
`18`	`16`	`}`