From e7d310fd769d8612a3ea75a67f963aa8aa991f40 Mon Sep 17 00:00:00 2001
From: Aabishkar Aryal <aabishkaryal@gmail.com>
Date: Thu, 25 Sep 2025 12:38:57 +0545
Subject: [PATCH] Add DNS cache management methods for TCPDialer (#2072)

* Add DNS cache management methods for TCPDialer

Resolves #2066

This commit introduces two new methods for managing DNS cache in TCPDialer:

1. FlushDNSCache() - Clears all cached DNS entries, forcing fresh lookups
2. CleanDNSCache() - Removes only expired entries based on DNSCacheDuration

Key changes:
- Add FlushDNSCache() and CleanDNSCache() methods to TCPDialer
- Add global FlushDNSCache() and CleanDNSCache() functions for default dialer
- Refactor tcpAddrsClean() to extract reusable cleanExpiredDNSEntries() method
- Add comprehensive tests with mock resolver to verify caching behavior

Use case: Users can now set longer cache durations (e.g., 30 minutes) and
manually refresh DNS when needed, providing better control over DNS resolution
timing while maintaining performance benefits of caching.

* Remove CleanDNSCache method to reduce the API surface layer and related tests from TCPDialer

* fix: resolve godot linter issue in client_test.go

Add missing period to comment to comply with godot linter rule requiring
comments to end with proper punctuation.
---
 client_test.go | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++
 tcpdialer.go   | 39 +++++++++++++++++++++++++-------
 2 files changed, 91 insertions(+), 8 deletions(-)

diff --git a/client_test.go b/client_test.go
index 1e6aaac..4e07ea9 100644
--- a/client_test.go
+++ b/client_test.go
@@ -3,6 +3,7 @@ package fasthttp
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -3559,3 +3560,62 @@ func (f F) Read(p []byte) (n int, err error) {
 	time.Sleep(500 * time.Microsecond)
 	return f.Reader.Read(p)
 }
+
+func TestTCPDialerFlushDNSCache(t *testing.T) {
+	resolver := &testResolver{
+		lookupCountByHost: make(map[string]int),
+		resolver:          net.DefaultResolver,
+	}
+
+	dialer := &TCPDialer{
+		DNSCacheDuration: 30 * time.Minute, // Long cache
+		Resolver:         resolver,
+	}
+
+	// First dial - should trigger DNS lookup
+	conn1, err := dialer.DialTimeout("httpbin.org:80", 5*time.Second)
+	if err != nil {
+		t.Skip("Dial failed:", err)
+	}
+	conn1.Close()
+
+	if resolver.lookupCountByHost["httpbin.org"] != 1 {
+		t.Errorf("Expected 1 DNS lookup after first dial, got %d", resolver.lookupCountByHost["httpbin.org"])
+	}
+
+	// Second dial - should use cache (no new DNS lookup)
+	conn2, err := dialer.DialTimeout("httpbin.org:80", 5*time.Second)
+	if err != nil {
+		t.Skip("Second dial failed:", err)
+	}
+	conn2.Close()
+
+	if resolver.lookupCountByHost["httpbin.org"] != 1 {
+		t.Errorf("Expected 1 DNS lookup after cached dial, got %d", resolver.lookupCountByHost["httpbin.org"])
+	}
+
+	// Flush cache - should clear all entries
+	dialer.FlushDNSCache()
+
+	// Third dial - should trigger new DNS lookup since cache was flushed
+	conn3, err := dialer.DialTimeout("httpbin.org:80", 5*time.Second)
+	if err != nil {
+		t.Skip("Third dial failed:", err)
+	}
+	conn3.Close()
+
+	if resolver.lookupCountByHost["httpbin.org"] != 2 {
+		t.Errorf("Expected 2 DNS lookups after cache flush, got %d", resolver.lookupCountByHost["httpbin.org"])
+	}
+}
+
+// Simple test resolver that implements the Resolver interface.
+type testResolver struct {
+	resolver          *net.Resolver
+	lookupCountByHost map[string]int
+}
+
+func (r *testResolver) LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) {
+	r.lookupCountByHost[host]++
+	return r.resolver.LookupIPAddr(ctx, host)
+}
diff --git a/tcpdialer.go b/tcpdialer.go
index f87008d..48e57fa 100644
--- a/tcpdialer.go
+++ b/tcpdialer.go
@@ -271,6 +271,22 @@ func (d *TCPDialer) DialDualStackTimeout(addr string, timeout time.Duration) (ne
 	return d.dial(addr, true, timeout)
 }
 
+// FlushDNSCache clears all cached DNS entries, forcing fresh DNS lookups on subsequent dials.
+// This is useful when you want to ensure fresh DNS resolution, for example after network changes.
+func (d *TCPDialer) FlushDNSCache() {
+	d.tcpAddrsMap.Range(func(k, v any) bool {
+		d.tcpAddrsMap.Delete(k)
+		return true
+	})
+}
+
+// FlushDNSCache clears all cached DNS entries for the default dialer,
+// forcing fresh DNS lookups on subsequent Dial* calls.
+// This is useful when you want to ensure fresh DNS resolution, for example after network changes.
+func FlushDNSCache() {
+	defaultDialer.FlushDNSCache()
+}
+
 func (d *TCPDialer) dial(addr string, dualStack bool, timeout time.Duration) (net.Conn, error) {
 	d.once.Do(func() {
 		if d.Concurrency > 0 {
@@ -406,17 +422,24 @@ type tcpAddrEntry struct {
 // by Dial* functions.
 const DefaultDNSCacheDuration = time.Minute
 
-func (d *TCPDialer) tcpAddrsClean() {
+// cleanExpiredDNSEntries removes expired DNS cache entries based on DNSCacheDuration.
+// This is the core cleanup logic used by both the background cleaner and manual cleanup.
+func (d *TCPDialer) cleanExpiredDNSEntries() {
 	expireDuration := 2 * d.DNSCacheDuration
+
+	t := time.Now()
+	d.tcpAddrsMap.Range(func(k, v any) bool {
+		if e, ok := v.(*tcpAddrEntry); ok && t.Sub(e.resolveTime) > expireDuration {
+			d.tcpAddrsMap.Delete(k)
+		}
+		return true
+	})
+}
+
+func (d *TCPDialer) tcpAddrsClean() {
 	for {
 		time.Sleep(time.Second)
-		t := time.Now()
-		d.tcpAddrsMap.Range(func(k, v any) bool {
-			if e, ok := v.(*tcpAddrEntry); ok && t.Sub(e.resolveTime) > expireDuration {
-				d.tcpAddrsMap.Delete(k)
-			}
-			return true
-		})
+		d.cleanExpiredDNSEntries()
 	}
 }