From 870a59089c42e2199f25e3a18ab3cd93047ab38a Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Tue, 4 Dec 2018 17:59:08 +1030
Subject: [PATCH] Prevent timeout in TestConcurrentExchanges (#862)

This eliminates the possibility of deadlocking the handler and also
simplifies the test considerably.
---
 client_test.go    | 38 ++++++++++++++------------------------
 singleinflight.go | 10 +++++++---
 2 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/client_test.go b/client_test.go
index 311c5cea..e7bb8ba3 100644
--- a/client_test.go
+++ b/client_test.go
@@ -7,7 +7,6 @@ import (
 	"net"
 	"strconv"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 )
@@ -521,17 +520,13 @@ func TestConcurrentExchanges(t *testing.T) {
 	cases[0] = new(Msg)
 	cases[1] = new(Msg)
 	cases[1].Truncated = true
-	for _, m := range cases {
-		block := make(chan struct{})
-		waiting := make(chan struct{})
 
+	for _, m := range cases {
 		mm := m // redeclare m so as not to trip the race detector
 		handler := func(w ResponseWriter, req *Msg) {
 			r := mm.Copy()
 			r.SetReply(req)
 
-			waiting <- struct{}{}
-			<-block
 			w.WriteMsg(r)
 		}
 
@@ -546,29 +541,24 @@ func TestConcurrentExchanges(t *testing.T) {
 
 		m := new(Msg)
 		m.SetQuestion("miek.nl.", TypeSRV)
+
 		c := &Client{
 			SingleInflight: true,
 		}
-		r := make([]*Msg, 2)
+		// Force this client to always return the same request,
+		// even though we're querying sequentially. Running the
+		// Exchange calls below concurrently can fail due to
+		// goroutine scheduling, but this simulates the same
+		// outcome.
+		c.group.dontDeleteForTesting = true
 
-		var wg sync.WaitGroup
-		wg.Add(len(r))
-		for i := 0; i < len(r); i++ {
-			go func(i int) {
-				defer wg.Done()
-				r[i], _, _ = c.Exchange(m.Copy(), addrstr)
-				if r[i] == nil {
-					t.Errorf("response %d is nil", i)
-				}
-			}(i)
+		r := make([]*Msg, 2)
+		for i := range r {
+			r[i], _, _ = c.Exchange(m.Copy(), addrstr)
+			if r[i] == nil {
+				t.Errorf("response %d is nil", i)
+			}
 		}
-		select {
-		case <-waiting:
-		case <-time.After(time.Second):
-			t.FailNow()
-		}
-		close(block)
-		wg.Wait()
 
 		if r[0] == r[1] {
 			t.Errorf("got same response, expected non-shared responses")
diff --git a/singleinflight.go b/singleinflight.go
index 9573c7d0..febcc300 100644
--- a/singleinflight.go
+++ b/singleinflight.go
@@ -23,6 +23,8 @@ type call struct {
 type singleflight struct {
 	sync.Mutex                  // protects m
 	m          map[string]*call // lazily initialized
+
+	dontDeleteForTesting bool // this is only to be used by TestConcurrentExchanges
 }
 
 // Do executes and returns the results of the given function, making
@@ -49,9 +51,11 @@ func (g *singleflight) Do(key string, fn func() (*Msg, time.Duration, error)) (v
 	c.val, c.rtt, c.err = fn()
 	c.wg.Done()
 
-	g.Lock()
-	delete(g.m, key)
-	g.Unlock()
+	if !g.dontDeleteForTesting {
+		g.Lock()
+		delete(g.m, key)
+		g.Unlock()
+	}
 
 	return c.val, c.rtt, c.err, c.dups > 0
 }