Repository: valyala/fasthttp
Pull request #1046 (2021-06-14 09:59:39)

TCPDialer :: DNSCacheDuration option
TCPDialer{} extended to support DNSCacheDuration

![image](https://user-images.githubusercontent.com/1286155/121874609-2e9acd80-cd10-11eb-937d-574c2ab805a6.png)



Patch:
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 22b46e4..ef1b5c5 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -1,7 +1,9 @@
-name: lint
-
-on: [push, pull_request]
-
+name: Lint
+on:
+  push:
+    branches:
+      - master
+  pull_request:
 jobs:
   lint:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml
index da67002..7763984 100644
--- a/.github/workflows/security.yml
+++ b/.github/workflows/security.yml
@@ -1,10 +1,14 @@
-on: [push, pull_request]
 name: Security
+on:
+  push:
+    branches:
+      - master
+  pull_request:
 jobs:
   test:
     strategy:
       matrix:
-        go-version: [1.13.x, 1.14.x, 1.15.x, 1.16.x]
+        go-version: [1.16.x]
         platform: [ubuntu-latest]
     runs-on: ${{ matrix.platform }}
     steps:
diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 55e6a5f..0000000
--- a/.travis.yml
+++ /dev/null
@@ -1,67 +0,0 @@
-language: go
-arch:
-  - amd64
-  - ppc64le
-
-dist: bionic
-
-os:
-  - linux
-  - osx
-  - windows
-go:
-  - 1.16.x
-  - 1.15.x
-  - 1.14.x
-  - 1.13.x
-  - 1.12.x
-  - tip
-
-script:
-  - go test -v ./...
-
-env:
-  global:
-    secure: 'v/F0oI9zE9mcpEp4AVdHzSSHbe5ZFtH6B0i/BiUXKdQRQ10+JMPDOFRJQti7yxjMwltyd/QSFmR50Fl108sQYpo4xdlEXMHp2Y6OAN6crrp6PuHbLYgDWu3df/cH7/BqDyIq1uX8KZEeQssnygYN8hN4tpJCUg+NIb40Lm57Zsodt8DVjjyDWQQFDL7soNyAwGwQIqEyJsn+NUieXWEB1Qnt0xUtPIReuLlrwXR8wC1nLEjG9yz4ftDHHQdhVbO2b+xGWyaJ7QB5ixztaQP8Jnny6kSW9j6zEhJVuzdZ6d3xz23ibCbzSXBHdIUEI9u6ifQj8BYXr8fFS0FB3++IxgAYSs3ybZ+qEwuAxSBBm6YNW+3FrfDknVwTQscjKqnXPisjUqaRC9b31hke0tXzBq1488hE+wxMXeDM4LwWT5IMEO2gz0WGQXxmdVit72DIjCZxJkf1TvZZ0YH7Y//6wJTYYP9xulsy4gqu8CuFdWiF3fiGc3p5DTIS75nJ/Yy76Sa1pRPASKCujfLxtHE6Mt0XKvSolIXklYIzBkjN6vn80N6JIrqtqlimBGPW/Ec6+dwbmRe2AcOKRl4y7pZsGYhJhqdue1mucUYO/e2QeBZJGkqqG+zF5AW0v8x29BHvMwViAonc8o9eelkJ8khYzc/Qeq05pZnR/N/Pqfc+68k='
-
-before_install:
-  - go get -t -v ./...
-
-jobs:
-  allow_failures:
-    - go: tip
-
-  include:
-    - stage: cross compilation
-      os:
-        - linux
-      script:
-        - GOOS=linux go build
-        - GOOS=darwin go build
-        - GOOS=freebsd go build
-        - GOOS=windows go build
-        - GOARCH=386 go build
-    - stage: cross compilation
-      os:
-        - osx
-      script:
-        - GOOS=linux go build
-        - GOOS=darwin go build
-        - GOOS=freebsd go build
-        - GOOS=windows go build
-    - stage: race detector
-      os:
-        - linux
-      script:
-        - go test -race -v ./...
-    - stage: race detector
-      os:
-        - osx
-      script:
-        - go test -race -v ./...
-    - stage: race detector
-      arch: ppc64le
-      os:
-        - linux
-      script:
-        - go test -race -v ./...
diff --git a/README.md b/README.md
index 8aee99d..894141e 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# fasthttp [![Build Status](https://travis-ci.org/valyala/fasthttp.svg?branch=master)](https://travis-ci.org/valyala/fasthttp?branch=master) [![GoDoc](https://godoc.org/github.com/valyala/fasthttp?status.svg)](http://godoc.org/github.com/valyala/fasthttp) [![Go Report](https://goreportcard.com/badge/github.com/valyala/fasthttp)](https://goreportcard.com/report/github.com/valyala/fasthttp) [![Sourcegraph](https://sourcegraph.com/github.com/valyala/fasthttp/-/badge.svg)](https://sourcegraph.com/github.com/valyala/fasthttp?badge)
+# fasthttp [![GoDoc](https://godoc.org/github.com/valyala/fasthttp?status.svg)](http://godoc.org/github.com/valyala/fasthttp) [![Go Report](https://goreportcard.com/badge/github.com/valyala/fasthttp)](https://goreportcard.com/report/github.com/valyala/fasthttp) [![Sourcegraph](https://sourcegraph.com/github.com/valyala/fasthttp/-/badge.svg)](https://sourcegraph.com/github.com/valyala/fasthttp?badge)
 
 ![FastHTTP – Fastest and reliable HTTP implementation in Go](https://github.com/fasthttp/docs-assets/raw/master/banner@0.5.png)
 
diff --git a/client.go b/client.go
index f135ec3..4736321 100644
--- a/client.go
+++ b/client.go
@@ -294,9 +294,9 @@ type Client struct {
 	// By default will use isIdempotent function
 	RetryIf RetryIfFunc
 
-	mLock sync.Mutex
-	m     map[string]*HostClient
-	ms    map[string]*HostClient
+	mLock      sync.Mutex
+	m          map[string]*HostClient
+	ms         map[string]*HostClient
 	readerPool sync.Pool
 	writerPool sync.Pool
 }
@@ -594,6 +594,9 @@ type DialFunc func(addr string) (net.Conn, error)
 // Request argument passed to RetryIfFunc, if there are any request errors.
 type RetryIfFunc func(request *Request) bool
 
+// TransportFunc wraps every request/response.
+type TransportFunc func(*Request, *Response) error
+
 // HostClient balances http requests among hosts listed in Addr.
 //
 // HostClient may be used for balancing load among multiple upstream hosts.
@@ -745,6 +748,9 @@ type HostClient struct {
 	// By default will use isIdempotent function
 	RetryIf RetryIfFunc
 
+	// Transport defines a transport-like mechanism that wraps every request/response.
+	Transport TransportFunc
+
 	clientName  atomic.Value
 	lastUseTime uint32
 
@@ -1360,8 +1366,16 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 	resp.Reset()
 	resp.SkipBody = customSkipBody
 
-	if c.DisablePathNormalizing {
-		req.URI().DisablePathNormalizing = true
+	req.URI().DisablePathNormalizing = c.DisablePathNormalizing
+
+	userAgentOld := req.Header.UserAgent()
+	if len(userAgentOld) == 0 {
+		req.Header.userAgent = append(req.Header.userAgent[:0], c.getClientName()...)
+	}
+
+	if c.Transport != nil {
+		err := c.Transport(req, resp)
+		return err == nil, err
 	}
 
 	cc, err := c.acquireConn(req.timeout, req.ConnectionClose())
@@ -1388,10 +1402,6 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 		resetConnection = true
 	}
 
-	userAgentOld := req.Header.UserAgent()
-	if len(userAgentOld) == 0 {
-		req.Header.userAgent = append(req.Header.userAgent[:0], c.getClientName()...)
-	}
 	bw := c.acquireWriter(conn)
 	err = req.Write(bw)
 
diff --git a/fasthttpproxy/http.go b/fasthttpproxy/http.go
index 0ba3092..6bc9c1e 100644
--- a/fasthttpproxy/http.go
+++ b/fasthttpproxy/http.go
@@ -22,7 +22,7 @@ func FasthttpHTTPDialer(proxy string) fasthttp.DialFunc {
 	return FasthttpHTTPDialerTimeout(proxy, 0)
 }
 
-// FasthttpHTTPDialer returns a fasthttp.DialFunc that dials using
+// FasthttpHTTPDialerTimeout returns a fasthttp.DialFunc that dials using
 // the provided HTTP proxy using the given timeout.
 //
 // Example usage:
diff --git a/fasthttpproxy/proxy_env.go b/fasthttpproxy/proxy_env.go
index 22cc76e..2457ad5 100644
--- a/fasthttpproxy/proxy_env.go
+++ b/fasthttpproxy/proxy_env.go
@@ -14,6 +14,12 @@ import (
 	"github.com/valyala/fasthttp"
 )
 
+const (
+	httpsScheme = "https"
+	httpScheme  = "http"
+	tlsPort     = "443"
+)
+
 // FasthttpProxyHTTPDialer returns a fasthttp.DialFunc that dials using
 // the the env(HTTP_PROXY, HTTPS_PROXY and NO_PROXY) configured HTTP proxy.
 //
@@ -32,13 +38,6 @@ func FasthttpProxyHTTPDialer() fasthttp.DialFunc {
 //	c := &fasthttp.Client{
 //		Dial: FasthttpProxyHTTPDialerTimeout(time.Second * 2),
 //	}
-
-const (
-	httpsScheme = "https"
-	httpScheme  = "http"
-	tlsPort     = "443"
-)
-
 func FasthttpProxyHTTPDialerTimeout(timeout time.Duration) fasthttp.DialFunc {
 	proxier := httpproxy.FromEnvironment().ProxyFunc()
 
diff --git a/prefork/prefork.go b/prefork/prefork.go
index 9ace2ed..dc08b76 100644
--- a/prefork/prefork.go
+++ b/prefork/prefork.go
@@ -22,7 +22,9 @@ var (
 	defaultLogger = Logger(log.New(os.Stderr, "", log.LstdFlags))
 	// ErrOverRecovery is returned when the times of starting over child prefork processes exceed
 	// the threshold.
-	ErrOverRecovery           = errors.New("exceeding the value of RecoverThreshold")
+	ErrOverRecovery = errors.New("exceeding the value of RecoverThreshold")
+
+	// ErrOnlyReuseportOnWindows is returned when Reuseport is false.
 	ErrOnlyReuseportOnWindows = errors.New("windows only supports Reuseport = true")
 )
 
diff --git a/tcpdialer.go b/tcpdialer.go
index 0023cf6..2317aca 100644
--- a/tcpdialer.go
+++ b/tcpdialer.go
@@ -15,7 +15,7 @@ import (
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -42,7 +42,7 @@ func Dial(addr string) (net.Conn, error) {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -67,7 +67,7 @@ func DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -96,7 +96,7 @@ func DialDualStack(addr string) (net.Conn, error) {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -153,6 +153,9 @@ type TCPDialer struct {
 	// }
 	Resolver Resolver
 
+	// DNSCacheDuration may be used to override the default DNS cache duration (DefaultDNSCacheDuration)
+	DNSCacheDuration time.Duration
+
 	tcpAddrsLock sync.Mutex
 	tcpAddrsMap  map[string]*tcpAddrEntry
 
@@ -166,7 +169,7 @@ type TCPDialer struct {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -193,7 +196,7 @@ func (d *TCPDialer) Dial(addr string) (net.Conn, error) {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -218,7 +221,7 @@ func (d *TCPDialer) DialTimeout(addr string, timeout time.Duration) (net.Conn, e
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -247,7 +250,7 @@ func (d *TCPDialer) DialDualStack(addr string) (net.Conn, error) {
 // This function has the following additional features comparing to net.Dial:
 //
 //   * It reduces load on DNS resolver by caching resolved TCP addressed
-//     for DefaultDNSCacheDuration.
+//     for DNSCacheDuration.
 //   * It dials all the resolved TCP addresses in round-robin manner until
 //     connection is established. This may be useful if certain addresses
 //     are temporarily unreachable.
@@ -272,6 +275,11 @@ func (d *TCPDialer) dial(addr string, dualStack bool, timeout time.Duration) (ne
 		if d.Concurrency > 0 {
 			d.concurrencyCh = make(chan struct{}, d.Concurrency)
 		}
+
+		if d.DNSCacheDuration == 0 {
+			d.DNSCacheDuration = DefaultDNSCacheDuration
+		}
+
 		d.tcpAddrsMap = make(map[string]*tcpAddrEntry)
 		go d.tcpAddrsClean()
 	})
@@ -361,7 +369,7 @@ type tcpAddrEntry struct {
 const DefaultDNSCacheDuration = time.Minute
 
 func (d *TCPDialer) tcpAddrsClean() {
-	expireDuration := 2 * DefaultDNSCacheDuration
+	expireDuration := 2 * d.DNSCacheDuration
 	for {
 		time.Sleep(time.Second)
 		t := time.Now()
@@ -379,7 +387,7 @@ func (d *TCPDialer) tcpAddrsClean() {
 func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, uint32, error) {
 	d.tcpAddrsLock.Lock()
 	e := d.tcpAddrsMap[addr]
-	if e != nil && !e.pending && time.Since(e.resolveTime) > DefaultDNSCacheDuration {
+	if e != nil && !e.pending && time.Since(e.resolveTime) > d.DNSCacheDuration {
 		e.pending = true
 		e = nil
 	}


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..489d983
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,21 @@
+name: Test
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+jobs:
+  test:
+    strategy:
+      matrix:
+        go-version: [1.15.x, 1.16.x]
+        os: [ubuntu-latest, macos-latest, windows-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-go@v2
+        with:
+          go-version: ${{ matrix.go-version }}
+      - run: go version
+      - run: go test ./...
+      - run: go test -race ./...
diff --git a/client_test.go b/client_test.go
index a6be61e..0a61f0d 100644
--- a/client_test.go
+++ b/client_test.go
@@ -1952,6 +1952,70 @@ func TestClientRetryRequestWithCustomDecider(t *testing.T) {
 	}
 }
 
+func TestHostClientTransport(t *testing.T) {
+	t.Parallel()
+
+	ln := fasthttputil.NewInmemoryListener()
+
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			ctx.WriteString("abcd") //nolint:errcheck
+		},
+	}
+	serverStopCh := make(chan struct{})
+	go func() {
+		if err := s.Serve(ln); err != nil {
+			t.Errorf("unexpected error: %s", err)
+		}
+		close(serverStopCh)
+	}()
+
+	c := &HostClient{
+		Addr: "foobar",
+		Transport: func() TransportFunc {
+			c, _ := ln.Dial()
+
+			br := bufio.NewReader(c)
+			bw := bufio.NewWriter(c)
+
+			return func(req *Request, res *Response) error {
+				if err := req.Write(bw); err != nil {
+					return err
+				}
+
+				if err := bw.Flush(); err != nil {
+					return err
+				}
+
+				return res.Read(br)
+			}
+		}(),
+	}
+
+	for i := 0; i < 5; i++ {
+		statusCode, body, err := c.Get(nil, "http://aaaa.com/bbb/cc")
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		if statusCode != StatusOK {
+			t.Fatalf("unexpected status code %d. Expecting %d", statusCode, StatusOK)
+		}
+		if string(body) != "abcd" {
+			t.Fatalf("unexpected body %q. Expecting %q", body, "abcd")
+		}
+	}
+
+	if err := ln.Close(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	select {
+	case <-serverStopCh:
+	case <-time.After(time.Second):
+		t.Fatalf("timeout")
+	}
+}
+
 type writeErrorConn struct {
 	net.Conn
 }
diff --git a/uri_test.go b/uri_test.go
index e2995f6..2faf825 100644
--- a/uri_test.go
+++ b/uri_test.go
@@ -395,6 +395,6 @@ func TestURIWithQuerystringOverride(t *testing.T) {
 	uriString := string(u.RequestURI())
 
 	if uriString != "/?q1=foo&q2=bar&q4=quux" {
-		t.Fatalf("Expected Querystring to be overriden but was %s ", uriString)
+		t.Fatalf("Expected Querystring to be overridden but was %s ", uriString)
 	}
 }