SSE cleanup refactoring

1 files changed, 9 insertions, 4 deletions

diff --git a/juick-server/src/main/java/com/juick/server/api/Service.java b/juick-server/src/main/java/com/juick/server/api/Service.java
index c65d89a1..2938e9da 100644
--- a/juick-server/src/main/java/com/juick/server/api/Service.java
+++ b/juick-server/src/main/java/com/juick/server/api/Service.java
@@ -39,7 +39,7 @@ import java.util.*;
 
 @Controller
 public class Service {
-    private static Logger logger = LoggerFactory.getLogger(Post.class);
+    private static Logger logger = LoggerFactory.getLogger(Service.class);
     @Inject
     private UserService userService;
     @Inject
@@ -57,7 +57,7 @@ public class Service {
     @Inject
     private ServerManager serverManager;
 
-    Session session = Session.getDefaultInstance(new Properties());
+    private Session session = Session.getDefaultInstance(new Properties());
 
     @ApiIgnore
     @PostMapping("/api/mail")
@@ -137,6 +137,11 @@ public class Service {
             throw new HttpForbiddenException();
         }
     }
+    private void endSession(SseEmitter emitter) {
+        serverManager.getSessions().stream()
+                .filter(s -> s.getEmitter().equals(emitter))
+                .forEach(session -> serverManager.getSessions().remove(session));
+    }
     @GetMapping("/api/events")
     public SseEmitter handle() throws IOException {
         logger.info("{} connected", UserUtils.getCurrentUser().getName());
@@ -146,8 +151,8 @@ public class Service {
         emitter.send(service);
         serverManager.getSessions().add(new ServerManager.EventSession(UserUtils.getCurrentUser(), emitter));
 
-        emitter.onCompletion(() -> serverManager.getSessions().remove(emitter));
-        emitter.onTimeout(() -> serverManager.getSessions().remove(emitter));
+        emitter.onCompletion(() -> endSession(emitter));
+        emitter.onTimeout(() -> endSession(emitter));
 
         return emitter;
     }