feat: rollout

internal/websocket/websocket.go

@@ -3,6 +3,7 @@ package websocket
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"sync"
 	"time"
@@ -90,13 +91,10 @@ func (wss *WebSocketServer) handleClient(conn *websocket.Conn) {
 		ticker := time.NewTicker(54 * time.Second)
 		defer ticker.Stop()
 
-		for {
-			select {
-			case <-ticker.C:
-				conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
-				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
-					return
-				}
+		for range ticker.C {
+			conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
 			}
 		}
 	}()
@@ -334,6 +332,99 @@ func (wss *WebSocketServer) BroadcastFirmwareUploadStatus(nodeIP, status, filena
 	}
 }
 
+// BroadcastRolloutProgress sends rollout progress updates to all clients
+func (wss *WebSocketServer) BroadcastRolloutProgress(rolloutID, nodeIP, status string, current, total int) {
+	wss.mutex.RLock()
+	clients := make([]*websocket.Conn, 0, len(wss.clients))
+	for client := range wss.clients {
+		clients = append(clients, client)
+	}
+	wss.mutex.RUnlock()
+
+	if len(clients) == 0 {
+		return
+	}
+
+	message := struct {
+		Type      string `json:"type"`
+		RolloutID string `json:"rolloutId"`
+		NodeIP    string `json:"nodeIp"`
+		Status    string `json:"status"`
+		Current   int    `json:"current"`
+		Total     int    `json:"total"`
+		Progress  int    `json:"progress"`
+		Timestamp string `json:"timestamp"`
+	}{
+		Type:      "rollout_progress",
+		RolloutID: rolloutID,
+		NodeIP:    nodeIP,
+		Status:    status,
+		Current:   current,
+		Total:     total,
+		Progress:  wss.calculateProgress(current, total, status),
+		Timestamp: time.Now().Format(time.RFC3339),
+	}
+
+	data, err := json.Marshal(message)
+	if err != nil {
+		wss.logger.WithError(err).Error("Failed to marshal rollout progress")
+		return
+	}
+
+	wss.logger.WithFields(log.Fields{
+		"rollout_id": rolloutID,
+		"node_ip":    nodeIP,
+		"status":     status,
+		"progress":   fmt.Sprintf("%d/%d", current, total),
+		"clients":    len(clients),
+	}).Debug("Broadcasting rollout progress to WebSocket clients")
+
+	// Send to all clients with write synchronization
+	wss.writeMutex.Lock()
+	defer wss.writeMutex.Unlock()
+
+	for _, client := range clients {
+		client.SetWriteDeadline(time.Now().Add(5 * time.Second))
+		if err := client.WriteMessage(websocket.TextMessage, data); err != nil {
+			wss.logger.WithError(err).Error("Failed to send rollout progress to client")
+		}
+	}
+}
+
+// calculateProgress calculates the correct progress percentage based on current status
+func (wss *WebSocketServer) calculateProgress(current, total int, status string) int {
+	if total == 0 {
+		return 0
+	}
+
+	// Base progress is based on completed nodes
+	completedNodes := current - 1
+	if status == "completed" {
+		completedNodes = current
+	}
+
+	// Calculate base progress (completed nodes / total nodes)
+	baseProgress := float64(completedNodes) / float64(total) * 100
+
+	// If currently updating labels or uploading, add partial progress for the current node
+	if status == "updating_labels" {
+		// Add 25% of one node's progress (label update is quick)
+		nodeProgress := 100.0 / float64(total) * 0.25
+		baseProgress += nodeProgress
+	} else if status == "uploading" {
+		// Add 50% of one node's progress (so uploading shows as halfway through that node)
+		nodeProgress := 100.0 / float64(total) * 0.5
+		baseProgress += nodeProgress
+	}
+
+	// Ensure we don't exceed 100%
+	if baseProgress > 100 {
+		baseProgress = 100
+	}
+
+	return int(baseProgress)
+}
+
 // getCurrentClusterMembers fetches real cluster data from SPORE nodes
 func (wss *WebSocketServer) getCurrentClusterMembers() ([]client.ClusterMember, error) {
 	nodes := wss.nodeDiscovery.GetNodes()