package com.ard.tianyu.websocket.handler; import com.alibaba.fastjson2.JSON; import com.ard.tianyu.model.receiver.BoxStreamUrlReceiver; import com.ard.tianyu.service.ITyService; import com.ard.tianyu.websocket.model.PullStreamDTO; import com.ard.tianyu.websocket.utils.WebSocketUtil; import com.ard.vtdu.domain.ArdVtdu; import com.ard.vtdu.service.IMediaService; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Component; import org.springframework.web.socket.CloseStatus; import org.springframework.web.socket.TextMessage; import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.handler.TextWebSocketHandler; import javax.annotation.Resource; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @Component @Slf4j public class TyStreamHandler extends TextWebSocketHandler { @Resource private ITyService tyService; @Resource private IMediaService mediaService; //盒子拉流地址缓存 private static final ConcurrentMap BOX_STREAM_PULL_MAP = new ConcurrentHashMap<>(); //盒子拉流用户集合缓存 private static final ConcurrentMap> STREAM_USER_MAP = new ConcurrentHashMap<>(); @Override public void handleTextMessage(WebSocketSession session, TextMessage message) { String userId = (String) session.getAttributes().get("userId"); String payload = message.getPayload(); PullStreamDTO pullStreamDTO = JSON.parseObject(payload, PullStreamDTO.class); if (pullStreamDTO != null) { String boxSn = pullStreamDTO.getBoxSn(); if (pullStreamDTO.getPullType().equals(1)) { startPullStream(boxSn, userId); } else { stopPullStream(boxSn, userId); } } log.info("【stream】Received message: {}", payload); log.info("【stream】size: {}", WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.size()); } //开始拉流 private void startPullStream(String boxSn, String userId) { STREAM_USER_MAP.computeIfAbsent(boxSn, k -> ConcurrentHashMap.newKeySet()).add(userId); String messageReply; if (BOX_STREAM_PULL_MAP.containsKey(boxSn)) { String webrtcUrl = BOX_STREAM_PULL_MAP.get(boxSn); messageReply = webrtcUrl; log.warn("【stream】Already pulled stream: {}", webrtcUrl); } else { BoxStreamUrlReceiver boxLiveUrl = tyService.getBoxLiveUrl(boxSn); if (boxLiveUrl != null) { ArdVtdu vtdu = ArdVtdu.builder().name(boxSn).mode(0).isCode(0).rtspSource(boxLiveUrl.getRtsp()).build(); ArdVtdu ardVtdu = mediaService.addPath(vtdu); BOX_STREAM_PULL_MAP.put(boxSn, ardVtdu.getWebrtcUrl()); messageReply = ardVtdu.getWebrtcUrl(); log.warn("【stream】First pulled stream: {}", ardVtdu.getWebrtcUrl()); } else { log.warn("【stream】Box stream not found: {}", boxSn); messageReply = "【stream】Box stream not found: " + boxSn; } } WebSocketUtil.sendMessage(WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.get(userId), messageReply); } //停止拉流 private void stopPullStream(String boxSn, String userId) { synchronized (this) { Set users = STREAM_USER_MAP.get(boxSn); if (users != null) { users.remove(userId); int userCount = users.size(); String messageReply; if (userCount == 0) { closeStream(boxSn); messageReply = "【stream】Stream closed due to no active users: " + boxSn; } else { messageReply = "【stream】Stream is still in use, remaining users: " + users; } WebSocketUtil.sendMessage(WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.get(userId), messageReply); } } } //关闭拉流 private void closeStream(String boxSn) { mediaService.removePath(boxSn); BOX_STREAM_PULL_MAP.remove(boxSn); STREAM_USER_MAP.remove(boxSn); log.warn("【stream】Stream closed due to no active users: {}", boxSn); } @Override public void afterConnectionEstablished(WebSocketSession session) { String userId = (String) session.getAttributes().get("userId"); WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.put(userId, session); log.info("【stream】Connected: {}", session.getId()); log.info("【stream】size: {}", WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.size()); } @Override public void afterConnectionClosed(WebSocketSession session, CloseStatus status) { String userId = (String) session.getAttributes().get("userId"); WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.remove(userId); STREAM_USER_MAP.forEach((boxSn, users) -> stopPullStream(boxSn, userId)); log.info("【stream】Disconnected: {}", session.getId()); log.info("【stream】size: {}", WebSocketUtil.STREAM_ONLINE_USER_SESSIONS.size()); } }