fix: stabilize datasource progress and websocket flow
This commit is contained in:
linkong
@@ -1,7 +1,38 @@
|
||||
import { useEffect, useRef, useState, useCallback } from 'react'
|
||||
import { useAuthStore } from '../stores/auth'
|
||||
|
||||
const WS_URL = (import.meta as any).env?.VITE_WS_URL || 'ws://localhost:8000/ws'
|
||||
const DEFAULT_WS_URL = (() => {
|
||||
if (typeof window === 'undefined') {
|
||||
return 'ws://localhost:8000/ws'
|
||||
}
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
|
||||
return `${protocol}//${window.location.host}/ws`
|
||||
})()
|
||||
|
||||
const WS_URL = (import.meta as any).env?.VITE_WS_URL || DEFAULT_WS_URL
|
||||
|
||||
function buildWebSocketCandidates(): string[] {
|
||||
if ((import.meta as any).env?.VITE_WS_URL) {
|
||||
return [WS_URL]
|
||||
}
|
||||
|
||||
if (typeof window === 'undefined') {
|
||||
return ['ws://localhost:8000/ws']
|
||||
}
|
||||
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
|
||||
const sameOrigin = `${protocol}//${window.location.host}/ws`
|
||||
const candidates = [sameOrigin]
|
||||
|
||||
if (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1') {
|
||||
const directBackend = `${protocol}//${window.location.hostname}:8000/ws`
|
||||
if (!candidates.includes(directBackend)) {
|
||||
candidates.push(directBackend)
|
||||
}
|
||||
}
|
||||
|
||||
return candidates
|
||||
}
|
||||
|
||||
interface WebSocketMessage {
|
||||
type: string
|
||||
@@ -41,71 +72,130 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
|
||||
} = options
|
||||
|
||||
const { token } = useAuthStore()
|
||||
console.log('[WebSocket] Token present:', !!token, token ? token.substring(0, 20) + '...' : 'none')
|
||||
console.log('[WebSocket] autoConnect:', autoConnect)
|
||||
|
||||
const wsRef = useRef<WebSocket | null>(null)
|
||||
const [connected, setConnected] = useState(false)
|
||||
const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null)
|
||||
const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
|
||||
const heartbeatTimerRef = useRef<ReturnType<typeof setInterval> | null>(null)
|
||||
const activeWsUrlRef = useRef<string | null>(null)
|
||||
const intentionalCloseRef = useRef(false)
|
||||
const pendingCloseSocketRef = useRef<WebSocket | null>(null)
|
||||
const autoSubscribeRef = useRef(autoSubscribe)
|
||||
const onMessageRef = useRef(onMessage)
|
||||
const onConnectRef = useRef(onConnect)
|
||||
const onDisconnectRef = useRef(onDisconnect)
|
||||
const onErrorRef = useRef(onError)
|
||||
|
||||
useEffect(() => {
|
||||
autoSubscribeRef.current = autoSubscribe
|
||||
onMessageRef.current = onMessage
|
||||
onConnectRef.current = onConnect
|
||||
onDisconnectRef.current = onDisconnect
|
||||
onErrorRef.current = onError
|
||||
}, [autoSubscribe, onMessage, onConnect, onDisconnect, onError])
|
||||
|
||||
const connect = useCallback(() => {
|
||||
if (!token) {
|
||||
console.log('[WebSocket] No token, skipping connect')
|
||||
return
|
||||
}
|
||||
|
||||
const wsUrl = `${WS_URL}?token=${token}`
|
||||
console.log('[WebSocket] Connecting to:', wsUrl)
|
||||
const ws = new WebSocket(wsUrl)
|
||||
intentionalCloseRef.current = false
|
||||
const candidates = buildWebSocketCandidates()
|
||||
let candidateIndex = 0
|
||||
let opened = false
|
||||
|
||||
ws.onopen = () => {
|
||||
console.log('[WebSocket] Connected!')
|
||||
setConnected(true)
|
||||
if (autoSubscribe.length > 0) {
|
||||
ws.send(JSON.stringify({ type: 'subscribe', data: { channels: autoSubscribe } }))
|
||||
const tryConnect = () => {
|
||||
const baseUrl = candidates[candidateIndex]
|
||||
const wsUrl = `${baseUrl}?token=${token}`
|
||||
activeWsUrlRef.current = baseUrl
|
||||
const ws = new WebSocket(wsUrl)
|
||||
|
||||
ws.onopen = () => {
|
||||
if (intentionalCloseRef.current || pendingCloseSocketRef.current === ws) {
|
||||
pendingCloseSocketRef.current = null
|
||||
ws.close()
|
||||
return
|
||||
}
|
||||
opened = true
|
||||
setConnected(true)
|
||||
if (autoSubscribeRef.current.length > 0) {
|
||||
ws.send(JSON.stringify({ type: 'subscribe', data: { channels: autoSubscribeRef.current } }))
|
||||
}
|
||||
onConnectRef.current?.()
|
||||
}
|
||||
onConnect?.()
|
||||
|
||||
ws.onmessage = (event) => {
|
||||
try {
|
||||
const message: WebSocketMessage = JSON.parse(event.data)
|
||||
setLastMessage(message)
|
||||
onMessageRef.current?.(message)
|
||||
} catch {
|
||||
console.error('Failed to parse WebSocket message')
|
||||
}
|
||||
}
|
||||
|
||||
ws.onclose = () => {
|
||||
if (wsRef.current === ws) {
|
||||
wsRef.current = null
|
||||
}
|
||||
if (pendingCloseSocketRef.current === ws) {
|
||||
pendingCloseSocketRef.current = null
|
||||
}
|
||||
setConnected(false)
|
||||
if (heartbeatTimerRef.current) {
|
||||
clearInterval(heartbeatTimerRef.current)
|
||||
heartbeatTimerRef.current = null
|
||||
}
|
||||
|
||||
if (!opened && candidateIndex < candidates.length - 1) {
|
||||
candidateIndex += 1
|
||||
tryConnect()
|
||||
return
|
||||
}
|
||||
|
||||
if (intentionalCloseRef.current) {
|
||||
return
|
||||
}
|
||||
|
||||
onDisconnectRef.current?.()
|
||||
|
||||
if (autoConnect && token) {
|
||||
reconnectTimeoutRef.current = setTimeout(() => {
|
||||
connect()
|
||||
}, 3000)
|
||||
}
|
||||
}
|
||||
|
||||
ws.onerror = (error) => {
|
||||
setConnected(false)
|
||||
if (intentionalCloseRef.current || ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED) {
|
||||
return
|
||||
}
|
||||
if (candidateIndex >= candidates.length - 1) {
|
||||
console.warn('[WebSocket] Connection error', { url: baseUrl, error })
|
||||
onErrorRef.current?.(new Error('WebSocket error'))
|
||||
}
|
||||
}
|
||||
|
||||
wsRef.current = ws
|
||||
}
|
||||
|
||||
ws.onmessage = (event) => {
|
||||
console.log('[WebSocket] Received:', event.data)
|
||||
try {
|
||||
const message: WebSocketMessage = JSON.parse(event.data)
|
||||
setLastMessage(message)
|
||||
onMessage?.(message)
|
||||
} catch {
|
||||
console.error('Failed to parse WebSocket message')
|
||||
}
|
||||
}
|
||||
|
||||
ws.onclose = (event) => {
|
||||
console.log('[WebSocket] Disconnected:', event.code, event.reason)
|
||||
try {
|
||||
tryConnect()
|
||||
} catch (error) {
|
||||
setConnected(false)
|
||||
if (heartbeatTimerRef.current) {
|
||||
clearInterval(heartbeatTimerRef.current)
|
||||
heartbeatTimerRef.current = null
|
||||
}
|
||||
onDisconnect?.()
|
||||
|
||||
console.warn('[WebSocket] Failed to initialize connection', { url: activeWsUrlRef.current, error })
|
||||
if (autoConnect && token) {
|
||||
reconnectTimeoutRef.current = setTimeout(() => {
|
||||
console.log('[WebSocket] Reconnecting...')
|
||||
connect()
|
||||
}, 3000)
|
||||
}
|
||||
}
|
||||
|
||||
ws.onerror = (error) => {
|
||||
console.error('[WebSocket] Error:', error)
|
||||
onError?.(new Error('WebSocket error'))
|
||||
}
|
||||
|
||||
wsRef.current = ws
|
||||
}, [token, autoConnect, autoSubscribe, onConnect, onDisconnect, onError])
|
||||
}, [token, autoConnect])
|
||||
|
||||
const disconnect = useCallback(() => {
|
||||
intentionalCloseRef.current = true
|
||||
if (reconnectTimeoutRef.current) {
|
||||
clearTimeout(reconnectTimeoutRef.current)
|
||||
reconnectTimeoutRef.current = null
|
||||
@@ -114,8 +204,14 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
|
||||
clearInterval(heartbeatTimerRef.current)
|
||||
heartbeatTimerRef.current = null
|
||||
}
|
||||
wsRef.current?.close()
|
||||
wsRef.current = null
|
||||
const socket = wsRef.current
|
||||
if (socket) {
|
||||
if (socket.readyState === WebSocket.CONNECTING) {
|
||||
pendingCloseSocketRef.current = socket
|
||||
} else if (socket.readyState === WebSocket.OPEN) {
|
||||
socket.close()
|
||||
}
|
||||
}
|
||||
setConnected(false)
|
||||
}, [])
|
||||
|
||||
@@ -131,7 +227,6 @@ export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketRet
|
||||
}, [])
|
||||
|
||||
useEffect(() => {
|
||||
console.log('[WebSocket] useEffect triggered, autoConnect:', autoConnect, 'token:', !!token)
|
||||
if (autoConnect && token) {
|
||||
connect()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user