BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 44e44dc4bb1431a293afd2d3cd9c57fccb65258e

Authored by cherishsince
1 parent 5044a581

增加注释

Inline Side-by-side
Showing 3 changed files with 113 additions and 53 deletions
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/constants/MjNotifyCode.java
  View file @44e44dc
... ... @@ -8,35 +8,7 @@ public final class MjNotifyCode {
8 8 * 成功.
9 9 */
10 10 public static final int SUCCESS = 1;
11   - /**
12   - * 数据未找到.
13   - */
14   - public static final int NOT_FOUND = 3;
15   - /**
16   - * 校验错误.
17   - */
18   - public static final int VALIDATION_ERROR = 4;
19   - /**
20   - * 系统异常.
21   - */
22   - public static final int FAILURE = 9;
23 11  
24   - /**
25   - * 已存在.
26   - */
27   - public static final int EXISTED = 21;
28   - /**
29   - * 排队中.
30   - */
31   - public static final int IN_QUEUE = 22;
32   - /**
33   - * 队列已满.
34   - */
35   - public static final int QUEUE_REJECTED = 23;
36   - /**
37   - * prompt包含敏感词.
38   - */
39   - public static final int BANNED_PROMPT = 24;
40 12  
41 13  
42 14 }
43 15 \ No newline at end of file
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MjWebSocketStarter.java
  View file @44e44dc
... ... @@ -7,9 +7,11 @@ import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyConfig;
7 7 import cn.iocoder.yudao.framework.ai.midjourney.constants.MjNotifyCode;
8 8 import cn.iocoder.yudao.framework.ai.midjourney.webSocket.handler.MjWebSocketHandler;
9 9 import cn.iocoder.yudao.framework.ai.midjourney.webSocket.listener.MjMessageListener;
  10 +import lombok.Getter;
10 11 import lombok.extern.slf4j.Slf4j;
11 12 import org.apache.tomcat.websocket.Constants;
12 13 import org.jetbrains.annotations.NotNull;
  14 +import org.springframework.util.concurrent.ListenableFuture;
