From ff7f2e4f733b520aa2b2625cede08883cd8bbc76 Mon Sep 17 00:00:00 2001
From: Jacky <me@jackyu.cn>
Date: Mon, 21 Apr 2025 13:57:49 +0000
Subject: [PATCH] feat(ota): enhance Docker upgrade process with progress
 tracking

---
 .devcontainer/Dockerfile    |  2 +-
 internal/docker/ota.go      | 56 +++++++++++++++++++++++++++++++++++--
 internal/upgrader/binary.go |  1 +
 internal/upgrader/docker.go | 44 +++++++++++++++++++++++------
 4 files changed, 91 insertions(+), 12 deletions(-)

diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index d7bc0643..4f023ddc 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,4 +1,4 @@
-FROM mcr.microsoft.com/devcontainers/base:jammy
+FROM mcr.microsoft.com/devcontainers/base:noble
 
 # Combine installation steps for Nginx and Go to avoid repetitive update/cleanup commands
 RUN apt-get update && \
diff --git a/internal/docker/ota.go b/internal/docker/ota.go
index 632326b7..c32f3447 100644
--- a/internal/docker/ota.go
+++ b/internal/docker/ota.go
@@ -2,6 +2,7 @@ package docker
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"os"
@@ -54,7 +55,7 @@ func removeAllTempContainers(ctx context.Context, cli *client.Client) (err error
 }
 
 // UpgradeStepOne Trigger in the OTA upgrade
-func UpgradeStepOne(channel string) (err error) {
+func UpgradeStepOne(channel string, progressChan chan<- float64) (err error) {
 	ctx := context.Background()
 
 	// 1. Get the tag of the latest release
@@ -78,8 +79,57 @@ func UpgradeStepOne(channel string) (err error) {
 	}
 	defer out.Close()
 
-	// Wait for pull to complete by reading the output
-	io.Copy(os.Stdout, out)
+	// Parse JSON stream and send progress updates through channel
+	decoder := json.NewDecoder(out)
+	type ProgressDetail struct {
+		Current int64 `json:"current"`
+		Total   int64 `json:"total"`
+	}
+	type PullStatus struct {
+		Status         string         `json:"status"`
+		ProgressDetail ProgressDetail `json:"progressDetail"`
+		ID             string         `json:"id"`
+	}
+
+	layers := make(map[string]float64)
+	var status PullStatus
+	var lastProgress float64
+
+	for {
+		if err := decoder.Decode(&status); err != nil {
+			if err == io.EOF {
+				break
+			}
+			logger.Error("Error decoding Docker pull status:", err)
+			continue
+		}
+
+		// Only process layers with progress information
+		if status.ProgressDetail.Total > 0 {
+			progress := float64(status.ProgressDetail.Current) / float64(status.ProgressDetail.Total) * 100
+			layers[status.ID] = progress
+
+			// Calculate overall progress (average of all layers)
+			var totalProgress float64
+			for _, p := range layers {
+				totalProgress += p
+			}
+			overallProgress := totalProgress / float64(len(layers))
+
+			// Only send progress updates when there's a meaningful change
+			if overallProgress > lastProgress+1 || overallProgress >= 100 {
+				if progressChan != nil {
+					progressChan <- overallProgress
+				}
+				lastProgress = overallProgress
+			}
+		}
+	}
+
+	// Ensure we send 100% at the end
+	if progressChan != nil && lastProgress < 100 {
+		progressChan <- 100
+	}
 
 	// 3. Create a temp container
 	// Clean up any existing temp containers
diff --git a/internal/upgrader/binary.go b/internal/upgrader/binary.go
index 9c573b4a..92b8f137 100644
--- a/internal/upgrader/binary.go
+++ b/internal/upgrader/binary.go
@@ -13,6 +13,7 @@ type Control struct {
 	Channel string `json:"channel"`
 }
 
+// BinaryUpgrade Upgrade the binary
 func BinaryUpgrade(ws *websocket.Conn, control *Control) {
 	_ = ws.WriteJSON(CoreUpgradeResp{
 		Status:  UpgradeStatusInfo,
diff --git a/internal/upgrader/docker.go b/internal/upgrader/docker.go
index 5755ff2d..6eb78f48 100644
--- a/internal/upgrader/docker.go
+++ b/internal/upgrader/docker.go
@@ -6,14 +6,42 @@ import (
 	"github.com/uozi-tech/cosy/logger"
 )
 
+// DockerUpgrade Upgrade the Docker container
 func DockerUpgrade(ws *websocket.Conn, control *Control) {
-	err := docker.UpgradeStepOne(control.Channel)
-	if err != nil {
-		_ = ws.WriteJSON(CoreUpgradeResp{
-			Status:  UpgradeStatusError,
-			Message: err.Error(),
-		})
-		logger.Error(err)
-		return
+	progressChan := make(chan float64)
+
+	// Start a goroutine to listen for progress updates and send them via WebSocket
+	go func() {
+		for progress := range progressChan {
+			err := ws.WriteJSON(CoreUpgradeResp{
+				Status:   UpgradeStatusProgress,
+				Progress: progress,
+				Message:  "Pulling Docker image...",
+			})
+			if err != nil {
+				logger.Error("Failed to send progress update:", err)
+				return
+			}
+		}
+	}()
+	defer close(progressChan)
+
+	if !control.DryRun {
+		err := docker.UpgradeStepOne(control.Channel, progressChan)
+		if err != nil {
+			_ = ws.WriteJSON(CoreUpgradeResp{
+				Status:  UpgradeStatusError,
+				Message: err.Error(),
+			})
+			logger.Error(err)
+			return
+		}
 	}
+
+	// Send completion message
+	_ = ws.WriteJSON(CoreUpgradeResp{
+		Status:   UpgradeStatusInfo,
+		Progress: 100,
+		Message:  "Docker image pull completed, upgrading...",
+	})
 }