From bb71be0271ff319bfa6b59b16eef4b0c61058b7b Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Sat, 24 Aug 2013 02:28:58 +0000
Subject: [PATCH] Implement outstanding query detection.

Copied from the standard library and adapted to Go DNS.
---
 client.go         |  7 +++----
 client_test.go    | 38 ++++++++++++++++++++++++++++++++++++++
 singleinflight.go |  2 +-
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/client.go b/client.go
index 26319b1b..847af4f5 100644
--- a/client.go
+++ b/client.go
@@ -48,15 +48,14 @@ func (c *Client) exchangeMerge(m *Msg, a string, s net.Conn) (r *Msg, rtt time.D
 	}
 	// This adds a bunch of garbage, TODO(miek).
 	t := "nop"
-	if t1, ok := TypeToString[r.Question[0].Qtype]; ok {
+	if t1, ok := TypeToString[m.Question[0].Qtype]; ok {
 		t = t1
 	}
 	cl := "nop"
-	if cl1, ok := ClassToString[r.Question[0].Qclass]; ok {
+	if cl1, ok := ClassToString[m.Question[0].Qclass]; ok {
 		cl = cl1
 	}
-	key := r.Question[0].Name + t + cl
-	r, rtt, err, shared := c.group.Do(key, func() (*Msg, time.Duration, error) {
+	r, rtt, err, shared := c.group.Do(m.Question[0].Name+t+cl, func() (*Msg, time.Duration, error) {
 		if s == nil {
 			return c.exchange(m, a)
 		}
diff --git a/client_test.go b/client_test.go
index 823dae56..04d43b7f 100644
--- a/client_test.go
+++ b/client_test.go
@@ -41,6 +41,44 @@ func TestClientEDNS0(t *testing.T) {
 	}
 }
 
+func TestInflight(t *testing.T) {
+	m := new(Msg)
+	m.SetQuestion("miek.nl.", TypeDNSKEY)
+
+	c := new(Client)
+	c.Inflight = true
+	nr := 10
+	ch := make(chan time.Duration)
+	for i := 0; i < nr; i++ {
+		go func() {
+			_, rtt, _ := c.Exchange(m, "37.251.95.53:53")
+			ch <- rtt
+		}()
+	}
+	i := 0
+	var first time.Duration
+	// With inflight *all* rtt are identical, and by doing actual lookups
+	// the changes that this is a coincidence is small.
+Loop:
+	for {
+		select {
+		case rtt := <-ch:
+			if i == 0 {
+				first = rtt
+			} else {
+				if first != rtt {
+					t.Log("All rtt should be equal")
+					t.Fail()
+				}
+			}
+			i++
+			if i == 10 {
+				break Loop
+			}
+		}
+	}
+}
+
 func TestClientTsigAXFR(t *testing.T) {
 	m := new(Msg)
 	m.SetAxfr("miek.nl.")
diff --git a/singleinflight.go b/singleinflight.go
index e16d2e9c..9573c7d0 100644
--- a/singleinflight.go
+++ b/singleinflight.go
@@ -46,7 +46,7 @@ func (g *singleflight) Do(key string, fn func() (*Msg, time.Duration, error)) (v
 	g.m[key] = c
 	g.Unlock()
 
-	c.val, rtt, c.err = fn()
+	c.val, c.rtt, c.err = fn()
 	c.wg.Done()
 
 	g.Lock()