From 4e953975d3726404dd796b3167a8b7f1747d8f8f Mon Sep 17 00:00:00 2001
From: Homayoon Alimohammadi <homayoonalimohammadi@gmail.com>
Date: Sun, 25 Feb 2024 21:06:31 +0330
Subject: [PATCH] Add MinDNSResolutionRate Option

---
 internal/resolver/dns/dns_resolver.go      | 14 +++++---
 internal/resolver/dns/dns_resolver_test.go | 38 ++++++++++++++++++++--
 internal/resolver/dns/internal/internal.go |  4 ---
 resolver/dns/dns_resolver.go               |  9 +++++
 4 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/internal/resolver/dns/dns_resolver.go b/internal/resolver/dns/dns_resolver.go
index abab35e250ef..c06ff7f40d83 100644
--- a/internal/resolver/dns/dns_resolver.go
+++ b/internal/resolver/dns/dns_resolver.go
@@ -41,9 +41,15 @@ import (
 	"google.golang.org/grpc/serviceconfig"
 )
 
-// EnableSRVLookups controls whether the DNS resolver attempts to fetch gRPCLB
-// addresses from SRV records.  Must not be changed after init time.
-var EnableSRVLookups = false
+var (
+	// EnableSRVLookups controls whether the DNS resolver attempts to fetch gRPCLB
+	// addresses from SRV records.  Must not be changed after init time.
+	EnableSRVLookups = false
+
+	// MinResolutionRate is the minimum rate at which re-resolutions are
+	// allowed. This helps to prevent excessive re-resolution.
+	MinResolutionRate = 30 * time.Second
+)
 
 // ResolvingTimeout specifies the maximum duration for a DNS resolution request.
 // If the timeout expires before a response is received, the request will be canceled.
@@ -208,7 +214,7 @@ func (d *dnsResolver) watcher() {
 			// Success resolving, wait for the next ResolveNow. However, also wait 30
 			// seconds at the very least to prevent constantly re-resolving.
 			backoffIndex = 1
-			waitTime = internal.MinResolutionRate
+			waitTime = MinResolutionRate
 			select {
 			case <-d.ctx.Done():
 				return
diff --git a/internal/resolver/dns/dns_resolver_test.go b/internal/resolver/dns/dns_resolver_test.go
index 498cf5b83e27..a431e38f2b03 100644
--- a/internal/resolver/dns/dns_resolver_test.go
+++ b/internal/resolver/dns/dns_resolver_test.go
@@ -70,9 +70,9 @@ func overrideNetResolver(t *testing.T, r *testNetResolver) {
 
 // Override the DNS Min Res Rate used by the resolver.
 func overrideResolutionRate(t *testing.T, d time.Duration) {
-	origMinResRate := dnsinternal.MinResolutionRate
-	dnsinternal.MinResolutionRate = d
-	t.Cleanup(func() { dnsinternal.MinResolutionRate = origMinResRate })
+	origMinResRate := dns.MinResolutionRate
+	dns.MinResolutionRate = d
+	t.Cleanup(func() { dns.MinResolutionRate = origMinResRate })
 }
 
 // Override the timer used by the DNS resolver to fire after a duration of d.
@@ -1258,3 +1258,35 @@ func (s) TestResolveTimeout(t *testing.T) {
 		}
 	}
 }
+
+// Test verifies that changing [MinResolutionRate] variable correctly effects
+// the resolution behaviour
+func (s) TestMinResolutionRate(t *testing.T) {
+	const target = "foo.bar.com"
+
+	overrideResolutionRate(t, 1*time.Millisecond)
+	tr := &testNetResolver{
+		hostLookupTable: map[string][]string{
+			"foo.bar.com": {"1.2.3.4", "5.6.7.8"},
+		},
+		txtLookupTable: map[string][]string{
+			"_grpc_config.foo.bar.com": txtRecordServiceConfig(txtRecordGood),
+		},
+	}
+	overrideNetResolver(t, tr)
+
+	r, stateCh, _ := buildResolverWithTestClientConn(t, target)
+
+	wantAddrs := []resolver.Address{{Addr: "1.2.3.4" + colonDefaultPort}, {Addr: "5.6.7.8" + colonDefaultPort}}
+	wantSC := scJSON
+
+	for i := 0; i < 5; i++ {
+		// set context timeout slightly higher than the resolution rate to make sure resolutions
+		// happen successfully
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+		defer cancel()
+
+		verifyUpdateFromResolver(ctx, t, stateCh, wantAddrs, nil, wantSC)
+		r.ResolveNow(resolver.ResolveNowOptions{})
+	}
+}
diff --git a/internal/resolver/dns/internal/internal.go b/internal/resolver/dns/internal/internal.go
index c7fc557d00c1..9cb3b5c58a6a 100644
--- a/internal/resolver/dns/internal/internal.go
+++ b/internal/resolver/dns/internal/internal.go
@@ -50,10 +50,6 @@ var (
 
 // The following vars are overridden from tests.
 var (
-	// MinResolutionRate is the minimum rate at which re-resolutions are
-	// allowed. This helps to prevent excessive re-resolution.
-	MinResolutionRate = 30 * time.Second
-
 	// TimeAfterFunc is used by the DNS resolver to wait for the given duration
 	// to elapse. In non-test code, this is implemented by time.After.  In test
 	// code, this can be used to control the amount of time the resolver is
diff --git a/resolver/dns/dns_resolver.go b/resolver/dns/dns_resolver.go
index b54a3a3225d4..756c0ee0fcf3 100644
--- a/resolver/dns/dns_resolver.go
+++ b/resolver/dns/dns_resolver.go
@@ -52,3 +52,12 @@ func SetResolvingTimeout(timeout time.Duration) {
 func NewBuilder() resolver.Builder {
 	return dns.NewBuilder()
 }
+
+// SetMinResolutionRate sets the default minimum rate at which DNS re-resolutions are
+// allowed. This helps to prevent excessive re-resolution.
+//
+// Using this option overwrites the default [MinResolutionRate] specified
+// in the internal dns resolver package.
+func SetMinResolutionRate(d time.Duration) {
+	dns.MinResolutionRate = d
+}