Correct node acknowledgement

SymphonyDatabaseNode/src/main/java/io/github/lumijiez/Main.java

@@ -1,36 +1,76 @@
 package io.github.lumijiez;
 
+import io.github.lumijiez.data.Data;
 import jakarta.persistence.EntityManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.LogManager;
 
-import java.io.IOException;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.net.http.HttpClient;
-import java.net.http.HttpRequest;
-import java.net.http.HttpResponse;
+import java.net.http.WebSocket;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.CountDownLatch;
 
 public class Main {
     public static Logger logger = LogManager.getLogger(Main.class);
+    private static final CountDownLatch waitForConnection = new CountDownLatch(1);
 
-    public static void main(String[] args) throws URISyntaxException, IOException, InterruptedException {
-        HttpClient client = HttpClient.newHttpClient();
-
+    public static void main(String[] args) {
+        try {
             logger.info("Node up.");
             EntityManager em = Data.getEntityManager();
 
             logger.info("Connected to database: << symphony >>");
             em.close();
 
-        HttpRequest request = HttpRequest.newBuilder()
-                .uri(new URI("http://symphony-discovery:8083/check"))
-                .build();
+            try (HttpClient client = HttpClient.newHttpClient()) {
+                CompletableFuture<WebSocket> wsFuture = client
+                        .newWebSocketBuilder()
+                        .buildAsync(new URI("ws://symphony-discovery:8083/discovery"), new WebSocket.Listener() {
+                            @Override
+                            public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) {
+                                try {
+                                    int nodeCount = Integer.parseInt(data.toString());
+                                    logger.info("Acknowledged system nodes: {}", nodeCount);
+                                } catch (NumberFormatException e) {
+                                    logger.error("Received invalid node count: {}", data);
+                                }
+                                return WebSocket.Listener.super.onText(webSocket, data, last);
+                            }
 
-        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
+                            @Override
+                            public void onOpen(WebSocket webSocket) {
+                                logger.info("Successfully registered to Discovery");
+                                waitForConnection.countDown();
+                                WebSocket.Listener.super.onOpen(webSocket);
+                            }
 
-        if (response.statusCode() == 200) {
-            logger.info("Node successfully registered to Discovery");
+                            @Override
+                            public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
+                                logger.info("Unregistered from Discovery: {}", reason);
+                                return WebSocket.Listener.super.onClose(webSocket, statusCode, reason);
+                            }
+
+                            @Override
+                            public void onError(WebSocket webSocket, Throwable error) {
+                                logger.error("Error: {}", error.getMessage());
+                                WebSocket.Listener.super.onError(webSocket, error);
+                            }
+                        });
+
+                WebSocket ws = wsFuture.join();
+
+                try {
+                    waitForConnection.await();
+                    Thread.currentThread().join();
+                } finally {
+                    ws.sendClose(WebSocket.NORMAL_CLOSURE, "Node shutting down").join();
+                }
+            }
+
+        } catch (Exception e) {
+            logger.error("Error in main: {}", e.getMessage(), e);
         }
     }
 }

SymphonyDatabaseNode/src/main/java/io/github/lumijiez/data/Data.java

@@ -1,4 +1,4 @@
-package io.github.lumijiez;
+package io.github.lumijiez.data;
 
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.EntityManagerFactory;

SymphonyDatabaseNode/src/main/java/io/github/lumijiez/data/entities/User.java

@@ -1,4 +1,4 @@
-package io.github.lumijiez;
+package io.github.lumijiez.data.entities;
 
 import jakarta.persistence.Entity;
 import jakarta.persistence.Id;

SymphonyDatabaseNode/src/main/resources/META-INF/persistence.xml

@@ -7,7 +7,7 @@
     <persistence-unit name="mainUnit" transaction-type="RESOURCE_LOCAL">
         <provider>org.hibernate.jpa.HibernatePersistenceProvider</provider>
 
-        <class>io.github.lumijiez.User</class>
+        <class>io.github.lumijiez.data.entities.User</class>
 
         <exclude-unlisted-classes>false</exclude-unlisted-classes>
 

SymphonyDiscovery/src/main/java/io/github/lumijiez/JavalinConfig.java

@@ -2,11 +2,34 @@ package io.github.lumijiez;
 
 import com.google.gson.Gson;
 import io.javalin.Javalin;
+import io.javalin.json.JavalinGson;
+import io.javalin.websocket.WsContext;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 public class JavalinConfig {
-    public static Gson gson = new Gson();
+    private static final Map<String, WsContext> users = new ConcurrentHashMap<>();
 
     public static void setup(Javalin app) {
         app.get("/check", ctx -> ctx.result("OK"));
+
+        app.ws("/discovery", ws -> {
+           ws.onConnect(ctx -> {
+               String id = ctx.sessionId();
+               users.put(id, ctx);
+               broadcast("Discovery", "Join");
+           });
+
+           ws.onClose(ctx -> {
+               String id = ctx.sessionId();
+               users.remove(id);
+               broadcast("Discovery", "Leave");
+           });
+        });
+    }
+
+    private static void broadcast(String sender, String message) {
+        users.values().forEach(ctx -> ctx.send(users.size()));
     }
 }

SymphonyDiscovery/src/main/java/io/github/lumijiez/Main.java

@@ -1,10 +1,14 @@
 package io.github.lumijiez;
 
 import io.javalin.Javalin;
+import io.javalin.json.JavalinGson;
 
 public class Main {
     public static void main(String[] args) {
-        Javalin app = Javalin.create().start(8083);
+        Javalin app = Javalin.create(config -> {
+            config.jsonMapper(new JavalinGson());
+        }).start(8083);
+
         JavalinConfig.setup(app);
 
         System.out.print("Discovery service up and running");