From fb10667f7976adae30404a4f3aa847afd628bd79 Mon Sep 17 00:00:00 2001 From: 127ewyo3 <226980226+127ewyo3@users.noreply.github.com> Date: Mon, 18 Aug 2025 17:06:09 +0800 Subject: [PATCH 1/2] feat: Add latency-based load balancing for upstream servers --- settings.go | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/settings.go b/settings.go index f7c9d67..5523369 100644 --- a/settings.go +++ b/settings.go @@ -1,3 +1,5 @@ +// File: settings.go + package godns import ( @@ -36,14 +38,15 @@ type Settings struct { // ResolvSettings holds resolver-specific configuration. type ResolvSettings struct { - Timeout int - Interval int - SetEDNS0 bool - ServerListFile string `toml:"server-list-file"` - ResolvFile string `toml:"resolv-file"` - DNSSECEnable bool `toml:"dnssec-enable"` + Timeout int + Interval int + SetEDNS0 bool + ServerListFile string `toml:"server-list-file"` + ResolvFile string `toml:"resolv-file"` + DNSSECEnable bool `toml:"dnssec-enable"` TrustAnchorFile string `toml:"trust-anchor-file"` - EnableLatencyBasedLoadBalancing bool `toml:"enable-latency-based-load-balancing"` + // LoadBalanceEnable controls whether latency-based load balancing is active. + LoadBalanceEnable bool `toml:"load-balance-enable"` } // DNSServerSettings holds DNS server-specific configuration. @@ -103,15 +106,15 @@ type HostsSettings struct { // BlocklistSettings holds blocklist configuration. type BlocklistSettings struct { - Enable bool - Backend string - File string - WhitelistFile string `toml:"whitelist-file"` - RefreshInterval int `toml:"refresh-interval"` - RedisEnable bool `toml:"redis-enable"` - RedisKey string `toml:"redis-key"` + Enable bool + Backend string + File string + WhitelistFile string `toml:"whitelist-file"` + RefreshInterval int `toml:"refresh-interval"` + RedisEnable bool `toml:"redis-enable"` + RedisKey string `toml:"redis-key"` RedisWhitelistKey string `toml:"redis-whitelist-key"` - TTL uint32 + TTL uint32 } // Global variable to hold the application's configuration. From fae8c616541cf82ef6e163fc6f09b2888199f028 Mon Sep 17 00:00:00 2001 From: 127ewyo3 <226980226+127ewyo3@users.noreply.github.com> Date: Mon, 18 Aug 2025 17:06:49 +0800 Subject: [PATCH 2/2] Update resolver.go --- resolver.go | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 81 insertions(+), 3 deletions(-) diff --git a/resolver.go b/resolver.go index 7a20a16..f3e06b8 100644 --- a/resolver.go +++ b/resolver.go @@ -1,3 +1,5 @@ +// File: resolver.go + package godns import ( @@ -9,6 +11,7 @@ import ( "strings" "sync" "time" + "sort" "github.com/miekg/dns" "github.com/ProxyFi/GoDNS/features/dnssec" @@ -34,10 +37,17 @@ type RResp struct { rtt time.Duration } +// nameserverState holds a nameserver address and its measured Round-Trip Time (RTT). +type nameserverState struct { + addr string + rtt time.Duration +} + // Resolver handles DNS queries to upstream servers. type Resolver struct { // servers stores a list of default upstream nameservers. - servers []string + // This field is used only during initial configuration and is deprecated in favor of rttServers. + servers []string // domain_server is a suffix tree used for specific domain rules. domain_server *suffixTreeNode // config holds the resolver's settings. @@ -47,6 +57,9 @@ type Resolver struct { // udpClient and tcpClient are reused for all DNS lookups to reduce overhead. udpClient *dns.Client tcpClient *dns.Client + // rttServers stores a map of nameservers and their measured RTTs for load balancing. + rttServers map[string]*nameserverState + serversMutex sync.RWMutex } // NewResolver creates and initializes a new Resolver instance. @@ -76,6 +89,8 @@ func NewResolver(c ResolvSettings) (*Resolver, error) { ReadTimeout: 5 * time.Second, WriteTimeout: 5 * time.Second, }, + rttServers: make(map[string]*nameserverState), + serversMutex: sync.RWMutex{}, } // Load server list from a dedicated file, if configured. @@ -96,6 +111,26 @@ func NewResolver(c ResolvSettings) (*Resolver, error) { r.servers = append(r.servers, net.JoinHostPort(server, clientConfig.Port)) } } + + // Populate rttServers map from the loaded server list. + // This map will be used for latency-based load balancing. + if len(r.servers) > 0 { + for _, server := range r.servers { + r.rttServers[server] = &nameserverState{ + addr: server, + rtt: time.Duration(c.Timeout) * time.Second, + } + } + } else { + // Fallback to hardcoded public DNS servers if none are configured. + defaultServers := []string{"8.8.8.8:53", "8.8.4.4:53"} + for _, server := range defaultServers { + r.rttServers[server] = &nameserverState{ + addr: server, + rtt: time.Duration(c.Timeout) * time.Second, + } + } + } return r, nil } @@ -158,6 +193,12 @@ func (r *Resolver) Lookup(net string, req *dns.Msg) (msg *dns.Msg, err error) { if lookupErr != nil { log.Warn("%s lookup on %s failed: %s", req.Question[0].String(), nameserver, lookupErr) + // Penalize the server by setting its RTT to a high value. + r.serversMutex.Lock() + if state, ok := r.rttServers[nameserver]; ok { + state.rtt = time.Duration(r.config.Timeout) * time.Second + } + r.serversMutex.Unlock() return } if a != nil && a.Rcode != dns.RcodeServerFailure { @@ -191,6 +232,17 @@ func (r *Resolver) Lookup(net string, req *dns.Msg) (msg *dns.Msg, err error) { } log.Debug("DNSSEC validation successful for %s from %s", req.Question[0].String(), re.nameserver) } + + // Update the RTT for the successful nameserver based on load balancing setting. + if r.config.LoadBalanceEnable { + r.serversMutex.Lock() + if state, ok := r.rttServers[re.nameserver]; ok { + // Use an Exponential Moving Average (EMA) for smoothing. + // This gives more weight to recent measurements. + state.rtt = time.Duration(float64(state.rtt)*0.9 + float64(re.rtt)*0.1) + } + r.serversMutex.Unlock() + } log.Debug("%s resolv on %s rtt: %v", UnFqdn(req.Question[0].Name), re.nameserver, re.rtt) return re.msg, nil @@ -203,6 +255,7 @@ func (r *Resolver) Lookup(net string, req *dns.Msg) (msg *dns.Msg, err error) { // Nameservers returns the list of nameservers to use for a given query name. // It prioritizes specific rules from the suffix tree over the default list. +// If load balancing is enabled, it sorts the default servers by RTT. func (r *Resolver) Nameservers(qname string) []string { // Split the domain name into parts and ignore the trailing dot. queryKeys := strings.Split(strings.Trim(qname, "."), ".") @@ -217,12 +270,37 @@ func (r *Resolver) Nameservers(qname string) []string { return ns } - // If no specific rule is found, use the default server list. + // If no specific rule is found and load balancing is enabled, sort by RTT. + if r.config.LoadBalanceEnable { + r.serversMutex.RLock() + defer r.serversMutex.RUnlock() + + var states []*nameserverState + for _, state := range r.rttServers { + states = append(states, state) + } + + // Sort the nameservers by their measured RTT. + sort.Slice(states, func(i, j int) bool { + return states[i].rtt < states[j].rtt + }) + + var sortedServers []string + for _, state := range states { + sortedServers = append(sortedServers, state.addr) + } + + if len(sortedServers) > 0 { + return sortedServers + } + } + + // Fallback to the original, unsorted list if no RTT data or load balancing is disabled. if len(r.servers) > 0 { return r.servers } - // Fallback to hardcoded public DNS servers if no other servers are configured. + // Final fallback to hardcoded public DNS servers if no other servers are configured. return []string{"8.8.8.8:53", "8.8.4.4:53"} }