|
@@ -1,76 +1,130 @@
|
|
|
|
|
+type MessageHandler = (msg: Record<string, unknown>) => void
|
|
|
|
|
+
|
|
|
|
|
+interface JobConnection {
|
|
|
|
|
+ ws: WebSocket
|
|
|
|
|
+ reconnectTimer: ReturnType<typeof setTimeout> | null
|
|
|
|
|
+ intentionalClose: boolean
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
class WSManager {
|
|
class WSManager {
|
|
|
- private ws: WebSocket | null = null
|
|
|
|
|
- private handlers: Map<string, Set<(msg: Record<string, unknown>) => void>> = new Map()
|
|
|
|
|
- private reconnectTimer: ReturnType<typeof setTimeout> | null = null
|
|
|
|
|
- private intentionalClose = false
|
|
|
|
|
-
|
|
|
|
|
- connect(baseUrl?: string) {
|
|
|
|
|
- if (this.ws) return
|
|
|
|
|
- this.intentionalClose = false
|
|
|
|
|
- const url = baseUrl || (import.meta.env.VITE_WS_BASE_URL as string) || '/ws'
|
|
|
|
|
- let wsUrl = url.startsWith('ws') ? url : `${window.location.protocol === 'https:' ? 'wss://' : 'ws://'}${window.location.host}${url}`
|
|
|
|
|
- const token = localStorage.getItem('token')
|
|
|
|
|
- if (token) {
|
|
|
|
|
- wsUrl += wsUrl.includes('?') ? '&' : '?'
|
|
|
|
|
- wsUrl += `token=${encodeURIComponent(token)}`
|
|
|
|
|
|
|
+ private connections: Map<string, JobConnection> = new Map()
|
|
|
|
|
+ private handlers: Map<string, Set<MessageHandler>> = new Map()
|
|
|
|
|
+
|
|
|
|
|
+ /**
|
|
|
|
|
+ * 为指定 job 建立 WebSocket 连接(/ws/training/{jobId})。
|
|
|
|
|
+ * 如果该 job 已有连接则跳过。
|
|
|
|
|
+ */
|
|
|
|
|
+ connect(jobId: string) {
|
|
|
|
|
+ if (this.connections.has(jobId)) return
|
|
|
|
|
+
|
|
|
|
|
+ const conn: JobConnection = {
|
|
|
|
|
+ ws: null as unknown as WebSocket,
|
|
|
|
|
+ reconnectTimer: null,
|
|
|
|
|
+ intentionalClose: false,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+ const protocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://'
|
|
|
|
|
+ const wsUrl = `${protocol}${window.location.host}/ws/training/${encodeURIComponent(jobId)}`
|
|
|
|
|
+ const token = localStorage.getItem('token')
|
|
|
|
|
+ const finalUrl = token ? `${wsUrl}?token=${encodeURIComponent(token)}` : wsUrl
|
|
|
|
|
+
|
|
|
try {
|
|
try {
|
|
|
- this.ws = new WebSocket(wsUrl)
|
|
|
|
|
|
|
+ conn.ws = new WebSocket(finalUrl)
|
|
|
} catch {
|
|
} catch {
|
|
|
- this.scheduleReconnect()
|
|
|
|
|
|
|
+ this.scheduleReconnect(jobId)
|
|
|
return
|
|
return
|
|
|
}
|
|
}
|
|
|
|
|
+ this.connections.set(jobId, conn)
|
|
|
|
|
|
|
|
- this.ws.onopen = () => {
|
|
|
|
|
- console.log('[WS] Connected')
|
|
|
|
|
|
|
+ conn.ws.onopen = () => {
|
|
|
|
|
+ console.log(`[WS] Connected for job ${jobId.slice(0, 8)}...`)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- this.ws.onmessage = (event) => {
|
|
|
|
|
|
|
+ conn.ws.onmessage = (event) => {
|
|
|
try {
|
|
try {
|
|
|
const msg = JSON.parse(event.data) as Record<string, unknown>
|
|
const msg = JSON.parse(event.data) as Record<string, unknown>
|
|
|
- this.handlers.get(msg.job_id as string)?.forEach(h => h(msg))
|
|
|
|
|
|
|
+ const msgJobId = (msg.job_id as string) || jobId
|
|
|
|
|
+ this.handlers.get(msgJobId)?.forEach(h => h(msg))
|
|
|
this.handlers.get('*')?.forEach(h => h(msg))
|
|
this.handlers.get('*')?.forEach(h => h(msg))
|
|
|
} catch {
|
|
} catch {
|
|
|
// ignore non-JSON messages
|
|
// ignore non-JSON messages
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- this.ws.onclose = () => {
|
|
|
|
|
- this.ws = null
|
|
|
|
|
- if (!this.intentionalClose) {
|
|
|
|
|
- this.scheduleReconnect()
|
|
|
|
|
|
|
+ conn.ws.onclose = () => {
|
|
|
|
|
+ this.connections.delete(jobId)
|
|
|
|
|
+ if (!conn.intentionalClose) {
|
|
|
|
|
+ this.scheduleReconnect(jobId)
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- this.ws.onerror = () => {
|
|
|
|
|
- this.ws?.close()
|
|
|
|
|
|
|
+ conn.ws.onerror = () => {
|
|
|
|
|
+ conn.ws?.close()
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- private scheduleReconnect() {
|
|
|
|
|
- if (this.reconnectTimer) return
|
|
|
|
|
- this.reconnectTimer = setTimeout(() => {
|
|
|
|
|
- this.reconnectTimer = null
|
|
|
|
|
- this.connect()
|
|
|
|
|
- }, 3000)
|
|
|
|
|
|
|
+ /** 是否已有该 job 的连接或正在重连 */
|
|
|
|
|
+ isConnected(jobId: string): boolean {
|
|
|
|
|
+ return this.connections.has(jobId)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- subscribe(jobId: string, handler: (msg: Record<string, unknown>) => void): () => void {
|
|
|
|
|
|
|
+ /**
|
|
|
|
|
+ * 断开指定 job 的 WebSocket 连接。
|
|
|
|
|
+ */
|
|
|
|
|
+ disconnect(jobId: string) {
|
|
|
|
|
+ const conn = this.connections.get(jobId)
|
|
|
|
|
+ if (!conn) return
|
|
|
|
|
+ conn.intentionalClose = true
|
|
|
|
|
+ if (conn.reconnectTimer) {
|
|
|
|
|
+ clearTimeout(conn.reconnectTimer)
|
|
|
|
|
+ conn.reconnectTimer = null
|
|
|
|
|
+ }
|
|
|
|
|
+ if (conn.ws) {
|
|
|
|
|
+ conn.ws.close()
|
|
|
|
|
+ }
|
|
|
|
|
+ this.connections.delete(jobId)
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ /**
|
|
|
|
|
+ * 断开所有 WebSocket 连接,清理所有 handler。
|
|
|
|
|
+ */
|
|
|
|
|
+ disconnectAll() {
|
|
|
|
|
+ for (const jobId of [...this.connections.keys()]) {
|
|
|
|
|
+ this.disconnect(jobId)
|
|
|
|
|
+ }
|
|
|
|
|
+ this.handlers.clear()
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ /**
|
|
|
|
|
+ * 订阅指定 job 的 WebSocket 消息。
|
|
|
|
|
+ * 返回取消订阅的函数。
|
|
|
|
|
+ */
|
|
|
|
|
+ subscribe(jobId: string, handler: MessageHandler): () => void {
|
|
|
if (!this.handlers.has(jobId)) this.handlers.set(jobId, new Set())
|
|
if (!this.handlers.has(jobId)) this.handlers.set(jobId, new Set())
|
|
|
this.handlers.get(jobId)!.add(handler)
|
|
this.handlers.get(jobId)!.add(handler)
|
|
|
return () => this.handlers.get(jobId)?.delete(handler)
|
|
return () => this.handlers.get(jobId)?.delete(handler)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- disconnect() {
|
|
|
|
|
- this.intentionalClose = true
|
|
|
|
|
- if (this.reconnectTimer) {
|
|
|
|
|
- clearTimeout(this.reconnectTimer)
|
|
|
|
|
- this.reconnectTimer = null
|
|
|
|
|
- }
|
|
|
|
|
- if (this.ws) {
|
|
|
|
|
- this.ws.close()
|
|
|
|
|
- this.ws = null
|
|
|
|
|
|
|
+ private scheduleReconnect(jobId: string) {
|
|
|
|
|
+ const conn = this.connections.get(jobId)
|
|
|
|
|
+ if (conn?.reconnectTimer) return
|
|
|
|
|
+ const timer = setTimeout(() => {
|
|
|
|
|
+ if (conn) conn.reconnectTimer = null
|
|
|
|
|
+ // 仅在 handler 仍在订阅时才重连(说明页面还关心这个 job)
|
|
|
|
|
+ if (this.handlers.get(jobId)?.size) {
|
|
|
|
|
+ this.connections.delete(jobId)
|
|
|
|
|
+ this.connect(jobId)
|
|
|
|
|
+ }
|
|
|
|
|
+ }, 3000)
|
|
|
|
|
+ if (conn) {
|
|
|
|
|
+ conn.reconnectTimer = timer
|
|
|
|
|
+ } else {
|
|
|
|
|
+ // conn 已删除,用临时对象跟踪 timer
|
|
|
|
|
+ this.connections.set(jobId, {
|
|
|
|
|
+ ws: null as unknown as WebSocket,
|
|
|
|
|
+ reconnectTimer: timer,
|
|
|
|
|
+ intentionalClose: false,
|
|
|
|
|
+ })
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|