better implementation of the response rate tracking

Author: caffix

cmd/resolve/main.go

@@ -168,6 +168,8 @@ func (p *params) SetupResolverPool(list []string, rpath string, timeout int, det
 		p.Pool.SetDetectionResolver(p.QPS, detector)
 		p.Detection = true
 	}
+
+	p.Pool.SetRateTracker(resolve.NewRateTracker())
 	return nil
 }
 

conn.go

@@ -168,13 +168,14 @@ func (r *connections) responses(c *connection) {
 		default:
 		}
 		if n, addr, err := c.conn.ReadFrom(b); err == nil && n >= headerSize {
+			at := time.Now()
 			m := new(dns.Msg)
 
 			if err := m.Unpack(b[:n]); err == nil && len(m.Question) > 0 {
 				r.resps.Append(&resp{
 					Msg:  m,
 					Addr: addr,
-					At:   time.Now(),
+					At:   at,
 				})
 			}
 		}

go.mod

@@ -9,6 +9,7 @@ require (
 	go.uber.org/ratelimit v0.3.1
 	golang.org/x/net v0.37.0
 	golang.org/x/sys v0.31.0
+	golang.org/x/time v0.11.0
 )
 
 require (

go.sum

@@ -26,6 +26,8 @@ golang.org/x/sync v0.12.0 h1:MHc5BpPuC30uJk597Ri8TV3CNZcTLu6B6z4lJy+g6Jw=
 golang.org/x/sync v0.12.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
 golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
 golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0=
+golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
 golang.org/x/tools v0.31.0 h1:0EedkvKDbh+qistFTd0Bcwe/YLh4vHwWEkiI0toFIBU=
 golang.org/x/tools v0.31.0/go.mod h1:naFTU+Cev749tSJRXJlna0T3WxKvb1kWEx15xA4SdmQ=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

rate.go

@@ -5,29 +5,30 @@
 package resolve
 
 import (
+	"context"
+	"math"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/miekg/dns"
-	"go.uber.org/ratelimit"
 	"golang.org/x/net/publicsuffix"
+	"golang.org/x/time/rate"
 )
 
 const (
-	startQPSPerNameserver = 10
-	maxQPSPerNameserver   = 100
-	numIntervalSeconds    = 2
-	rateUpdateInterval    = numIntervalSeconds * time.Second
-	maxTimeoutPercentage  = 0.5
+	maxQPSPerNameserver = 100
+	numIntervalSeconds  = 5
+	rateUpdateInterval  = numIntervalSeconds * time.Second
+	minUpdateSampleSize = 10
 )
 
 type rateTrack struct {
 	sync.Mutex
-	qps     int
-	rate    ratelimit.Limiter
-	success int
-	timeout int
+	rate  *rate.Limiter
+	avg   time.Duration
+	count int
+	first bool
 }
 
 type RateTracker struct {
@@ -52,9 +53,11 @@ func NewRateTracker() *RateTracker {
 }
 
 func newRateTrack() *rateTrack {
+	limit := rate.Every(100 * time.Millisecond)
+
 	return &rateTrack{
-		qps:  startQPSPerNameserver,
-		rate: ratelimit.New(startQPSPerNameserver),
+		rate:  rate.NewLimiter(limit, 1),
+		first: true,
 	}
 }
 
@@ -75,24 +78,24 @@ func (r *RateTracker) Take(sub string) {
 	rate := tracker.rate
 	tracker.Unlock()
 
-	rate.Take()
-}
-
-// Success signals to the RateTracker that a request for the provided subdomain name was successful.
-func (r *RateTracker) Success(sub string) {
-	tracker := r.getDomainRateTracker(sub)
-
-	tracker.Lock()
-	tracker.success++
-	tracker.Unlock()
+	_ = rate.Wait(context.TODO())
 }
 
-// Timeout signals to the RateTracker that a request for the provided subdomain name timed out.
-func (r *RateTracker) Timeout(sub string) {
+// ReportResponseTime provides the response time for a request for the domain name provided in the sub parameter.
+func (r *RateTracker) ReportResponseTime(sub string, delta time.Duration) {
+	var average, count float64
 	tracker := r.getDomainRateTracker(sub)
 
 	tracker.Lock()
-	tracker.timeout++
+	tracker.count++
+	count = float64(tracker.count)
+	average = float64(tracker.avg.Milliseconds())
+	average = ((average * (count - 1)) + float64(delta.Milliseconds())) / count
+	tracker.avg = time.Duration(math.Round(average)) * time.Millisecond
+
+	if tracker.first {
+		defer tracker.update()
+	}
 	tracker.Unlock()
 }
 
@@ -124,26 +127,18 @@ func (r *RateTracker) updateAllRateLimiters() {
 func (rt *rateTrack) update() {
 	rt.Lock()
 	defer rt.Unlock()
-	// check if this rate tracker has already been updated
-	if rt.success == 0 && rt.timeout == 0 {
+
+	if rt.first {
+		rt.first = false
+	} else if rt.count < minUpdateSampleSize {
 		return
 	}
-	// timeouts in excess of maxTimeoutPercentage indicate a need to slow down
-	if float64(rt.timeout)/float64(rt.success+rt.timeout) > maxTimeoutPercentage {
-		rt.qps -= 1
-		if rt.qps <= 0 {
-			rt.qps = 1
-		}
-	} else {
-		rt.qps += 1
-		if rt.qps > maxQPSPerNameserver {
-			rt.qps = maxQPSPerNameserver
-		}
-	}
+
+	limit := rate.Every(rt.avg)
 	// update the QPS rate limiter and reset counters
-	rt.rate = ratelimit.New(rt.qps)
-	rt.success = 0
-	rt.timeout = 0
+	rt.rate = rate.NewLimiter(limit, 1)
+	rt.avg = 0
+	rt.count = 0
 }
 
 func (r *RateTracker) getDomainRateTracker(sub string) *rateTrack {

rate_test.go

@@ -18,52 +18,27 @@ func TestUpdateRateLimiters(t *testing.T) {
 	rt.Take(domain)
 	tracker := rt.getDomainRateTracker(domain)
 
-	tracker.Lock()
-	qps := tracker.qps
-	tracker.Unlock()
-	num := qps / 3
-	// set a large number of timeouts
-	for i := 0; i < num; i++ {
-		rt.Success(domain)
-	}
-
-	max := tracker.qps - num
-	for i := 0; i < max; i++ {
-		rt.Timeout(domain)
-	}
-	time.Sleep(rateUpdateInterval + (rateUpdateInterval / 2))
+	rt.ReportResponseTime(domain, 500*time.Millisecond)
 
 	tracker.Lock()
-	qps2 := tracker.qps
+	limit := tracker.rate.Limit()
 	tracker.Unlock()
 	// the QPS should now be lower
-	if qps2 >= qps {
-		t.Errorf("Unexpected QPS, expected QPS lower than %d, got %d", qps, qps2)
+	if limit > 3 {
+		t.Errorf("Unexpected QPS, expected QPS lower than %d, got %f", 3, limit)
 	}
 
 	tracker.Lock()
-	succ := tracker.success
-	tout := tracker.timeout
+	tracker.avg = 50 * time.Millisecond
+	tracker.count = minUpdateSampleSize
 	tracker.Unlock()
-	// check that the counters have been cleared
-	if succ != 0 || tout != 0 {
-		t.Errorf("Unexpected counter values, Success Counter %d, Timeout Counter %d", succ, tout)
-	}
-
-	tracker.Lock()
-	qps = tracker.qps
-	tracker.Unlock()
-	// set a large number of successes
-	for i := 0; i < qps; i++ {
-		rt.Success(domain)
-	}
-	time.Sleep(rateUpdateInterval + (rateUpdateInterval / 2))
+	tracker.update()
 
 	tracker.Lock()
-	qps2 = tracker.qps
+	limit = tracker.rate.Limit()
 	tracker.Unlock()
 	// the QPS should now be higher
-	if qps2 <= qps {
-		t.Errorf("Unexpected QPS, expected QPS higher than %d, got %d", qps, qps2)
+	if limit < 20 || limit > 21 {
+		t.Errorf("Unexpected QPS, expected QPS of %d, got %f", 20, limit)
 	}
 }

resolvers.go

@@ -357,7 +357,8 @@ func (r *Resolvers) processSingleResp(response *resp) {
 			req.Result <- req.Resp
 			req.Res.collectStats(req.Resp)
 			if r.servRates != nil {
-				r.servRates.Success(name)
+				delta := req.RecvAt.Sub(req.SentAt)
+				r.servRates.ReportResponseTime(name, delta)
 			}
 			req.release()
 		}
@@ -392,9 +393,6 @@ func (r *Resolvers) timeouts() {
 				for _, req := range res.xchgs.removeExpired() {
 					req.errNoResponse()
 					res.collectStats(req.Msg)
-					if r.servRates != nil {
-						r.servRates.Timeout(req.Msg.Question[0].Name)
-					}
 					req.release()
 				}
 			}

xchg_test.go

@@ -44,18 +44,18 @@ func TestXchgUpdateTimestamp(t *testing.T) {
 	msg := QueryMsg(name, dns.TypeA)
 	req := &request{Msg: msg}
 
-	if !req.Timestamp.IsZero() {
+	if !req.SentAt.IsZero() {
 		t.Errorf("Expected the new request to have a zero value timestamp")
 	}
 	if err := xchg.add(req); err != nil {
 		t.Errorf("Failed to add the request")
 	}
-	xchg.updateTimestamp(msg.Id, name)
+	xchg.updateSentAt(msg.Id, name)
 	// For complete coverage
-	xchg.updateTimestamp(msg.Id, "Bad Name")
+	xchg.updateSentAt(msg.Id, "Bad Name")
 
 	req = xchg.remove(msg.Id, msg.Question[0].Name)
-	if req == nil || req.Timestamp.IsZero() {
+	if req == nil || req.SentAt.IsZero() {
 		t.Errorf("Expected the updated request to not have a zero value timestamp")
 	}
 }
@@ -67,8 +67,8 @@ func TestXchgRemoveExpired(t *testing.T) {
 	for _, name := range names {
 		msg := QueryMsg(name, dns.TypeA)
 		if err := xchg.add(&request{
-			Msg:       msg,
-			Timestamp: time.Now(),
+			Msg:    msg,
+			SentAt: time.Now(),
 		}); err != nil {
 			t.Errorf("Failed to add the request")
 		}
@@ -77,8 +77,8 @@ func TestXchgRemoveExpired(t *testing.T) {
 	name := "vpn.caffix.net"
 	msg := QueryMsg(name, dns.TypeA)
 	if err := xchg.add(&request{
-		Msg:       msg,
-		Timestamp: time.Now().Add(3 * time.Second),
+		Msg:    msg,
+		SentAt: time.Now().Add(3 * time.Second),
 	}); err != nil {
 		t.Errorf("Failed to add the request")
 	}