package handlers import ( "encoding/json" "net/http" "sync" "time" "github.com/gorilla/websocket" ) var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { // 允许所有来源(生产环境应该做更严格的检查) return true }, } // ProgressMessage WebSocket进度消息结构 type ProgressMessage struct { Status string `json:"status"` // 状态: checking, cloning, pulling, success, error Message string `json:"message"` // 消息文本 Percent int `json:"percent"` // 进度百分比 (0-100) } // ProgressHub 进度广播中心 type ProgressHub struct { clients map[string]*websocket.Conn clientsMux sync.RWMutex broadcast chan ProgressMessage } var ( hub *ProgressHub hubOnce sync.Once ) // GetProgressHub 获取进度广播中心单例 func GetProgressHub() *ProgressHub { hubOnce.Do(func() { hub = &ProgressHub{ clients: make(map[string]*websocket.Conn), broadcast: make(chan ProgressMessage, 10), } go hub.run() }) return hub } // run 运行广播循环 func (h *ProgressHub) run() { for msg := range h.broadcast { h.clientsMux.RLock() for clientID, conn := range h.clients { err := conn.WriteJSON(msg) if err != nil { // 写入失败,移除客户端 conn.Close() h.clientsMux.RUnlock() h.removeClient(clientID) h.clientsMux.RLock() } } h.clientsMux.RUnlock() } } // addClient 添加客户端 func (h *ProgressHub) addClient(clientID string, conn *websocket.Conn) { h.clientsMux.Lock() defer h.clientsMux.Unlock() h.clients[clientID] = conn } // removeClient 移除客户端 func (h *ProgressHub) removeClient(clientID string) { h.clientsMux.Lock() defer h.clientsMux.Unlock() delete(h.clients, clientID) } // SendProgress 发送进度消息 func (h *ProgressHub) SendProgress(msg ProgressMessage) { select { case h.broadcast <- msg: case <-time.After(1 * time.Second): // 超时,丢弃消息 } } // GitProgressWS WebSocket处理器 func GitProgressWS() http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { // 升级HTTP连接到WebSocket conn, err := upgrader.Upgrade(w, r, nil) if err != nil { http.Error(w, "Failed to upgrade to WebSocket", http.StatusInternalServerError) return } defer conn.Close() // 生成客户端ID clientID := r.RemoteAddr + "-" + time.Now().Format("20060102150405") // 添加到hub hub := GetProgressHub() hub.addClient(clientID, conn) defer hub.removeClient(clientID) // 发送欢迎消息 welcomeMsg := ProgressMessage{ Status: "connected", Message: "WebSocket连接已建立", Percent: 0, } if err := conn.WriteJSON(welcomeMsg); err != nil { return } // 保持连接,读取客户端消息(主要用于心跳) for { _, _, err := conn.ReadMessage() if err != nil { // 连接断开 break } } } } // SendProgressToClients 辅助函数:向所有客户端发送进度 func SendProgressToClients(status, message string, percent int) { hub := GetProgressHub() hub.SendProgress(ProgressMessage{ Status: status, Message: message, Percent: percent, }) } // SendProgressJSON 辅助函数:发送JSON格式的进度消息 func SendProgressJSON(w http.ResponseWriter, status, message string, percent int) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(ProgressMessage{ Status: status, Message: message, Percent: percent, }) }