From e2c8b68b11e58762495113e11212ea8ef12d4a85 Mon Sep 17 00:00:00 2001
From: zyxkad <zyxkad@gmail.com>
Date: Wed, 13 Dec 2023 19:59:41 -0700
Subject: [PATCH] fix shutdown process

---
 src/cluster.go | 39 +++++++++++++++++++++------------------
 src/main.go    |  5 ++---
 2 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/cluster.go b/src/cluster.go
index ecb6d4dd..70197abd 100644
--- a/src/cluster.go
+++ b/src/cluster.go
@@ -178,9 +178,20 @@ func (cr *Cluster) Enable() (err error) {
 	cr.enabled = true
 
 	var keepaliveCtx context.Context
-	keepaliveCtx, cr.keepalive = context.WithCancel(context.TODO())
+	keepaliveCtx, cr.keepalive = context.WithCancel(cr.ctx)
 	createInterval(keepaliveCtx, func() {
-		cr.KeepAlive()
+		if !cr.KeepAlive() {
+			logDebug("Reconnecting due to keepalive failed")
+			cr.Disable()
+			if !cr.Connect() {
+				logError("Cannot reconnect to server, exit.")
+				os.Exit(1)
+			}
+			if err := cr.Enable(); err != nil {
+				logError("Cannot enable cluster:", err, "; exit.")
+				os.Exit(1)
+			}
+		}
 	}, KeepAliveInterval)
 	return
 }
@@ -196,27 +207,19 @@ func (cr *Cluster) KeepAlive() (ok bool) {
 		logError("Error when keep-alive:", err)
 		return false
 	}
-	if len(data) > 1 && data.Get(0) == nil {
-		logInfo("Keep-alive success:", hits, bytesToUnit((float64)(hbytes)), data)
-	} else {
-		logInfo("Keep-alive failed:", data.Get(0))
-		cr.Disable()
-		if !cr.Connect() {
-			logError("Cannot reconnect to server, exit.")
-			os.Exit(1)
-		}
-		if err := cr.Enable(); err != nil {
-			logError("Cannot enable cluster:", err, "; exit.")
-			os.Exit(1)
-		}
+	if erro := data.Get(0); len(data) <= 1 || erro != nil {
+		logError("Keep-alive failed:", erro)
+		return false
 	}
+	logInfo("Keep-alive success:", hits, bytesToUnit((float64)(hbytes)), data)
 	return true
 }
 
-func (cr *Cluster) Disable() (successed bool) {
+func (cr *Cluster) Disable() (ok bool) {
 	cr.mux.Lock()
 	defer cr.mux.Unlock()
 
+	cr.KeepAlive()
 	if !cr.enabled {
 		logDebug("Extra disable")
 		return true
@@ -236,7 +239,7 @@ func (cr *Cluster) Disable() (successed bool) {
 	if err != nil {
 		return false
 	}
-	logInfo("disable ack:", data)
+	logDebug("disable ack:", data)
 	if !data.GetBool(1) {
 		logError("Disable failed: " + data.String())
 		return false
@@ -270,7 +273,7 @@ func (pair *CertKeyPair) SaveAsFile() (cert, key string, err error) {
 }
 
 func (cr *Cluster) RequestCert() (ckp *CertKeyPair, err error) {
-	logInfo("Requesting cert, please wait ...")
+	logInfo("Requesting certificates, please wait ...")
 	data, err := cr.socket.EmitAck("request-cert")
 	if err != nil {
 		return
diff --git a/src/main.go b/src/main.go
index 4930eab2..37c8ed4f 100644
--- a/src/main.go
+++ b/src/main.go
@@ -256,14 +256,13 @@ START:
 	case <-exitCh:
 		return
 	case s := <-signalCh:
-		shutCtx, cancelShut := context.WithTimeout(bgctx, 16*time.Second)
+		shutCtx, _ := context.WithTimeout(ctx, 16*time.Second)
 		logWarn("Closing server ...")
-		cancel()
 		shutExit := make(chan struct{}, 0)
 		go hjServer.Shutdown(shutCtx)
 		go func() {
 			defer close(shutExit)
-			defer cancelShut()
+			defer cancel()
 			cluster.Disable()
 			cluster.Server.Shutdown(shutCtx)
 		}()