Fix race condition in getTCPAddrs

1 files changed, 9 insertions, 6 deletions

diff --git a/tcpdialer.go b/tcpdialer.go
index 261f62e..20a766f 100644
--- a/tcpdialer.go
+++ b/tcpdialer.go
@@ -358,8 +358,8 @@ type tcpAddrEntry struct {
 	addrs    []net.TCPAddr
 	addrsIdx uint32
 
+	pending     int32
 	resolveTime time.Time
-	pending     bool
 }
 
 // DefaultDNSCacheDuration is the duration for caching resolved TCP addresses
@@ -384,9 +384,11 @@ func (d *TCPDialer) tcpAddrsClean() {
 func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, uint32, error) {
 	item, exist := d.tcpAddrsMap.Load(addr)
 	e, ok := item.(*tcpAddrEntry)
-	if exist && ok && e != nil && !e.pending && time.Since(e.resolveTime) > d.DNSCacheDuration {
-		e.pending = true
-		e = nil
+	if exist && ok && e != nil && time.Since(e.resolveTime) > d.DNSCacheDuration {
+		// Only let one goroutine re-resolve at a time.
+		if atomic.SwapInt32(&e.pending, 1) == 0 {
+			e = nil
+		}
 	}
 
 	if e == nil {
@@ -394,8 +396,9 @@ func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, uin
 		if err != nil {
 			item, exist := d.tcpAddrsMap.Load(addr)
 			e, ok = item.(*tcpAddrEntry)
-			if exist && ok && e != nil && e.pending {
-				e.pending = false
+			if exist && ok && e != nil {
+				// Set pending to 0 so another goroutine can retry.
+				atomic.StoreInt32(&e.pending, 0)
 			}
 			return nil, 0, err
 		}