From ce9611ea818efbc5601130d13a5411a83858e70b Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Thu, 15 Apr 2021 11:16:58 +0800
Subject: [PATCH] =?UTF-8?q?=E5=90=AF=E5=8A=A8=E6=97=B6=E5=A6=82=E6=9E=9C?=
 =?UTF-8?q?=E6=97=A0=E6=B3=95=E8=BF=9E=E6=8E=A5=E6=95=B0=E6=8D=AE=E5=BA=93?=
 =?UTF-8?q?=EF=BC=8C=E5=88=99=E5=B0=9D=E8=AF=95N=E5=88=86=E9=92=9F?=
 =?UTF-8?q?=EF=BC=8C=E9=98=B2=E6=AD=A2=E5=9B=A0=E4=B8=BA=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=BA=93=E8=BF=9E=E6=8E=A5=E5=A4=B1=E8=B4=A5=E8=80=8C=E5=BC=82?=
 =?UTF-8?q?=E5=B8=B8=E9=80=80=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/api_node.go | 43 +++++++++++++++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/internal/nodes/api_node.go b/internal/nodes/api_node.go
index 85afdb1f..829a86cc 100644
--- a/internal/nodes/api_node.go
+++ b/internal/nodes/api_node.go
@@ -3,6 +3,7 @@ package nodes
 import (
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"github.com/TeaOSLab/EdgeAPI/internal/configs"
 	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
@@ -43,8 +44,15 @@ func NewAPINode() *APINode {
 func (this *APINode) Start() {
 	logs.Println("[API_NODE]start api node, pid: " + strconv.Itoa(os.Getpid()))
 
+	// 检查数据库连接
+	err := this.checkDB()
+	if err != nil {
+		logs.Println("[API_NODE]" + err.Error())
+		return
+	}
+
 	// 本地Sock
-	err := this.listenSock()
+	err = this.listenSock()
 	if err != nil {
 		logs.Println("[API_NODE]" + err.Error())
 		return
@@ -108,7 +116,7 @@ func (this *APINode) Start() {
 	select {}
 }
 
-// 实现守护进程
+// Daemon 实现守护进程
 func (this *APINode) Daemon() {
 	path := os.TempDir() + "/edge-api.sock"
 	isDebug := lists.ContainsString(os.Args, "debug")
@@ -153,7 +161,7 @@ func (this *APINode) Daemon() {
 	}
 }
 
-// 安装系统服务
+// InstallSystemService 安装系统服务
 func (this *APINode) InstallSystemService() error {
 	shortName := teaconst.SystemdServiceName
 
@@ -263,6 +271,35 @@ func (this *APINode) listenRPC(listener net.Listener, tlsConfig *tls.Config) err
 	return nil
 }
 
+// 检查数据库
+func (this *APINode) checkDB() error {
+	logs.Println("checking database connection ...")
+
+	db, err := dbs.Default()
+	if err != nil {
+		return err
+	}
+
+	maxTries := 600
+	for i := 0; i <= maxTries; i++ {
+		_, err := db.Exec("SELECT 1")
+		if err != nil {
+			if i == maxTries-1 {
+				return err
+			} else {
+				if i%10 == 0 { // 这让提示不会太多
+					logs.Println("[API_NODE]reconnecting to database (" + fmt.Sprintf("%.1f", float32(i * 100)/float32(maxTries+1)) + "%) ...")
+				}
+				time.Sleep(1 * time.Second)
+			}
+		} else {
+			return nil
+		}
+	}
+
+	return nil
+}
+
 // 自动升级
 func (this *APINode) autoUpgrade() error {
 	if Tea.IsTesting() {