close() implemented for correct componen shutdown

1 files changed, 16 insertions, 1 deletions

diff --git a/juick-server-web/src/main/java/com/juick/server/component/JuickServerComponent.java b/juick-server-web/src/main/java/com/juick/server/component/JuickServerComponent.java
index 96b8c398..71e81f2d 100644
--- a/juick-server-web/src/main/java/com/juick/server/component/JuickServerComponent.java
+++ b/juick-server-web/src/main/java/com/juick/server/component/JuickServerComponent.java
@@ -11,8 +11,11 @@ import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.handler.TextWebSocketHandler;
 
+import javax.annotation.PostConstruct;
+import javax.annotation.PreDestroy;
 import javax.inject.Inject;
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class JuickServerComponent extends TextWebSocketHandler {
     private static Logger logger = LoggerFactory.getLogger(JuickServerComponent.class);
@@ -22,7 +25,12 @@ public class JuickServerComponent extends TextWebSocketHandler {
     private ObjectMapper jsonMapper;
 
     private WebSocketSession session;
+    private final AtomicBoolean closeFlag = new AtomicBoolean(false);
 
+    @PostConstruct
+    public void init() {
+        closeFlag.set(false);
+    }
     @Override
     public void afterConnectionEstablished(WebSocketSession session) throws Exception {
         logger.info("WebSocket connected");
@@ -41,7 +49,9 @@ public class JuickServerComponent extends TextWebSocketHandler {
 
         if (logger.isInfoEnabled()) // prevent writeValueAsString execution if logger disabled
             logger.info("got jmsg: {}", jsonMapper.writeValueAsString(jmsg));
-        applicationEventPublisher.publishEvent(new MessageEvent(this, jmsg));
+        if (!closeFlag.get()) {
+            applicationEventPublisher.publishEvent(new MessageEvent(this, jmsg));
+        }
     }
 
     @Scheduled(fixedRate = 30000, initialDelay = 30000)
@@ -53,4 +63,9 @@ public class JuickServerComponent extends TextWebSocketHandler {
             applicationEventPublisher.publishEvent(new DisconnectedEvent(this));
         }
     }
+    @PreDestroy
+    public void close() {
+        closeFlag.set(true);
+    }
+
 }