diff --git a/config.go b/config.go new file mode 100644 index 0000000..fd23af8 --- /dev/null +++ b/config.go @@ -0,0 +1,206 @@ +package main + +import ( + lru "github.com/hashicorp/golang-lru" + "github.com/oschwald/maxminddb-golang" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + log "github.com/sirupsen/logrus" + "github.com/spf13/viper" + "net" + "net/url" + "strings" + "sync" +) + +func reloadConfig() { + err := viper.ReadInConfig() // Find and read the config file + + if err != nil { // Handle errors reading the config file + log.WithError(err).Fatalln("Unable to load config file") + } + + // db will never be reloaded. + if db == nil { + // Load maxmind database + db, err = maxminddb.Open(viper.GetString("geodb")) + + if err != nil { + log.WithError(err).Fatalln("Unable to open database") + } + } + + // Refresh server cache if size changed + if serverCache == nil { + serverCache, err = lru.New(viper.GetInt("cacheSize")) + } else { + serverCache.Resize(viper.GetInt("cacheSize")) + } + + // Set top choice count + topChoices = viper.GetInt("topChoices") + + // Reload map file + reloadMap() + + // Reload server list + reloadServers() + + // Check top choices size + if topChoices > len(servers) { + topChoices = len(servers) + } + + // Force check + go servers.Check() +} + +func reloadServers() { + var serverList []ServerConfig + viper.UnmarshalKey("servers", &serverList) + + var wg sync.WaitGroup + + existing := make(map[string]int) + + for i, server := range servers { + existing[server.Host] = i + } + + hosts := make(map[string]bool) + + for _, server := range serverList { + wg.Add(1) + + var prefix string + + if !strings.HasPrefix(server.Server, "http") { + prefix = "https://" + } + + u, err := url.Parse(prefix + server.Server) + + if err != nil { + log.WithFields(log.Fields{ + "error": err, + "server": server, + }).Warning("Server is invalid") + return + } + + hosts[u.Host] = true + + i := -1 + + if v, exists := existing[u.Host]; exists { + i = v + } + + go func(i int, server ServerConfig, u *url.URL) { + defer wg.Done() + + s := addServer(server, u) + + if _, ok := existing[u.Host]; ok { + s.Redirects = servers[i].Redirects + + servers[i] = s + } else { + s.Redirects = promauto.NewCounter(prometheus.CounterOpts{ + Name: "armbian_router_redirects_" + metricReplacer.Replace(u.Host), + Help: "The number of redirects for server " + u.Host, + }) + + servers = append(servers, s) + + log.WithFields(log.Fields{ + "server": u.Host, + "path": u.Path, + "latitude": s.Latitude, + "longitude": s.Longitude, + }).Info("Added server") + } + }(i, server, u) + } + + wg.Wait() + + // Remove servers that no longer exist in the config + for i := len(servers) - 1; i >= 0; i-- { + if _, exists := hosts[servers[i].Host]; exists { + continue + } + + log.WithFields(log.Fields{ + "server": servers[i].Host, + }).Info("Removed server") + + servers = append(servers[:i], servers[i+1:]...) + } +} + +var metricReplacer = strings.NewReplacer(".", "_", "-", "_") + +// addServer takes ServerConfig and constructs a server. +// This will create duplicate servers, but it will overwrite existing ones when changed. +func addServer(server ServerConfig, u *url.URL) *Server { + s := &Server{ + Available: true, + Host: u.Host, + Path: u.Path, + Latitude: server.Latitude, + Longitude: server.Longitude, + Weight: server.Weight, + } + + if s.Weight == 0 { + s.Weight = 1 + } + + if s.Latitude == 0 && s.Longitude == 0 { + ips, err := net.LookupIP(u.Host) + + if err != nil { + log.WithFields(log.Fields{ + "error": err, + "server": s.Host, + }).Warning("Could not resolve address") + return nil + } + + var city City + err = db.Lookup(ips[0], &city) + + if err != nil { + log.WithFields(log.Fields{ + "error": err, + "server": s.Host, + "ip": ips[0], + }).Warning("Could not geolocate address") + return nil + } + + s.Latitude = city.Location.Latitude + s.Longitude = city.Location.Longitude + } + + return s +} + +func reloadMap() { + mapFile := viper.GetString("dl_map") + + if mapFile == "" { + return + } + + log.WithField("file", mapFile).Info("Loading download map") + + newMap, err := loadMap(mapFile) + + if err != nil { + return + } + + dlMap = newMap +} diff --git a/dlrouter.yaml b/dlrouter.yaml index 0e73de0..41431af 100644 --- a/dlrouter.yaml +++ b/dlrouter.yaml @@ -1,16 +1,43 @@ -geodb: GeoIP2-City.mmdb +# GeoIP Database Path +geodb: GeoLite2-City.mmdb +# LRU Cache Size (in items) +cacheSize: 1024 + +# Server definition +# Weights are just like nginx, where if it's > 1 it'll be chosen x out of x + total times +# By default, the top 3 servers are used for choosing the best. +# server = full url or host+path +# weight = int +# optional: latitude, longitude (float) servers: - - https://mirrors.tuna.tsinghua.edu.cn/armbian-releases/ - - https://mirrors.bfsu.edu.cn/armbian-releases/ - - https://mirrors.nju.edu.cn/armbian-releases/ - - https://mirrors.ustc.edu.cn/armbian-dl/ - - https://mirror.12z.eu/pub/linux/armbian/dl/ - - https://armbian.tnahosting.net/dl/ - - https://stpete-mirror.armbian.com/dl/ - - https://mirror.armbian.de/dl/ - - https://mirrors.netix.net/armbian/dl/ - - https://mirrors.dotsrc.org/armbian-dl/ - - https://armbian.hosthatch.com/dl/ - - https://xogium.performanceservers.nl/dl/ - - https://github.com/armbian/mirror/releases/download/ + - server: armbian.12z.eu/apt/ + - server: armbian.chi.auroradev.org/apt/ + weight: 5 + latitude: 41.8879 + longitude: -88.1995 + - server: armbian.hosthatch.com/apt/ + - server: armbian.lv.auroradev.org/apt/ + weight: 5 + - server: armbian.site-meganet.com/apt/ + - server: armbian.systemonachip.net/apt/ + - server: armbian.tnahosting.net/apt/ + weight: 5 + - server: au-mirror.bret.dk/armbian/apt/ + - server: es-mirror.bret.dk/armbian/apt/ + - server: imola.armbian.com/apt/ + - server: mirror.iscas.ac.cn/armbian/ + - server: mirror.sjtu.edu.cn/armbian/ + - server: mirrors.aliyun.com/armbian/ + - server: mirrors.bfsu.edu.cn/armbian/ + - server: mirrors.dotsrc.org/armbian-apt/ + weight: 5 + - server: mirrors.netix.net/armbian/apt/ + - server: mirrors.nju.edu.cn/armbian/ + - server: mirrors.sustech.edu.cn/armbian/ + - server: mirrors.tuna.tsinghua.edu.cn/armbian/ + - server: mirrors.ustc.edu.cn/armbian/ + - server: mirrors.xtom.de/armbian/ + - server: sg-mirror.bret.dk/armbian/apt/ + - server: stpete-mirror.armbian.com/apt/ + - server: xogium.performanceservers.nl/apt/ diff --git a/go.mod b/go.mod index a5b1731..38c545f 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,8 @@ go 1.17 require ( github.com/chi-middleware/logrus-logger v0.2.0 github.com/go-chi/chi/v5 v5.0.7 + github.com/hashicorp/golang-lru v0.5.4 + github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff github.com/oschwald/maxminddb-golang v1.8.0 github.com/prometheus/client_golang v1.11.0 github.com/sirupsen/logrus v1.8.1 diff --git a/go.sum b/go.sum index 3c22734..ce8227a 100644 --- a/go.sum +++ b/go.sum @@ -224,6 +224,7 @@ github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/b github.com/hashicorp/go-uuid v1.0.1/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8= github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8= +github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc= github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= @@ -234,6 +235,8 @@ github.com/hashicorp/serf v0.9.6/go.mod h1:TXZNMjZQijwlDvp+r0b63xZ45H7JmCmgg4gpT github.com/iancoleman/strcase v0.2.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47ZCWhYzw7ho= github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= +github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff h1:6NvhExg4omUC9NfA+l4Oq3ibNNeJUdiAF3iBVB0PlDk= +github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff/go.mod h1:ddfPX8Z28YMjiqoaJhNBzWHapTHXejnB5cDCUWDwriw= github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX5e0EB2j4= github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU= github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= diff --git a/main.go b/main.go index 5a50620..1afad98 100644 --- a/main.go +++ b/main.go @@ -4,19 +4,16 @@ import ( "flag" "github.com/chi-middleware/logrus-logger" "github.com/go-chi/chi/v5" + lru "github.com/hashicorp/golang-lru" "github.com/oschwald/maxminddb-golang" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" "github.com/prometheus/client_golang/prometheus/promhttp" log "github.com/sirupsen/logrus" "github.com/spf13/viper" - "net" "net/http" - "net/url" "os" "os/signal" - "strings" - "sync" "syscall" ) @@ -35,6 +32,10 @@ var ( Name: "armbian_router_download_maps", Help: "The total number of mapped download paths", }) + + serverCache *lru.Cache + + topChoices int ) // City represents a MaxmindDB city @@ -45,6 +46,13 @@ type City struct { } `maxminddb:"location"` } +type ServerConfig struct { + Server string `mapstructure:"server" yaml:"server"` + Latitude float64 `mapstructure:"latitude" yaml:"latitude"` + Longitude float64 `mapstructure:"longitude" yaml:"longitude"` + Weight int `mapstructure:"weight" yaml:"weight"` +} + var ( configFlag = flag.String("config", "", "configuration file path") ) @@ -53,6 +61,8 @@ func main() { flag.Parse() viper.SetDefault("bind", ":8080") + viper.SetDefault("cacheSize", 1024) + viper.SetDefault("topChoices", 3) viper.SetConfigName("dlrouter") // name of config file (without extension) viper.SetConfigType("yaml") // REQUIRED if the config file does not have the extension in the name @@ -64,43 +74,7 @@ func main() { viper.SetConfigFile(*configFlag) } - err := viper.ReadInConfig() // Find and read the config file - - if err != nil { // Handle errors reading the config file - log.WithError(err).Fatalln("Unable to load config file") - } - - db, err = maxminddb.Open(viper.GetString("geodb")) - - if err != nil { - log.WithError(err).Fatalln("Unable to open database") - } - - if mapFile := viper.GetString("dl_map"); mapFile != "" { - log.WithField("file", mapFile).Info("Loading download map") - - dlMap, err = loadMap(mapFile) - - if err != nil { - log.WithError(err).Fatalln("Unable to load download map") - } - } - - serverList := viper.GetStringSlice("servers") - - var wg sync.WaitGroup - - for _, server := range serverList { - wg.Add(1) - - go func(server string) { - defer wg.Done() - - addServer(server) - }(server) - } - - wg.Wait() + reloadConfig() log.Info("Servers added, checking statuses") @@ -135,85 +109,6 @@ func main() { break } - reloadMap() + reloadConfig() } } - -var metricReplacer = strings.NewReplacer(".", "_", "-", "_") - -func addServer(server string) { - var prefix string - - if !strings.HasPrefix(server, "http") { - prefix = "https://" - } - - u, err := url.Parse(prefix + server) - - if err != nil { - log.WithFields(log.Fields{ - "error": err, - "server": server, - }).Warning("Server is invalid") - return - } - - ips, err := net.LookupIP(u.Host) - - if err != nil { - log.WithFields(log.Fields{ - "error": err, - "server": server, - }).Warning("Could not resolve address") - return - } - - var city City - err = db.Lookup(ips[0], &city) - - if err != nil { - log.WithFields(log.Fields{ - "error": err, - "server": server, - }).Warning("Could not geolocate address") - return - } - - redirects := promauto.NewCounter(prometheus.CounterOpts{ - Name: "armbian_router_redirects_" + metricReplacer.Replace(u.Host), - Help: "The number of redirects for server " + u.Host, - }) - - servers = append(servers, &Server{ - Host: u.Host, - Path: u.Path, - Latitude: city.Location.Latitude, - Longitude: city.Location.Longitude, - Redirects: redirects, - }) - - log.WithFields(log.Fields{ - "server": u.Host, - "path": u.Path, - "latitude": city.Location.Latitude, - "longitude": city.Location.Longitude, - }).Info("Added server") -} - -func reloadMap() { - mapFile := viper.GetString("dl_map") - - if mapFile == "" { - return - } - - log.WithField("file", mapFile).Info("Loading download map") - - newMap, err := loadMap(mapFile) - - if err != nil { - return - } - - dlMap = newMap -} diff --git a/servers.go b/servers.go index 7ca7d0f..d4156f3 100644 --- a/servers.go +++ b/servers.go @@ -1,12 +1,14 @@ package main import ( + "github.com/jmcvetta/randutil" "github.com/prometheus/client_golang/prometheus" log "github.com/sirupsen/logrus" "math" "net" "net/http" "runtime" + "sort" "strings" "sync" "time" @@ -24,6 +26,7 @@ type Server struct { Path string Latitude float64 Longitude float64 + Weight int Redirects prometheus.Counter } @@ -86,8 +89,8 @@ func (s ServerList) checkLoop() { t := time.NewTicker(60 * time.Second) for { - s.Check() <-t.C + s.Check() } } @@ -110,9 +113,38 @@ func (s ServerList) Check() { wg.Wait() } -// Closest will use GeoIP on the IP provided and find the closest server. +// ComputedDistance is a wrapper that contains a Server and Distance. +type ComputedDistance struct { + Server *Server + Distance float64 +} + +// DistanceList is a list of Computed Distances with an easy "Choices" func +type DistanceList []ComputedDistance + +func (d DistanceList) Choices() []randutil.Choice { + c := make([]randutil.Choice, len(d)) + + for i, item := range d { + c[i] = randutil.Choice{ + Weight: item.Server.Weight, + Item: item, + } + } + + return c +} + +// Closest will use GeoIP on the IP provided and find the closest servers. +// When we have a list of x servers closest, we can choose a random or weighted one. // Return values are the closest server, the distance, and if an error occurred. func (s ServerList) Closest(ip net.IP) (*Server, float64, error) { + i, exists := serverCache.Get(ip.String()) + + if exists { + return i.(ComputedDistance).Server, i.(ComputedDistance).Distance, nil + } + var city City err := db.Lookup(ip, &city) @@ -120,23 +152,35 @@ func (s ServerList) Closest(ip net.IP) (*Server, float64, error) { return nil, -1, err } - var closest *Server - var closestDistance float64 = -1 + c := make(DistanceList, len(s)) - for _, server := range s { + for i, server := range s { if !server.Available { continue } - distance := Distance(city.Location.Latitude, city.Location.Longitude, server.Latitude, server.Longitude) - - if closestDistance == -1 || distance < closestDistance { - closestDistance = distance - closest = server + c[i] = ComputedDistance{ + Server: server, + Distance: Distance(city.Location.Latitude, city.Location.Longitude, server.Latitude, server.Longitude), } } - return closest, closestDistance, nil + // Sort by distance + sort.Slice(s, func(i int, j int) bool { + return c[i].Distance < c[j].Distance + }) + + choice, err := randutil.WeightedChoice(c[0:topChoices].Choices()) + + if err != nil { + return nil, -1, err + } + + dist := choice.Item.(ComputedDistance) + + serverCache.Add(ip.String(), dist) + + return dist.Server, dist.Distance, nil } // haversin(θ) function