Improve SSEClient and SSEPushManager.

m0ver · m0ver · commit bf337c583a62 · 2025-05-23T21:12:23.000+08:00
diff --git a/src/main/java/org/tinystruct/http/SSEClient.java b/src/main/java/org/tinystruct/http/SSEClient.java
@@ -1,59 +1,67 @@
 package org.tinystruct.http;
 
-import org.tinystruct.ApplicationException;
+import org.tinystruct.data.component.Builder;
 
 import java.nio.charset.StandardCharsets;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 public class SSEClient implements Runnable {
     private final Response out;
-    private final LinkedBlockingQueue<String> messageQueue = new LinkedBlockingQueue<>();
+    private final BlockingQueue<Builder> messageQueue;
     private volatile boolean active = true;
+    private static final Logger logger = Logger.getLogger(SSEClient.class.getName());
 
     public SSEClient(Response out) {
+        this(out, new LinkedBlockingQueue<>());
+    }
+
+    public SSEClient(Response out, BlockingQueue<Builder> messageQueue) {
         this.out = out;
+        this.messageQueue = messageQueue;
+        // Set SSE headers
+        if (out instanceof ResponseBuilder) {
+            ResponseBuilder builder = (ResponseBuilder) out;
+            builder.setContentType("text/event-stream");
+            builder.addHeader("Cache-Control", "no-cache");
+            builder.addHeader("Connection", "keep-alive");
+        }
     }
 
-    public void send(String message) {
+    public void send(Builder message) {
         messageQueue.offer(message);
     }
 
     public void close() {
         this.active = false;
-        // 清空消息队列，防止阻塞
         messageQueue.clear();
     }
 
     @Override
     public void run() {
         try {
             while (active) {
-                String message = messageQueue.poll();
-                try {
-                    if (message != null) {
-                        writeEvent("message", message);
-                    } else {
-                        writeEvent("heartbeat", "{\"time\": " + System.currentTimeMillis() + "}");
-                    }
-                } catch (ApplicationException e) {
-                    // 写入失败，通常是客户端断开或流已关闭
-                    this.active = false;
-                    break;
+                Builder message = messageQueue.take();
+                if (message != null) {
+                    String event = formatSSEMessage(message);
+                    out.writeAndFlush(event.getBytes(StandardCharsets.UTF_8));
                 }
-                Thread.sleep(2000); // 每2秒一次
             }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
-            this.active = false;
         } catch (Exception e) {
-            this.active = false;
+            logger.log(Level.SEVERE, "Error in SSE client: " + e.getMessage(), e);
+        } finally {
+            close();
         }
     }
 
-    private void writeEvent(String event, String data) throws ApplicationException {
-        String payload = "event: " + event + "\n" +
-                "data: " + data.replace("\n", "\ndata: ") + "\n\n";
-        out.writeAndFlush(payload.getBytes(StandardCharsets.UTF_8));
+    private String formatSSEMessage(Builder message) {
+        StringBuilder sb = new StringBuilder();
+        sb.append("data: ").append(message.toString()).append("\n\n");
+        return sb.toString();
     }
 
     public boolean isActive() {
diff --git a/src/main/java/org/tinystruct/http/SSEPushManager.java b/src/main/java/org/tinystruct/http/SSEPushManager.java
@@ -1,13 +1,21 @@
 package org.tinystruct.http;
 
+import org.tinystruct.data.component.Builder;
+
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.BlockingQueue;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 public class SSEPushManager {
+    private static final Logger logger = Logger.getLogger(SSEPushManager.class.getName());
     private final ConcurrentHashMap<String, SSEClient> clients = new ConcurrentHashMap<>();
     private final ExecutorService executor = Executors.newCachedThreadPool();
+    private final AtomicBoolean isShutdown = new AtomicBoolean(false);
 
     private static final SSEPushManager instance = new SSEPushManager();
 
@@ -18,24 +26,59 @@ public static SSEPushManager getInstance() {
     }
 
     public void register(String sessionId, Response out) {
-        SSEClient client = new SSEClient(out);
+        register(sessionId, out, null);
+    }
+
+    public void register(String sessionId, Response out, BlockingQueue<Builder> messageQueue) {
+        if (isShutdown.get()) {
+            logger.log(Level.WARNING, "SSEPushManager is shutting down, cannot register new client");
+            return;
+        }
+
+        SSEClient oldClient = clients.get(sessionId);
+        if (oldClient != null) {
+            oldClient.close();
+        }
+
+        SSEClient client = messageQueue != null ? 
+            new SSEClient(out, messageQueue) : 
+            new SSEClient(out);
+            
         clients.put(sessionId, client);
         executor.submit(client);
     }
 
-    public void push(String sessionId, String message) {
+    public void push(String sessionId, Builder message) {
+        if (isShutdown.get()) {
+            logger.log(Level.WARNING, "SSEPushManager is shutting down, cannot push message");
+            return;
+        }
+
         SSEClient client = clients.get(sessionId);
-        if (client != null && client.isActive()) {
-            client.send(message);
+        if (client != null) {
+            if (client.isActive()) {
+                client.send(message);
+            } else {
+                // Clean up inactive client
+                clients.remove(sessionId, client);
+            }
         }
     }
 
-    public void broadcast(String message) {
-        for (SSEClient client : clients.values()) {
+    public void broadcast(Builder message) {
+        if (isShutdown.get()) {
+            logger.log(Level.WARNING, "SSEPushManager is shutting down, cannot broadcast message");
+            return;
+        }
+
+        clients.forEach((sessionId, client) -> {
             if (client.isActive()) {
                 client.send(message);
+            } else {
+                // Clean up inactive client
+                clients.remove(sessionId, client);
             }
-        }
+        });
     }
 
     public void remove(String sessionId) {
@@ -48,4 +91,15 @@ public void remove(String sessionId) {
     public Set<String> getClientIds() {
         return clients.keySet();
     }
+
+    public void shutdown() {
+        if (isShutdown.compareAndSet(false, true)) {
+            // Close all clients
+            clients.forEach((sessionId, client) -> client.close());
+            clients.clear();
+            
+            // Shutdown executor
+            executor.shutdown();
+        }
+    }
 }
