package ws import ( "encoding/json" "fmt" "os" "path/filepath" "sync" "time" "github.com/gorilla/websocket" ) type Client struct { Conn *websocket.Conn Send chan []byte } type Subscription struct { Client *Client Room string done chan struct{} // used by RegisterSync to wait for completion } type RoomMessage struct { Room string Message []byte } // Hub 管理 WebSocket client 訂閱與訊息廣播。 // // M8-4b:Hub 額外負責「第一個 client 連上時寫 sentinel file」, // 讓 Wails 端的 StartupPipeline 知道階段 6(Wait for Web UI WebSocket)已完成。 // 詳細設計見 .autoflow/04-architecture/v2/startup-pipeline.md §3。 // // dataDir 由 main.go 在初始化 Hub 後透過 SetStartupSentinel(dataDir) 注入。 // 若 dataDir 為空,sentinel 寫入會被跳過(單元測試或缺少資料目錄時的安全行為)。 type Hub struct { rooms map[string]map[*Client]bool register chan *Subscription unregister chan *Subscription broadcast chan *RoomMessage mu sync.RWMutex // M8-4b: 啟動 sentinel file sentinelDataDir string // ,由 SetStartupSentinel 設定 sentinelOnce sync.Once // 確保只在「第一個」client 連上時寫一次 bootID string // 寫入 sentinel 內容供 debug } func NewHub() *Hub { return &Hub{ rooms: make(map[string]map[*Client]bool), register: make(chan *Subscription, 10), unregister: make(chan *Subscription, 10), broadcast: make(chan *RoomMessage, 100), bootID: fmt.Sprintf("boot-%d", time.Now().UnixNano()), } } // SetStartupSentinel 設定 sentinel file 的根目錄。 // main.go 在 NewHub() 之後、Run() 之前呼叫一次,dataDir 應為完整路徑。 // // 寫入路徑:/.first-ws-connected // 內容:boot-id + timestamp(用於 debug,內容對 Wails 端的判斷沒有意義,存在即可) // // dataDir 為空字串時 sentinel 機制完全停用。 func (h *Hub) SetStartupSentinel(dataDir string) { h.mu.Lock() h.sentinelDataDir = dataDir h.mu.Unlock() } // writeStartupSentinel 在第一個 WebSocket client 連上時呼叫一次。 // 由 sentinelOnce 確保只執行一次;後續連線完全 no-op。 // // 寫入失敗不會 panic 也不會回 error:sentinel 是 best-effort 機制, // 若 disk 滿/權限錯,Wails 端會走 hard timeout 路徑進 Error state。 func (h *Hub) writeStartupSentinel() { h.sentinelOnce.Do(func() { h.mu.RLock() dir := h.sentinelDataDir bootID := h.bootID h.mu.RUnlock() if dir == "" { return } path := filepath.Join(dir, ".first-ws-connected") // 確保父目錄存在(dataDir 通常已存在,但保險起見) _ = os.MkdirAll(dir, 0o755) f, err := os.Create(path) if err != nil { return } _, _ = fmt.Fprintf(f, "bootId=%s\nts=%d\n", bootID, time.Now().UnixMilli()) _ = f.Close() }) } func (h *Hub) Run() { for { select { case sub := <-h.register: h.mu.Lock() if h.rooms[sub.Room] == nil { h.rooms[sub.Room] = make(map[*Client]bool) } h.rooms[sub.Room][sub.Client] = true h.mu.Unlock() // M8-4b:第一次有 client 加入任何 room → 寫 sentinel file // (sync.Once 保證後續呼叫 no-op) h.writeStartupSentinel() if sub.done != nil { close(sub.done) } case sub := <-h.unregister: h.mu.Lock() if clients, ok := h.rooms[sub.Room]; ok { if _, exists := clients[sub.Client]; exists { delete(clients, sub.Client) close(sub.Client.Send) } } h.mu.Unlock() case msg := <-h.broadcast: h.mu.RLock() if clients, ok := h.rooms[msg.Room]; ok { for client := range clients { select { case client.Send <- msg.Message: default: close(client.Send) delete(clients, client) } } } h.mu.RUnlock() } } } func (h *Hub) Register(sub *Subscription) { h.register <- sub } // RegisterSync registers a subscription and blocks until the Hub has processed it, // ensuring the client is in the room before returning. func (h *Hub) RegisterSync(sub *Subscription) { sub.done = make(chan struct{}) h.register <- sub <-sub.done } func (h *Hub) Unregister(sub *Subscription) { h.unregister <- sub } func (h *Hub) BroadcastToRoom(room string, data interface{}) { jsonData, err := json.Marshal(data) if err != nil { return } h.broadcast <- &RoomMessage{Room: room, Message: jsonData} }