From 1bc41e7e58e99e23cd8be522cc1c2d30090975d2 Mon Sep 17 00:00:00 2001
From: Yuhang Wei <weiyuhang3@huawei.com>
Date: Tue, 15 Aug 2023 16:33:11 +0800
Subject: [PATCH 13/17] KubeOS: fix clean space problems

add clean space function in server
change docker rm position in docker_image

Signed-off-by: Yuhang Wei <weiyuhang3@huawei.com>
---
 cmd/agent/main.go                |  4 +---
 cmd/agent/server/docker_image.go |  6 +++---
 cmd/agent/server/server.go       |  9 +++++++++
 cmd/agent/server/utils.go        | 27 +++++++++++++--------------
 4 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 67c7e2d..7b1ed3d 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -13,8 +13,6 @@
 package main
 
 import (
-	"fmt"
-
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 
@@ -24,7 +22,7 @@ import (
 )
 
 func main() {
-	fmt.Println("Version is:", version.Version)
+	logrus.Infoln("Version is:", version.Version)
 	l, err := server.NewListener(server.SockDir, server.SockName)
 	if err != nil {
 		logrus.Errorln("listen error" + err.Error())
diff --git a/cmd/agent/server/docker_image.go b/cmd/agent/server/docker_image.go
index 0b6ee35..16bcea5 100644
--- a/cmd/agent/server/docker_image.go
+++ b/cmd/agent/server/docker_image.go
@@ -61,13 +61,13 @@ func (d dockerImageHandler) getRootfsArchive(req *pb.UpdateRequest, neededPath p
 	if err != nil {
 		return "", err
 	}
-	if err := runCommand("docker", "cp", containerId+":/"+neededPath.rootfsFile, neededPath.updatePath); err != nil {
-		return "", err
-	}
 	defer func() {
 		if err := runCommand("docker", "rm", containerId); err != nil {
 			logrus.Errorln("remove kubeos-temp container error", err)
 		}
 	}()
+	if err := runCommand("docker", "cp", containerId+":/"+neededPath.rootfsFile, neededPath.updatePath); err != nil {
+		return "", err
+	}
 	return neededPath.tarPath, nil
 }
diff --git a/cmd/agent/server/server.go b/cmd/agent/server/server.go
index 8ac6ffd..f8cbb41 100644
--- a/cmd/agent/server/server.go
+++ b/cmd/agent/server/server.go
@@ -112,6 +112,15 @@ func (s *Server) update(req *pb.UpdateRequest) error {
 		return fmt.Errorf("image type %s cannot be recognized", action)
 	}
 	imagePath, err := handler.downloadImage(req)
+	defer func() {
+		if err != nil {
+			path := newPreparePath()
+			if err := cleanSpace(path.updatePath, path.mountPath, path.imagePath); err != nil {
+				logrus.Errorln("clean space error " + err.Error())
+			}
+			logrus.Infoln("clean space success")
+		}
+	}()
 	if err != nil {
 		return err
 	}
diff --git a/cmd/agent/server/utils.go b/cmd/agent/server/utils.go
index b42db18..d2d0946 100644
--- a/cmd/agent/server/utils.go
+++ b/cmd/agent/server/utils.go
@@ -190,26 +190,25 @@ func prepareEnv() (preparePath, error) {
 	if err := checkDiskSize(needGBSize, PersistDir); err != nil {
 		return preparePath{}, err
 	}
-	rootfsFile := rootfsArchive
-	updatePath := splicePath(PersistDir, updateDir)
-	mountPath := splicePath(updatePath, mountDir)
-	tarPath := splicePath(updatePath, rootfsFile)
-	imagePath := splicePath(PersistDir, osImageName)
-
-	if err := cleanSpace(updatePath, mountPath, imagePath); err != nil {
+	upgradePath := newPreparePath()
+	if err := cleanSpace(upgradePath.updatePath, upgradePath.mountPath, upgradePath.imagePath); err != nil {
 		return preparePath{}, err
 	}
-	if err := os.MkdirAll(mountPath, imgPermission); err != nil {
+	if err := os.MkdirAll(upgradePath.mountPath, imgPermission); err != nil {
 		return preparePath{}, err
 	}
-	upgradePath := preparePath{
+	return upgradePath, nil
+}
+
+func newPreparePath() preparePath {
+	updatePath := splicePath(PersistDir, updateDir)
+	return preparePath{
 		updatePath: updatePath,
-		mountPath:  mountPath,
-		tarPath:    tarPath,
-		imagePath:  imagePath,
-		rootfsFile: rootfsFile,
+		mountPath:  splicePath(updatePath, mountDir),
+		tarPath:    splicePath(updatePath, rootfsArchive),
+		imagePath:  splicePath(PersistDir, osImageName),
+		rootfsFile: rootfsArchive,
 	}
-	return upgradePath, nil
 }
 
 func checkDiskSize(needGBSize int, path string) error {
-- 
2.39.0