From 9a54e15ae9e67a2bdf2e7413d48303659269dd33 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Mon, 16 Jan 2023 15:15:44 +0800
Subject: [PATCH] [optimize] close the DB connections when reloading

---
 main_test.go       |  9 +++++----
 service/service.go |  1 +
 share/db.go        | 24 ++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/main_test.go b/main_test.go
index 6cb2572c..f54d9659 100644
--- a/main_test.go
+++ b/main_test.go
@@ -44,6 +44,7 @@ func TestCommandStart(t *testing.T) {
 		service.Stop(func() {})
 		log.Println("服务已关闭")
 	}()
+
 	go func() {
 		os.Args = append(os.Args, "start")
 		main()
@@ -68,9 +69,9 @@ func TestCommandStart(t *testing.T) {
 	}
 
 	// 等待服务启动
-	time.Sleep(time.Second * 2)
 	times := 0
-	for times < 20 { // 2秒超时
+	for times < 30 { // 2秒超时
+		time.Sleep(time.Second * 2)
 		times++
 		fmt.Printf("Trying(%d)...", times)
 		res, err := request()
@@ -115,9 +116,9 @@ func TestCommandStop(t *testing.T) {
 	}
 
 	// 等待服务启动
-	time.Sleep(time.Second * 2)
 	times := 0
-	for times < 20 { // 2秒超时
+	for times < 30 { // 2秒超时
+		time.Sleep(time.Second * 2)
 		times++
 		res, err := request()
 		if err != nil {
diff --git a/service/service.go b/service/service.go
index bc54eb4c..b5f24f00 100644
--- a/service/service.go
+++ b/service/service.go
@@ -69,6 +69,7 @@ func Stop(onComplete func()) {
 	select {
 	case <-shutdownComplete:
 		share.SessionStop()
+		share.DBClose()
 		onComplete()
 	}
 }
diff --git a/share/db.go b/share/db.go
index b5719c40..93992833 100644
--- a/share/db.go
+++ b/share/db.go
@@ -2,6 +2,7 @@ package share
 
 import (
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/yaoapp/kun/log"
@@ -45,3 +46,26 @@ func DBConnect(dbconfig config.DBConfig) (err error) {
 
 	return err
 }
+
+// DBClose close the database connections
+func DBClose() error {
+	messages := []string{}
+	capsule.Global.Connections.Range(func(key, value any) bool {
+		log.Trace("[DBClose] %s", key)
+		if conn, ok := value.(*capsule.Connection); ok {
+			err := conn.Close()
+			if err != nil {
+				messages = append(messages, err.Error())
+			}
+		}
+		return true
+	})
+
+	if len(messages) > 0 {
+		msg := fmt.Sprintf("[DBClose] %s ", strings.Join(messages, ";"))
+		log.Error(msg)
+		return fmt.Errorf(msg)
+	}
+
+	return nil
+}