13 15 import org.springframework.util.concurrent.ListenableFutureCallback;
14 16 import org.springframework.web.socket.CloseStatus;
15 17 import org.springframework.web.socket.WebSocketHttpHeaders;
... ... @@ -22,19 +24,43 @@ import java.util.concurrent.TimeoutException;
22 24  
23 25 @Slf4j
24 26 public class MjWebSocketStarter implements WebSocketStarter {
  27 + /**
  28 + * 链接重试次数
  29 + */
25 30 private static final int CONNECT_RETRY_LIMIT = 5;
26   -
  31 + /**
  32 + * mj 配置文件
  33 + */
27 34 private final MidjourneyConfig midjourneyConfig;
  35 + /**
  36 + * mj 监听(所有message 都会 callback到这里)
  37 + */
28 38 private final MjMessageListener userMessageListener;
  39 + /**
  40 + * wss 服务器
  41 + */
29 42 private final String wssServer;
  43 + /**
  44 + *
  45 + */
30 46 private final String resumeWss;
31   -
  47 + /**
  48 + *
  49 + */
  50 + private ResumeData resumeData = null;
  51 + /**
  52 + * 是否运行成功
  53 + */
32 54 private boolean running = false;
33   -
  55 + /**
  56 + * 链接成功的 session
  57 + */
34 58 private WebSocketSession webSocketSession = null;
35   - private ResumeData resumeData = null;
36 59  
37   - public MjWebSocketStarter(String wssServer, String resumeWss, MidjourneyConfig midjourneyConfig, MjMessageListener userMessageListener) {
  60 + public MjWebSocketStarter(String wssServer,
  61 + String resumeWss,
  62 + MidjourneyConfig midjourneyConfig,
  63 + MjMessageListener userMessageListener) {
38 64 this.wssServer = wssServer;
39 65 this.resumeWss = resumeWss;
40 66 this.midjourneyConfig = midjourneyConfig;
... ... @@ -42,11 +68,12 @@ public class MjWebSocketStarter implements WebSocketStarter {
42 68 }
43 69  
44 70 @Override
45   - public void start() throws Exception {
  71 + public void start() {
46 72 start(false);
47 73 }
48 74  
49 75 private void start(boolean reconnect) {
  76 + // 设置header
50 77 WebSocketHttpHeaders headers = new WebSocketHttpHeaders();
51 78 headers.add("Accept-Encoding", "gzip, deflate, br");
52 79 headers.add("Accept-Language", "zh-CN,zh;q=0.9");
... ... @@ -54,19 +81,26 @@ public class MjWebSocketStarter implements WebSocketStarter {
54 81 headers.add("Pragma", "no-cache");
55 82 headers.add("Sec-Websocket-Extensions", "permessage-deflate; client_max_window_bits");
56 83 headers.add("User-Agent", this.midjourneyConfig.getUserAage());
57   - var handler = new MjWebSocketHandler(this.midjourneyConfig, this.userMessageListener, this::onSocketSuccess, this::onSocketFailure);
  84 + // 创建 mjHeader
  85 + MjWebSocketHandler mjWebSocketHandler = new MjWebSocketHandler(
  86 + this.midjourneyConfig, this.userMessageListener, this::onSocketSuccess, this::onSocketFailure);
  87 + //
58 88 String gatewayUrl;
59 89 if (reconnect) {
60   - gatewayUrl = getGatewayServer(this.resumeData.resumeGatewayUrl()) + "/?encoding=json&v=9&compress=zlib-stream";
61   - handler.setSessionId(this.resumeData.sessionId());
62   - handler.setSequence(this.resumeData.sequence());
63   - handler.setResumeGatewayUrl(this.resumeData.resumeGatewayUrl());
  90 + gatewayUrl = getGatewayServer(this.resumeData.getResumeGatewayUrl()) + "/?encoding=json&v=9&compress=zlib-stream";
  91 + mjWebSocketHandler.setSessionId(this.resumeData.getSessionId());
  92 + mjWebSocketHandler.setSequence(this.resumeData.getSequence());
  93 + mjWebSocketHandler.setResumeGatewayUrl(this.resumeData.getResumeGatewayUrl());
64 94 } else {
65 95 gatewayUrl = getGatewayServer(null) + "/?encoding=json&v=9&compress=zlib-stream";
66 96 }
67   - var webSocketClient = new StandardWebSocketClient();
  97 + // 创建 StandardWebSocketClient
  98 + StandardWebSocketClient webSocketClient = new StandardWebSocketClient();
  99 + // 设置 io timeout 时间
68 100 webSocketClient.getUserProperties().put(Constants.IO_TIMEOUT_MS_PROPERTY, "10000");
69   - var socketSessionFuture = webSocketClient.doHandshake(handler, headers, URI.create(gatewayUrl));
  101 + //
  102 + ListenableFuture<WebSocketSession> socketSessionFuture = webSocketClient.doHandshake(mjWebSocketHandler, headers, URI.create(gatewayUrl));
  103 + // 添加 callback 进行回调
70 104 socketSessionFuture.addCallback(new ListenableFutureCallback<>() {
71 105 @Override
72 106 public void onFailure(@NotNull Throwable e) {
... ... @@ -87,14 +121,18 @@ public class MjWebSocketStarter implements WebSocketStarter {
87 121 }
88 122  
89 123 private void onSocketFailure(int code, String reason) {
  124 + // 1001异常可以忽略
90 125 if (code == 1001) {
91 126 return;
92 127 }
  128 + // 关闭 socket
93 129 closeSocketSessionWhenIsOpen();
  130 + // 没有运行通知
94 131 if (!this.running) {
95 132 notifyWssLock(code, reason);
96 133 return;
97 134 }
  135 + // 已经运行先设置为false，发起
98 136 this.running = false;
99 137 if (code >= 4000) {
100 138 log.warn("[wss-{}] Can't reconnect! Account disabled. Closed by {}({}).", this.midjourneyConfig.getChannelId(), code, reason);
... ... @@ -107,36 +145,34 @@ public class MjWebSocketStarter implements WebSocketStarter {
107 145 }
108 146 }
109 147  
  148 + /**
  149 + * 重连
  150 + */
110 151 private void tryReconnect() {
111 152 try {
112 153 tryStart(true);
113 154 } catch (Exception e) {
114   - if (e instanceof TimeoutException) {
115   - closeSocketSessionWhenIsOpen();
116   - }
117   - log.warn("[wss-{}] Reconnect fail: {}, Try new connection...", this.midjourneyConfig.getChannelId(), e.getMessage());
  155 + log.warn("[wss-{}] Reconnect fail: {}, Try new connection...", this.midjourneyConfig.getChannelId(), e.getMessage());
118 156 ThreadUtil.sleep(1000);
119 157 tryNewConnect();
120 158 }
121 159 }
122 160  
123 161 private void tryNewConnect() {
  162 + // 链接重试次数5
124 163 for (int i = 1; i <= CONNECT_RETRY_LIMIT; i++) {
125 164 try {
126 165 tryStart(false);
127 166 return;
128 167 } catch (Exception e) {
129   - if (e instanceof TimeoutException) {
130   - closeSocketSessionWhenIsOpen();
131   - }
132   - log.warn("[wss-{}] New connect fail ({}): {}", this.midjourneyConfig.getChannelId(), i, e.getMessage());
  168 + log.warn("[wss-{}] New connect fail ({}): {}", this.midjourneyConfig.getChannelId(), i, e.getMessage());
133 169 ThreadUtil.sleep(5000);
134 170 }
135 171 }
136 172 log.error("[wss-{}] Account disabled", this.midjourneyConfig.getChannelId());
137 173 }
138 174  
139   - public void tryStart(boolean reconnect) throws Exception {
  175 + public void tryStart(boolean reconnect) {
140 176 start(reconnect);
141 177 }
142 178  
... ... @@ -144,6 +180,9 @@ public class MjWebSocketStarter implements WebSocketStarter {
144 180 System.err.println("notifyWssLock: " + code + " - " + reason);
145 181 }
146 182  
  183 + /**
  184 + * 关闭 socket session
  185 + */
147 186 private void closeSocketSessionWhenIsOpen() {
148 187 try {
149 188 if (this.webSocketSession != null && this.webSocketSession.isOpen()) {
... ... @@ -161,6 +200,20 @@ public class MjWebSocketStarter implements WebSocketStarter {
161 200 return this.wssServer;
162 201 }
163 202  
164   - public record ResumeData(String sessionId, Object sequence, String resumeGatewayUrl) {
  203 + @Getter
  204 + public static class ResumeData {
  205 +
  206 + public ResumeData(String sessionId, Object sequence, String resumeGatewayUrl) {
  207 + this.sessionId = sessionId;
  208 + this.sequence = sequence;
  209 + this.resumeGatewayUrl = resumeGatewayUrl;
  210 + }
  211 +
  212 + /**
  213 + * socket session
  214 + */
  215 + private final String sessionId;
  216 + private final Object sequence;
  217 + private final String resumeGatewayUrl;
165 218 }
166 219 }
167 220 \ No newline at end of file
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/handler/MjWebSocketHandler.java
  View file @44e44dc
... ... @@ -30,16 +30,41 @@ import java.util.concurrent.TimeUnit;
30 30  
31 31 @Slf4j
32 32 public class MjWebSocketHandler implements WebSocketHandler {
  33 + /**
  34 + * close 错误码：重连
  35 + */
33 36 public static final int CLOSE_CODE_RECONNECT = 2001;
  37 + /**
  38 + * close 错误码：无效、作废
  39 + */
34 40 public static final int CLOSE_CODE_INVALIDATE = 1009;
  41 + /**
  42 + * close 错误码：异常
  43 + */
35 44 public static final int CLOSE_CODE_EXCEPTION = 1011;
36   -
  45 + /**
  46 + * mj配置文件
  47 + */
37 48 private final MidjourneyConfig midjourneyConfig;
  49 + /**
  50 + * mj 消息监听
  51 + */
38 52 private final MjMessageListener userMessageListener;
  53 + /**
  54 + * 成功回调
  55 + */
39 56 private final SuccessCallback successCallback;
  57 + /**
  58 + * 失败回调
  59 + */
40 60 private final FailureCallback failureCallback;
41   -
  61 + /**
  62 + * 心跳执行器
  63 + */
42 64 private final ScheduledExecutorService heartExecutor;
  65 + /**
  66 + * auth数据
  67 + */
43 68 private final DataObject authData;
44 69  
45 70 @Setter
... ... @@ -55,6 +80,9 @@ public class MjWebSocketHandler implements WebSocketHandler {
55 80 private Future<?> heartbeatInterval;
56 81 private Future<?> heartbeatTimeout;
57 82  
  83 + /**
  84 + * 处理 message 消息的 Decompressor
  85 + */
58 86 private final Decompressor decompressor = new ZlibDecompressor(2048);
59 87  
60 88 public MjWebSocketHandler(MidjourneyConfig account,
... ... @@ -77,11 +105,13 @@ public class MjWebSocketHandler implements WebSocketHandler {
77 105 @Override
78 106 public void handleTransportError(@NotNull WebSocketSession session, @NotNull Throwable e) throws Exception {
79 107 log.error("[wss-{}] Transport error", this.midjourneyConfig.getChannelId(), e);
  108 + // 通知链接异常
80 109 onFailure(CLOSE_CODE_EXCEPTION, "transport error");
81 110 }
82 111  
83 112 @Override
84 113 public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus closeStatus) throws Exception {
  114 + // 链接关闭
85 115 onFailure(closeStatus.getCode(), closeStatus.getReason());
86 116 }
87 117  
... ... @@ -92,13 +122,18 @@ public class MjWebSocketHandler implements WebSocketHandler {
92 122  
93 123 @Override
94 124 public void handleMessage(@NotNull WebSocketSession session, WebSocketMessage<?> message) throws Exception {
  125 + // 获取 message 消息
95 126 ByteBuffer buffer = (ByteBuffer) message.getPayload();
  127 + // 解析 message
96 128 byte[] decompressed = decompressor.decompress(buffer.array());
97 129 if (decompressed == null) {
98 130 return;
99 131 }
  132 + // 转换 json
100 133 String json = new String(decompressed, StandardCharsets.UTF_8);
  134 + // 转换 jda 自带的 dataObject(和json object 差不多)
101 135 DataObject data = DataObject.fromJson(json);
  136 + // 获取消息类型
102 137 int opCode = data.getInt("op");
103 138 switch (opCode) {
104 139 case WebSocketCode.HEARTBEAT -> handleHeartbeat(session);
... ...