From 4243a7428478b786bde56421d57cc15d857c3cfc Mon Sep 17 00:00:00 2001
From: Larvan2 <78135608+Larvan2@users.noreply.github.com>
Date: Sun, 19 May 2024 17:30:39 +0800
Subject: [PATCH] chore: auto update geo

---
 component/updater/update_geo.go | 25 +++++++++------
 hub/route/configs.go            | 54 ++++++++++-----------------------
 main.go                         | 26 +++++++---------
 3 files changed, 42 insertions(+), 63 deletions(-)

diff --git a/component/updater/update_geo.go b/component/updater/update_geo.go
index 4bb57e06..b07cd315 100644
--- a/component/updater/update_geo.go
+++ b/component/updater/update_geo.go
@@ -98,11 +98,13 @@ func updateGeoDatabases() error {
 	return nil
 }
 
-func UpdateGeoDatabases(updateNotification chan struct{}) error {
+var ErrGetDatabaseUpdateSkip = errors.New("GEO database is updating, skip")
+
+func UpdateGeoDatabases() error {
 	log.Infoln("[GEO] Start updating GEO database")
 
 	if UpdatingGeo.Load() {
-		return errors.New("GEO database is updating, skip")
+		return ErrGetDatabaseUpdateSkip
 	}
 
 	UpdatingGeo.Store(true)
@@ -115,7 +117,6 @@ func UpdateGeoDatabases(updateNotification chan struct{}) error {
 		return err
 	}
 
-	updateNotification <- struct{}{}
 	return nil
 }
 
@@ -136,17 +137,16 @@ func getUpdateTime() (err error, time time.Time) {
 	return nil, fileInfo.ModTime()
 }
 
-func RegisterGeoUpdater(updateNotification chan struct{}) {
+func RegisterGeoUpdater(onSuccess func()) {
 	if C.GeoUpdateInterval <= 0 {
 		log.Errorln("[GEO] Invalid update interval: %d", C.GeoUpdateInterval)
 		return
 	}
 
-	ticker := time.NewTicker(time.Duration(C.GeoUpdateInterval) * time.Hour)
-	defer ticker.Stop()
-
-	log.Infoln("[GEO] update GEO database every %d hours", C.GeoUpdateInterval)
 	go func() {
+		ticker := time.NewTicker(time.Duration(C.GeoUpdateInterval) * time.Hour)
+		defer ticker.Stop()
+
 		err, lastUpdate := getUpdateTime()
 		if err != nil {
 			log.Errorln("[GEO] Get GEO database update time error: %s", err.Error())
@@ -156,15 +156,20 @@ func RegisterGeoUpdater(updateNotification chan struct{}) {
 		log.Infoln("[GEO] last update time %s", lastUpdate)
 		if lastUpdate.Add(time.Duration(C.GeoUpdateInterval) * time.Hour).Before(time.Now()) {
 			log.Infoln("[GEO] Database has not been updated for %v, update now", time.Duration(C.GeoUpdateInterval)*time.Hour)
-			if err := UpdateGeoDatabases(updateNotification); err != nil {
+			if err := UpdateGeoDatabases(); err != nil {
 				log.Errorln("[GEO] Failed to update GEO database: %s", err.Error())
 				return
+			} else {
+				onSuccess()
 			}
 		}
 
 		for range ticker.C {
-			if err := UpdateGeoDatabases(updateNotification); err != nil {
+			log.Infoln("[GEO] updating database every %d hours", C.GeoUpdateInterval)
+			if err := UpdateGeoDatabases(); err != nil {
 				log.Errorln("[GEO] Failed to update GEO database: %s", err.Error())
+			} else {
+				onSuccess()
 			}
 		}
 	}()
diff --git a/hub/route/configs.go b/hub/route/configs.go
index 35977885..f2cf298a 100644
--- a/hub/route/configs.go
+++ b/hub/route/configs.go
@@ -364,47 +364,25 @@ func updateConfigs(w http.ResponseWriter, r *http.Request) {
 }
 
 func updateGeoDatabases(w http.ResponseWriter, r *http.Request) {
-	updateNotification := make(chan struct{})
-	errorChannel := make(chan error, 1)
-	done := make(chan struct{})
-	defer func() {
-		close(updateNotification)
-		close(errorChannel)
-	}()
+	err := updater.UpdateGeoDatabases()
+	if err != nil {
+		log.Errorln("[REST-API] update GEO databases failed: %v", err)
+		render.Status(r, http.StatusInternalServerError)
+		render.JSON(w, r, newError(err.Error()))
+		return
+	}
 
-	go func() {
-		defer close(done)
-		for {
-			select {
-			case <-updateNotification:
-				cfg, err := executor.ParseWithPath(C.Path.Config())
-				if err != nil {
-					log.Errorln("[REST-API] update GEO databases failed: %v", err)
-					render.Status(r, http.StatusInternalServerError)
-					render.JSON(w, r, newError("Error parsing configuration"))
-					return
-				}
+	cfg, err := executor.ParseWithPath(C.Path.Config())
+	if err != nil {
+		log.Errorln("[REST-API] update GEO databases failed: %v", err)
+		render.Status(r, http.StatusInternalServerError)
+		render.JSON(w, r, newError("Error parsing configuration"))
+		return
+	}
 
-				log.Warnln("[REST-API] update GEO databases success, applying config")
-				executor.ApplyConfig(cfg, false)
-				return
-			case err := <-errorChannel:
-				log.Errorln("[REST-API] update GEO databases failed: %v", err)
-				render.Status(r, http.StatusInternalServerError)
-				render.JSON(w, r, err.Error())
-				return
-			}
-		}
-	}()
+	log.Warnln("[GEO] update GEO databases success, applying config")
 
-	go func() {
-		err := updater.UpdateGeoDatabases(updateNotification)
-		if err != nil {
-			errorChannel <- err
-		}
-	}()
-
-	<-done
+	executor.ApplyConfig(cfg, false)
 
 	render.NoContent(w, r)
 }
diff --git a/main.go b/main.go
index 081a7eb5..61f1d683 100644
--- a/main.go
+++ b/main.go
@@ -113,23 +113,19 @@ func main() {
 	}
 
 	if C.GeoAutoUpdate {
-		updateNotification := make(chan struct{})
-		go updater.RegisterGeoUpdater(updateNotification)
-
-		go func() {
-			for range updateNotification {
-				cfg, err := executor.ParseWithPath(C.Path.Config())
-				if err != nil {
-					log.Errorln("[GEO] update GEO databases failed: %v", err)
-					return
-				}
-
-				log.Warnln("[GEO] update GEO databases success, applying config")
-
-				executor.ApplyConfig(cfg, false)
+		updater.RegisterGeoUpdater(func() {
+			cfg, err := executor.ParseWithPath(C.Path.Config())
+			if err != nil {
+				log.Errorln("[GEO] update GEO databases failed: %v", err)
+				return
 			}
-		}()
+
+			log.Warnln("[GEO] update GEO databases success, applying config")
+
+			executor.ApplyConfig(cfg, false)
+		})
 	}
+
 	defer executor.Shutdown()
 
 	termSign := make(chan os.Signal, 1)