|
|
@@ -0,0 +1,184 @@
|
|
|
+package com.yhlxj.web.wss;
|
|
|
+
|
|
|
+import com.alibaba.fastjson.JSON;
|
|
|
+import com.alibaba.fastjson.JSONObject;
|
|
|
+import com.baomidou.mybatisplus.core.toolkit.Wrappers;
|
|
|
+import com.cyksj.common.exception.BusinessRuntimeException;
|
|
|
+import com.cyksj.common.task.GlobalThreadPoolTaskExecutor;
|
|
|
+import com.cyksj.common.util.QueryWrapperUtils;
|
|
|
+import com.cyksj.common.util.SpringCtxUtils;
|
|
|
+import com.yhlxj.dao.mapper.midjourney.MidjourneyUserConversationMapper;
|
|
|
+import com.yhlxj.dao.mapper.midjourney.MidjourneyUserMapper;
|
|
|
+import com.yhlxj.dao.model.entity.GroupsRelation;
|
|
|
+import com.yhlxj.dao.model.entity.MidjourneyUser;
|
|
|
+import com.yhlxj.dao.model.entity.MidjourneyUserConversation;
|
|
|
+import com.yhlxj.dao.model.enums.WsMessageTypeEnum;
|
|
|
+import com.yhlxj.redis.RedisService;
|
|
|
+import com.yhlxj.service.midjourney.MidjourneyService;
|
|
|
+import lombok.RequiredArgsConstructor;
|
|
|
+import org.apache.commons.lang3.StringUtils;
|
|
|
+import org.apache.tomcat.websocket.WsSession;
|
|
|
+import org.slf4j.Logger;
|
|
|
+import org.slf4j.LoggerFactory;
|
|
|
+import org.springframework.stereotype.Component;
|
|
|
+
|
|
|
+import javax.websocket.*;
|
|
|
+import javax.websocket.server.PathParam;
|
|
|
+import javax.websocket.server.ServerEndpoint;
|
|
|
+import java.io.IOException;
|
|
|
+
|
|
|
+/**
|
|
|
+ * @author chan
|
|
|
+ * @date 2024/6/28 13:59
|
|
|
+ */
|
|
|
+@Component
|
|
|
+@ServerEndpoint("/ws/midjourney/{userToken}/{taskId}")
|
|
|
+@RequiredArgsConstructor
|
|
|
+public class MidjourneyServerEndpoint {
|
|
|
+
|
|
|
+ private static Logger logger = LoggerFactory.getLogger(MidjourneyServerEndpoint.class);
|
|
|
+
|
|
|
+ private static final GlobalThreadPoolTaskExecutor TASK_EXECUTOR = GlobalThreadPoolTaskExecutor.getInstance();
|
|
|
+
|
|
|
+ private final MidjourneyService midjourneyService;
|
|
|
+
|
|
|
+ private final MidjourneyUserConversationMapper midjourneyUserConversationMapper;
|
|
|
+
|
|
|
+ private final MidjourneyUserMapper midjourneyUserMapper;
|
|
|
+
|
|
|
+ private final RedisService redisService;
|
|
|
+
|
|
|
+
|
|
|
+ /**
|
|
|
+ * 连接建立回调
|
|
|
+ *
|
|
|
+ * @param session
|
|
|
+ */
|
|
|
+ @OnOpen
|
|
|
+ public void onOpen(Session session, @PathParam("userToken") String userToken, @PathParam("taskId") String taskId) throws IOException {
|
|
|
+ logger.info("[WS建立连接],token:{}, taskId:{}", userToken, taskId);
|
|
|
+
|
|
|
+ //检查uniodId是否合法
|
|
|
+ if (StringUtils.isBlank(userToken) || StringUtils.isBlank(taskId)) {
|
|
|
+
|
|
|
+ logger.error("[WS建立连接失败],[参数非法],{},{}", userToken, taskId);
|
|
|
+
|
|
|
+ session.getBasicRemote().sendText("入参错误");
|
|
|
+ session.close();
|
|
|
+ }
|
|
|
+ MidjourneyUser user = getUser(userToken);
|
|
|
+ if (user == null) {
|
|
|
+ logger.error("[WS建立连接失败],[用户过期],{},{}", userToken, taskId);
|
|
|
+ session.getBasicRemote().sendText("用户过期");
|
|
|
+ session.close();
|
|
|
+ } else {
|
|
|
+ MidjourneyUserConversation conversation = midjourneyUserConversationMapper.selectOne(QueryWrapperUtils.buildWrapper((wrapper) -> {
|
|
|
+ wrapper.eq(MidjourneyUserConversation::getUserId, user.getId()).eq(MidjourneyUserConversation::getTaskId, taskId);
|
|
|
+ }));
|
|
|
+ WssSession wsSession = new WssSession(session);
|
|
|
+ if (conversation == null) {
|
|
|
+ wsSession.sendMessage(WsMessageTypeEnum.SERVER_TASK_NOT_FIND, "任务不存在.");
|
|
|
+ session.close();
|
|
|
+ } else {
|
|
|
+ if (!redisService.hHasKey(RedisService.key.MIDJOURNEY_WS_SESSION.getName(), taskId)) {
|
|
|
+ redisService.hset(RedisService.key.MIDJOURNEY_WS_SESSION.getName(), taskId, wsSession, RedisService.key.MIDJOURNEY_WS_SESSION.getTimeout());
|
|
|
+ }
|
|
|
+ TASK_EXECUTOR.execute(() -> {
|
|
|
+ redisService.set(RedisService.key.MIDJOURNEY_CONVERSATION.getName() + taskId, conversation, RedisService.key.MIDJOURNEY_CONVERSATION.getTimeout());
|
|
|
+ });
|
|
|
+ }
|
|
|
+
|
|
|
+ }
|
|
|
+
|
|
|
+ }
|
|
|
+
|
|
|
+ /**
|
|
|
+ * 收到客户端消息回调
|
|
|
+ *
|
|
|
+ * @param message
|
|
|
+ * @param session
|
|
|
+ */
|
|
|
+ @OnMessage
|
|
|
+ public void OnMessage(String message, Session session, @PathParam("userToken") String userToken, @PathParam("taskId") String taskId) {
|
|
|
+ try {
|
|
|
+ logger.info("[WS收到消息],{},{},{}", userToken, taskId, message);
|
|
|
+
|
|
|
+ WssSession wsSession = new WssSession(session);
|
|
|
+ JSONObject body = null;//消息具体内容
|
|
|
+
|
|
|
+ //处理消息类型，是轮询还是消息
|
|
|
+ switch (WsMessageTypeEnum.getMessageType(message.split(":")[0])) {
|
|
|
+ //ping pong
|
|
|
+ case PING:
|
|
|
+ wsSession.pong();
|
|
|
+ return;
|
|
|
+ case MESSAGE:
|
|
|
+ body = JSON.parseObject(message.substring(8));
|
|
|
+ break;
|
|
|
+ default:
|
|
|
+ logger.error("[WS消息类型异常],{},{},{}", userToken, taskId, message);
|
|
|
+ return;
|
|
|
+ }
|
|
|
+
|
|
|
+ //根据子类型进行处理
|
|
|
+ switch (WsMessageTypeEnum.getMessageType(body.getString("type"))) {
|
|
|
+
|
|
|
+ case CLIENT_TASK_ID: //首次查看任务状态
|
|
|
+ midjourneyService.getConversationById(body.getLong("taskId"));
|
|
|
+ break;
|
|
|
+
|
|
|
+ case CLIENT_CLOSE:
|
|
|
+ logger.info("[WS连接关闭],{},{}", userToken, taskId);
|
|
|
+ //dailyQuizService.close(uniodId);
|
|
|
+ break;
|
|
|
+
|
|
|
+ default:
|
|
|
+ logger.error("[WS消息子类型异常],{},{},{}", userToken, taskId, message);
|
|
|
+
|
|
|
+ }
|
|
|
+
|
|
|
+ } catch (Exception e) {
|
|
|
+ logger.error("[WS接收消息失败],{},{},{},{}", userToken, taskId, message, e);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ /**
|
|
|
+ * 发生错误回调
|
|
|
+ *
|
|
|
+ * @param session
|
|
|
+ * @param error
|
|
|
+ */
|
|
|
+ @OnError
|
|
|
+ public void OnError(Session session, Throwable error) throws IOException {
|
|
|
+ logger.error("[WS ONERROR],{}", error.getMessage());
|
|
|
+ if (session.isOpen()) {
|
|
|
+ session.close();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ /**
|
|
|
+ * 关闭连接回调
|
|
|
+ *
|
|
|
+ * @param session
|
|
|
+ */
|
|
|
+ @OnClose
|
|
|
+ public void OnClose(Session session, @PathParam("userToken") String userToken, @PathParam("taskId") String taskId) throws IOException {
|
|
|
+
|
|
|
+ logger.info("[WS连接关闭],{},{}", userToken, taskId);
|
|
|
+ if (session.isOpen()) {
|
|
|
+ session.close();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ public MidjourneyUser getUser(String userToken) {
|
|
|
+
|
|
|
+ MidjourneyUser midjourneyUser = midjourneyUserMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUser.class)
|
|
|
+ .eq(MidjourneyUser::getUserToken, userToken).last("limit 1"));
|
|
|
+ if (midjourneyUser == null) {
|
|
|
+ throw new BusinessRuntimeException("账号不存在,请重新登录");
|
|
|
+ }
|
|
|
+ return midjourneyUser;
|
|
|
+
|
|
|
+ }
|
|
|
+
|
|
|
+}
|
chenbiao