chore: auto update geo

MetaCubeX · May 19, 2024 · 094537b · 094537b
1 parent e749c7e
commit 094537b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 60 deletions.
diff --git a/component/updater/update_geo.go b/component/updater/update_geo.go
@@ -98,7 +98,7 @@ func updateGeoDatabases() error {
 	return nil
 }
 
-func UpdateGeoDatabases(updateNotification chan struct{}) error {
+func UpdateGeoDatabases(onSuccess func()) error {
 	log.Infoln("[GEO] Start updating GEO database")
 
 	if UpdatingGeo.Load() {
@@ -115,7 +115,7 @@ func UpdateGeoDatabases(updateNotification chan struct{}) error {
 		return err
 	}
 
-	updateNotification <- struct{}{}
+	onSuccess()
 	return nil
 }
 
@@ -136,17 +136,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,14 +155,15 @@ 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(onSuccess); err != nil {
 				log.Errorln("[GEO] Failed to update GEO database: %s", err.Error())
 				return
 			}
 		}
 
 		for range ticker.C {
-			if err := UpdateGeoDatabases(updateNotification); err != nil {
+			log.Infoln("[GEO] updating database every %d hours", C.GeoUpdateInterval)
+			if err := UpdateGeoDatabases(onSuccess); err != nil {
 				log.Errorln("[GEO] Failed to update GEO database: %s", err.Error())
 			}
 		}

diff --git a/hub/route/configs.go b/hub/route/configs.go
@@ -364,47 +364,30 @@ 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)
-	}()
-
-	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
-				}
-
-				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
-			}
-		}
-	}()
-
-	go func() {
-		err := updater.UpdateGeoDatabases(updateNotification)
+	if updater.UpdatingGeo.Load() {
+		render.Status(r, http.StatusConflict)
+		render.JSON(w, r, newError("GEO database is updating, skip"))
+		return
+	}
+
+	onSuccess := func() {
+		cfg, err := executor.ParseWithPath(C.Path.Config())
 		if err != nil {
-			errorChannel <- err
+			log.Errorln("[GEO] update GEO databases failed: %v", err)
+			return
 		}
-	}()
 
-	<-done
+		log.Warnln("[GEO] update GEO databases success, applying config")
+
+		executor.ApplyConfig(cfg, false)
+	}
+
+	err := updater.UpdateGeoDatabases(onSuccess)
+	if err != nil {
+		render.Status(r, http.StatusInternalServerError)
+		render.JSON(w, r, newError(err.Error()))
+		return
+	}
 
 	render.NoContent(w, r)
 }
diff --git a/main.go b/main.go
@@ -113,23 +113,21 @@ func main() {
 	}
 
 	if C.GeoAutoUpdate {
-		updateNotification := make(chan struct{})
-		go updater.RegisterGeoUpdater(updateNotification)
+		onSuccess := func() {
+			cfg, err := executor.ParseWithPath(C.Path.Config())
+			if err != nil {
+				log.Errorln("[GEO] update GEO databases failed: %v", err)
+				return
+			}
 
-		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")
 
-				log.Warnln("[GEO] update GEO databases success, applying config")
+			executor.ApplyConfig(cfg, false)
+		}
 
-				executor.ApplyConfig(cfg, false)
-			}
-		}()
+		updater.RegisterGeoUpdater(onSuccess)
 	}
+
 	defer executor.Shutdown()
 
 	termSign := make(chan os.Signal, 1)