LocationSessionRegistry.java
1.27 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
package com.diligrp.rider.websocket;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;
import java.util.Collections;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
@Component
public class LocationSessionRegistry {
private final ConcurrentHashMap<Long, Set<WebSocketSession>> citySessions = new ConcurrentHashMap<>();
private final ConcurrentHashMap<String, Long> sessionCities = new ConcurrentHashMap<>();
public void register(Long cityId, WebSocketSession session) {
citySessions.computeIfAbsent(cityId, key -> ConcurrentHashMap.newKeySet()).add(session);
sessionCities.put(session.getId(), cityId);
}
public void unregister(WebSocketSession session) {
Long cityId = sessionCities.remove(session.getId());
if (cityId == null) {
return;
}
Set<WebSocketSession> sessions = citySessions.get(cityId);
if (sessions == null) {
return;
}
sessions.remove(session);
if (sessions.isEmpty()) {
citySessions.remove(cityId, sessions);
}
}
public Set<WebSocketSession> getSessions(Long cityId) {
return citySessions.getOrDefault(cityId, Collections.emptySet());
}
}