|
@@ -1,5 +1,8 @@
|
|
|
package org.hswebframework.web.socket.message;
|
|
|
|
|
|
+import org.hswebframework.web.concurrent.counter.Counter;
|
|
|
+import org.hswebframework.web.concurrent.counter.CounterManager;
|
|
|
+import org.hswebframework.web.concurrent.counter.SimpleCounterManager;
|
|
|
import org.hswebframework.web.message.MessageSubscribe;
|
|
|
import org.hswebframework.web.message.Messager;
|
|
|
import org.hswebframework.web.message.support.ObjectMessage;
|
|
@@ -7,8 +10,11 @@ import org.springframework.web.socket.TextMessage;
|
|
|
import org.springframework.web.socket.WebSocketSession;
|
|
|
|
|
|
import java.io.IOException;
|
|
|
+import java.util.Collection;
|
|
|
import java.util.Map;
|
|
|
+import java.util.Objects;
|
|
|
import java.util.concurrent.ConcurrentHashMap;
|
|
|
+import java.util.stream.Collectors;
|
|
|
|
|
|
import static org.hswebframework.web.message.builder.StaticMessageBuilder.object;
|
|
|
import static org.hswebframework.web.message.builder.StaticMessageSubjectBuilder.*;
|
|
@@ -23,61 +29,96 @@ public class DefaultWebSocketMessager implements WebSocketMessager {
|
|
|
private Messager messager;
|
|
|
|
|
|
public DefaultWebSocketMessager(Messager messager) {
|
|
|
+ this(messager, new SimpleCounterManager());
|
|
|
+ }
|
|
|
+
|
|
|
+ public DefaultWebSocketMessager(Messager messager, CounterManager counterManager) {
|
|
|
this.messager = messager;
|
|
|
+ this.counterManager = counterManager == null ? new SimpleCounterManager() : counterManager;
|
|
|
}
|
|
|
|
|
|
- // command, userId, sessionId
|
|
|
+ // command, type, sessionId
|
|
|
private final Map<String, Map<String, Map<String, MessageSubscribeSession>>> store = new ConcurrentHashMap<>(32);
|
|
|
|
|
|
+ private CounterManager counterManager = new SimpleCounterManager();
|
|
|
+
|
|
|
+
|
|
|
@Override
|
|
|
public void onSessionConnect(WebSocketSession session) {
|
|
|
|
|
|
}
|
|
|
|
|
|
+ private String getSubTotalKey(String command, String type) {
|
|
|
+ return "sub_".concat(command)
|
|
|
+ .concat("_")
|
|
|
+ .concat(type)
|
|
|
+ .concat("_total");
|
|
|
+ }
|
|
|
+
|
|
|
@Override
|
|
|
- public void onSessionClose(WebSocketSession session) {
|
|
|
+ public int getSubscribeTotal(String command, String type) {
|
|
|
+ return (int) counterManager.getCounter(getSubTotalKey(command, type)).get();
|
|
|
+ }
|
|
|
|
|
|
+ @Override
|
|
|
+ public void onSessionClose(WebSocketSession session) {
|
|
|
+ store.values() //command
|
|
|
+ .stream().map(Map::values).flatMap(Collection::stream)
|
|
|
+ .map(sessionStore -> sessionStore.get(session.getId()))
|
|
|
+ .filter(Objects::nonNull)
|
|
|
+ .forEach(MessageSubscribeSession::cancel);
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
- public void publish(String toUser, WebSocketMessage message) {
|
|
|
+ public void publish(String command, String type, WebSocketMessage message) {
|
|
|
messager.publish(object(message))
|
|
|
- .to(user(toUser))
|
|
|
+ .to(TYPE_QUEUE.equals(type) ? queue("queue_" + command) : topic("topic_" + command))
|
|
|
.send();
|
|
|
}
|
|
|
|
|
|
- private Map<String, MessageSubscribeSession> getSubSession(String command, String userId) {
|
|
|
+ private Map<String, MessageSubscribeSession> getSubSession(String command, String type) {
|
|
|
return store.computeIfAbsent(command, cmd -> new ConcurrentHashMap<>(128))
|
|
|
- .computeIfAbsent(userId, uid -> new ConcurrentHashMap<>());
|
|
|
+ .computeIfAbsent(type, t -> new ConcurrentHashMap<>());
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
- public boolean subscribe(String command, String userId, WebSocketSession socketSession) {
|
|
|
- Map<String, MessageSubscribeSession> subscribeSessionStore = getSubSession(command, userId);
|
|
|
+ public boolean subscribe(String command, String type, WebSocketSession socketSession) {
|
|
|
+ Map<String, MessageSubscribeSession> subscribeSessionStore = getSubSession(command, type);
|
|
|
subscribeSessionStore.computeIfAbsent(socketSession.getId(), sessionId -> {
|
|
|
- MessageSubscribe<ObjectMessage<WebSocketMessage>> subscribe = messager.subscribe(user(userId));
|
|
|
+ MessageSubscribe<ObjectMessage<WebSocketMessage>> subscribe = messager
|
|
|
+ .subscribe(TYPE_QUEUE.equals(type) ? queue("queue_" + command) : topic("topic_" + command));
|
|
|
subscribe.onMessage(message -> {
|
|
|
try {
|
|
|
if (!socketSession.isOpen()) {
|
|
|
- deSubscribe(command, userId, socketSession);
|
|
|
+ deSubscribe(command, type, socketSession);
|
|
|
+ return;
|
|
|
}
|
|
|
socketSession.sendMessage(new TextMessage(((ObjectMessage) message).getObject().toString()));
|
|
|
} catch (IOException e) {
|
|
|
e.printStackTrace();
|
|
|
}
|
|
|
});
|
|
|
- return new MessageSubscribeSession(subscribe, socketSession);
|
|
|
+ return new MessageSubscribeSession(subscribe, socketSession) {
|
|
|
+ @Override
|
|
|
+ public void cancel() {
|
|
|
+ super.cancel();
|
|
|
+ counterManager.getCounter(getSubTotalKey(command, type)).decrement();
|
|
|
+ }
|
|
|
+ };
|
|
|
});
|
|
|
- return false;
|
|
|
+ counterManager.getCounter(getSubTotalKey(command, type)).increment();
|
|
|
+ return true;
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
- public boolean deSubscribe(String command, String userId, WebSocketSession socketSession) {
|
|
|
- Map<String, MessageSubscribeSession> subscribeSessionStore = getSubSession(command, userId);
|
|
|
+ public boolean deSubscribe(String command, String type, WebSocketSession socketSession) {
|
|
|
+ Map<String, MessageSubscribeSession> subscribeSessionStore = getSubSession(command, type);
|
|
|
MessageSubscribeSession subscribeSession = subscribeSessionStore.get(socketSession.getId());
|
|
|
if (null != subscribeSession) {
|
|
|
subscribeSession.getSubscribe().cancel();
|
|
|
subscribeSessionStore.remove(socketSession.getId());
|
|
|
+ counterManager.getCounter(getSubTotalKey(command, type)).decrement();
|
|
|
+ return true;
|
|
|
}
|
|
|
return false;
|
|
|
}
|
|
@@ -107,5 +148,9 @@ public class DefaultWebSocketMessager implements WebSocketMessager {
|
|
|
public void setSession(WebSocketSession session) {
|
|
|
this.session = session;
|
|
|
}
|
|
|
+
|
|
|
+ public void cancel() {
|
|
|
+ subscribe.cancel();
|
|
|
+ }
|
|
|
}
|
|
|
}
|