auction/bid/src/main/java/cn/hobbystocks/auc/websocket/LotWebSocketHandler.java

package cn.hobbystocks.auc.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.net.URI;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Component
@Slf4j
public class LotWebSocketHandler extends TextWebSocketHandler {

    private static final Pattern LOT_PATH_PATTERN = Pattern.compile(".*/(?:bid/)?ws/lot/(\\d+)$");

    private final LotWebSocketSessionRegistry sessionRegistry;
    private final ConcurrentMap<String, Long> sessionLots = new ConcurrentHashMap<>();

    @Autowired
    public LotWebSocketHandler(LotWebSocketSessionRegistry sessionRegistry) {
        this.sessionRegistry = sessionRegistry;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long lotId = resolveLotId(session.getUri());
        if (lotId == null) {
            session.close(CloseStatus.BAD_DATA);
            return;
        }
        sessionLots.put(session.getId(), lotId);
        sessionRegistry.register(lotId, session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Long lotId = sessionLots.remove(session.getId());
        if (lotId != null) {
            sessionRegistry.unregister(lotId, session);
        }
    }

    static Long resolveLotId(URI uri) {
        if (uri == null) {
            return null;
        }
        Matcher matcher = LOT_PATH_PATTERN.matcher(uri.getPath());
        if (!matcher.matches()) {
            return null;
        }
        return Long.valueOf(matcher.group(1));
    }
}