chore: update geo in a batch

This commit is contained in:
wwqgtxx 2024-09-09 10:05:38 +08:00
parent 595a575cde
commit ef244b896a
2 changed files with 95 additions and 58 deletions

View File

@ -209,8 +209,11 @@ func LoadGeoIPMatcher(country string) (router.IPMatcher, error) {
return matcher, nil return matcher, nil
} }
func ClearCache() { func ClearGeoSiteCache() {
loadGeoSiteMatcherListSF.Reset() loadGeoSiteMatcherListSF.Reset()
loadGeoSiteMatcherSF.Reset() loadGeoSiteMatcherSF.Reset()
}
func ClearGeoIPCache() {
loadGeoIPMatcherSF.Reset() loadGeoIPMatcherSF.Reset()
} }

View File

@ -1,6 +1,7 @@
package updater package updater
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"os" "os"
@ -8,6 +9,7 @@ import (
"time" "time"
"github.com/metacubex/mihomo/common/atomic" "github.com/metacubex/mihomo/common/atomic"
"github.com/metacubex/mihomo/common/batch"
"github.com/metacubex/mihomo/component/geodata" "github.com/metacubex/mihomo/component/geodata"
_ "github.com/metacubex/mihomo/component/geodata/standard" _ "github.com/metacubex/mihomo/component/geodata/standard"
"github.com/metacubex/mihomo/component/mmdb" "github.com/metacubex/mihomo/component/mmdb"
@ -18,37 +20,15 @@ import (
) )
var ( var (
UpdatingGeo atomic.Bool updatingGeo atomic.Bool
) )
func updateGeoDatabases() error { func UpdateMMDB() (err error) {
defer runtime.GC()
geoLoader, err := geodata.GetGeoDataLoader("standard")
if err != nil {
return err
}
if C.GeodataMode {
data, err := downloadForBytes(C.GeoIpUrl)
if err != nil {
return fmt.Errorf("can't download GeoIP database file: %w", err)
}
if _, err = geoLoader.LoadIPByBytes(data, "cn"); err != nil {
return fmt.Errorf("invalid GeoIP database file: %s", err)
}
if err = saveFile(data, C.Path.GeoIP()); err != nil {
return fmt.Errorf("can't save GeoIP database file: %w", err)
}
} else {
defer mmdb.ReloadIP() defer mmdb.ReloadIP()
data, err := downloadForBytes(C.MmdbUrl) data, err := downloadForBytes(C.MmdbUrl)
if err != nil { if err != nil {
return fmt.Errorf("can't download MMDB database file: %w", err) return fmt.Errorf("can't download MMDB database file: %w", err)
} }
instance, err := maxminddb.FromBytes(data) instance, err := maxminddb.FromBytes(data)
if err != nil { if err != nil {
return fmt.Errorf("invalid MMDB database file: %s", err) return fmt.Errorf("invalid MMDB database file: %s", err)
@ -59,9 +39,10 @@ func updateGeoDatabases() error {
if err = saveFile(data, C.Path.MMDB()); err != nil { if err = saveFile(data, C.Path.MMDB()); err != nil {
return fmt.Errorf("can't save MMDB database file: %w", err) return fmt.Errorf("can't save MMDB database file: %w", err)
} }
return nil
} }
if C.ASNEnable { func UpdateASN() (err error) {
defer mmdb.ReloadASN() defer mmdb.ReloadASN()
data, err := downloadForBytes(C.ASNUrl) data, err := downloadForBytes(C.ASNUrl)
if err != nil { if err != nil {
@ -74,12 +55,32 @@ func updateGeoDatabases() error {
} }
_ = instance.Close() _ = instance.Close()
mmdb.ASNInstance().Reader.Close() mmdb.ASNInstance().Reader.Close() // mmdb is loaded with mmap, so it needs to be closed before overwriting the file
if err = saveFile(data, C.Path.ASN()); err != nil { if err = saveFile(data, C.Path.ASN()); err != nil {
return fmt.Errorf("can't save ASN database file: %w", err) return fmt.Errorf("can't save ASN database file: %w", err)
} }
return nil
} }
func UpdateGeoIp() (err error) {
defer geodata.ClearGeoIPCache()
geoLoader, err := geodata.GetGeoDataLoader("standard")
data, err := downloadForBytes(C.GeoIpUrl)
if err != nil {
return fmt.Errorf("can't download GeoIP database file: %w", err)
}
if _, err = geoLoader.LoadIPByBytes(data, "cn"); err != nil {
return fmt.Errorf("invalid GeoIP database file: %s", err)
}
if err = saveFile(data, C.Path.GeoIP()); err != nil {
return fmt.Errorf("can't save GeoIP database file: %w", err)
}
return nil
}
func UpdateGeoSite() (err error) {
defer geodata.ClearGeoSiteCache()
geoLoader, err := geodata.GetGeoDataLoader("standard")
data, err := downloadForBytes(C.GeoSiteUrl) data, err := downloadForBytes(C.GeoSiteUrl)
if err != nil { if err != nil {
return fmt.Errorf("can't download GeoSite database file: %w", err) return fmt.Errorf("can't download GeoSite database file: %w", err)
@ -92,8 +93,41 @@ func updateGeoDatabases() error {
if err = saveFile(data, C.Path.GeoSite()); err != nil { if err = saveFile(data, C.Path.GeoSite()); err != nil {
return fmt.Errorf("can't save GeoSite database file: %w", err) return fmt.Errorf("can't save GeoSite database file: %w", err)
} }
return nil
}
geodata.ClearCache() func updateGeoDatabases() error {
defer runtime.GC()
b, _ := batch.New[interface{}](context.Background())
if C.GeodataMode {
b.Go("UpdateGeoIp", func() (_ interface{}, err error) {
err = UpdateGeoIp()
return
})
} else {
b.Go("UpdateMMDB", func() (_ interface{}, err error) {
err = UpdateMMDB()
return
})
}
if C.ASNEnable {
b.Go("UpdateASN", func() (_ interface{}, err error) {
err = UpdateASN()
return
})
}
b.Go("UpdateGeoSite", func() (_ interface{}, err error) {
err = UpdateGeoSite()
return
})
if e := b.Wait(); e != nil {
return e.Err
}
return nil return nil
} }
@ -103,12 +137,12 @@ var ErrGetDatabaseUpdateSkip = errors.New("GEO database is updating, skip")
func UpdateGeoDatabases() error { func UpdateGeoDatabases() error {
log.Infoln("[GEO] Start updating GEO database") log.Infoln("[GEO] Start updating GEO database")
if UpdatingGeo.Load() { if updatingGeo.Load() {
return ErrGetDatabaseUpdateSkip return ErrGetDatabaseUpdateSkip
} }
UpdatingGeo.Store(true) updatingGeo.Store(true)
defer UpdatingGeo.Store(false) defer updatingGeo.Store(false)
log.Infoln("[GEO] Updating GEO database") log.Infoln("[GEO] Updating GEO database")