From b9741c92fc93b51883bfbd5c526faa8b66f8724b Mon Sep 17 00:00:00 2001
From: mzz <2017@duck.com>
Date: Wed, 27 Mar 2024 13:01:37 +0800
Subject: [PATCH] feat: support reloading progress and error for `dae reload`
 (#470)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Sumire (菫) <151038614+sumire88@users.noreply.github.com>
---
 cmd/reload.go           | 46 +++++++++++++++++++++++++++++++++++++++++
 cmd/run.go              |  9 +++++++-
 common/consts/reload.go |  8 +++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 common/consts/reload.go

diff --git a/cmd/reload.go b/cmd/reload.go
index 0eadd6d..4f8815e 100644
--- a/cmd/reload.go
+++ b/cmd/reload.go
@@ -11,11 +11,27 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/daeuniverse/dae/cmd/internal"
+	"github.com/daeuniverse/dae/common/consts"
 	"github.com/spf13/cobra"
 )
 
+func readSignalProgressFile() (code byte, content string, err error) {
+	b, err := os.ReadFile(SignalProgressFilePath)
+	if err != nil {
+		return 0, "", err
+	}
+	var firstLine string
+	firstLine, content, _ = strings.Cut(string(b), "\n")
+	if len(firstLine) != 1 {
+		return 0, "", fmt.Errorf("unexpected format: %v", string(b))
+	}
+	code = firstLine[0]
+	return code, content, nil
+}
+
 var (
 	abort     bool
 	reloadCmd = &cobra.Command{
@@ -41,10 +57,40 @@ var (
 					f.Close()
 				}
 			}
+			// Read the first line of SignalProgressFilePath.
+			code, _, err := readSignalProgressFile()
+			if err == nil && code != consts.ReloadDone && code != consts.ReloadError {
+				// In progress.
+				fmt.Printf("%v shows another reload operation is in progress.\n", SignalProgressFilePath)
+				return
+			}
+			// Set the progress as ReloadSend.
+			os.WriteFile(SignalProgressFilePath, []byte{consts.ReloadSend}, 0644)
+			// Send signal.
 			if err = syscall.Kill(pid, syscall.SIGUSR1); err != nil {
 				fmt.Println(err)
 				os.Exit(1)
 			}
+			time.Sleep(500 * time.Millisecond)
+			code, _, _ = readSignalProgressFile()
+			if code == consts.ReloadSend {
+				// Old version dae is running.
+				goto fallback
+			}
+
+			for {
+				time.Sleep(200 * time.Millisecond)
+				code, content, err := readSignalProgressFile()
+				if err != nil {
+					// Unexpecetd case.
+					goto fallback
+				}
+				if code == consts.ReloadDone || code == consts.ReloadError {
+					fmt.Println(content)
+					return
+				}
+			}
+		fallback:
 			fmt.Println("OK")
 		},
 	}
diff --git a/cmd/run.go b/cmd/run.go
index 65b86fe..5480628 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -27,6 +27,7 @@ import (
 
 	"github.com/daeuniverse/dae/cmd/internal"
 	"github.com/daeuniverse/dae/common"
+	"github.com/daeuniverse/dae/common/consts"
 	"github.com/daeuniverse/dae/common/subscription"
 	"github.com/daeuniverse/dae/config"
 	"github.com/daeuniverse/dae/control"
@@ -39,7 +40,8 @@ import (
 )
 
 const (
-	PidFilePath = "/var/run/dae.pid"
+	PidFilePath            = "/var/run/dae.pid"
+	SignalProgressFilePath = "/var/run/dae.progress"
 )
 
 var (
@@ -134,6 +136,7 @@ func Run(log *logrus.Logger, conf *config.Config, externGeoDataDirs []string) (e
 			if !disablePidFile {
 				_ = os.WriteFile(PidFilePath, []byte(strconv.Itoa(os.Getpid())), 0644)
 			}
+			_ = os.WriteFile(SignalProgressFilePath, []byte{consts.ReloadDone}, 0644)
 		}()
 		control.GetDaeNetns().With(func() error {
 			if listener, err = c.ListenAndServe(readyChan, conf.Global.TproxyPort); err != nil {
@@ -167,6 +170,7 @@ loop:
 				}()
 				<-readyChan
 				sdnotify.Ready()
+				_ = os.WriteFile(SignalProgressFilePath, append([]byte{consts.ReloadDone}, []byte("\nOK")...), 0644)
 				log.Warnln("[Reload] Finished")
 			} else {
 				// Listening error.
@@ -183,6 +187,7 @@ loop:
 				log.Warnln("[Reload] Received reload signal; prepare to reload")
 			}
 			sdnotify.Reloading()
+			_ = os.WriteFile(SignalProgressFilePath, []byte{consts.ReloadProcessing}, 0644)
 
 			// Load new config.
 			abortConnections = os.Remove(AbortFile) == nil
@@ -196,6 +201,7 @@ loop:
 						"err": err,
 					}).Errorln("[Reload] Failed to reload")
 					sdnotify.Ready()
+					_ = os.WriteFile(SignalProgressFilePath, append([]byte{consts.ReloadError}, []byte("\n"+err.Error())...), 0644)
 					continue
 				}
 				newConf.Global = deepcopy.Copy(conf.Global).(config.Global)
@@ -210,6 +216,7 @@ loop:
 						"err": err,
 					}).Errorln("[Reload] Failed to reload")
 					sdnotify.Ready()
+					_ = os.WriteFile(SignalProgressFilePath, append([]byte{consts.ReloadError}, []byte("\n"+err.Error())...), 0644)
 					continue
 				}
 				log.Infof("Include config files: [%v]", strings.Join(includes, ", "))
diff --git a/common/consts/reload.go b/common/consts/reload.go
new file mode 100644
index 0000000..1284881
--- /dev/null
+++ b/common/consts/reload.go
@@ -0,0 +1,8 @@
+package consts
+
+const (
+	ReloadSend = '0' + iota
+	ReloadProcessing
+	ReloadDone
+	ReloadError
+)