From beda7fa230f31c0353a11ee90edb8cb783a4a571 Mon Sep 17 00:00:00 2001
From: Liam <liam@fascinated.cc>
Date: Sun, 21 Apr 2024 18:37:21 +0100
Subject: [PATCH] fix websocket metrics

---
 .../xyz/mcutils/backend/websocket/WebSocket.java     | 12 ++++++++++--
 .../backend/websocket/impl/MetricsWebSocket.java     |  8 ++++----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/main/java/xyz/mcutils/backend/websocket/WebSocket.java b/src/main/java/xyz/mcutils/backend/websocket/WebSocket.java
index 5fd6374..0b54009 100644
--- a/src/main/java/xyz/mcutils/backend/websocket/WebSocket.java
+++ b/src/main/java/xyz/mcutils/backend/websocket/WebSocket.java
@@ -40,14 +40,22 @@ public abstract class WebSocket extends TextWebSocketHandler {
         session.sendMessage(new TextMessage(message));
     }
 
+    /**
+     * Called when a session connects to the WebSocket.
+     *
+     * @param session the session that connected
+     */
+    abstract public void onSessionConnect(WebSocketSession session);
+
     @Override
-    public void afterConnectionEstablished(@NotNull WebSocketSession session) throws Exception {
+    public final void afterConnectionEstablished(@NotNull WebSocketSession session) {
         this.sessions.add(session);
         log.info("Connection established: {}", session.getId());
+        this.onSessionConnect(session);
     }
 
     @Override
-    public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus status) throws Exception {
+    public final void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus status) {
         this.sessions.remove(session);
         log.info("Connection closed: {}", session.getId());
     }
diff --git a/src/main/java/xyz/mcutils/backend/websocket/impl/MetricsWebSocket.java b/src/main/java/xyz/mcutils/backend/websocket/impl/MetricsWebSocket.java
index 45fb885..308c754 100644
--- a/src/main/java/xyz/mcutils/backend/websocket/impl/MetricsWebSocket.java
+++ b/src/main/java/xyz/mcutils/backend/websocket/impl/MetricsWebSocket.java
@@ -32,8 +32,8 @@ public class MetricsWebSocket extends WebSocket {
     }
 
     @Override
-    public void afterConnectionEstablished(@NotNull WebSocketSession session) {
-        sendMetrics(session); // Send metrics to the client when they connect
+    public void onSessionConnect(WebSocketSession session) {
+        sendMetrics(session);
     }
 
     /**
@@ -43,11 +43,11 @@ public class MetricsWebSocket extends WebSocket {
      */
     private void sendMetrics(WebSocketSession session) {
         try {
-            session.sendMessage(new TextMessage(Main.GSON.toJson(Map.of(
+            this.sendMessage(session, Main.GSON.toJson(Map.of(
                     "totalRequests", metricService.getMetric(TotalRequestsMetric.class).getValue(),
                     "uniqueServerLookups", metricService.getMetric(UniqueServerLookupsMetric.class).getValue(),
                     "uniquePlayerLookups", metricService.getMetric(UniquePlayerLookupsMetric.class).getValue()
-            ))));
+            )));
         } catch (Exception e) {
             log.error("An error occurred while sending metrics to the client", e);
         }