From 32448f39cded212c821fc3e114ec6e6cf68ada7a Mon Sep 17 00:00:00 2001
From: Michael Haro <mharo@google.com>
Date: Tue, 5 May 2015 23:23:53 -0700
Subject: [PATCH] Cleanup Client.exchange

Make Client's exchange function easier to read by moving timeout logic
into separate functions.

Start the timers closer to where they're used so that time from other
logic doesn't impact the deadlines.
---
 client.go | 44 ++++++++++++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/client.go b/client.go
index cdab4432..0c27214d 100644
--- a/client.go
+++ b/client.go
@@ -122,31 +122,39 @@ func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err erro
 	return r, rtt, nil
 }
 
-func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
-	timeout := dnsTimeout
-	var co *Conn
+func (c *Client) dialTimeout() time.Duration {
 	if c.DialTimeout != 0 {
-		timeout = c.DialTimeout
+		return c.DialTimeout
 	}
+	return dnsTimeout
+}
+
+func (c *Client) readTimeout() time.Duration {
+	if c.ReadTimeout != 0 {
+		return c.ReadTimeout
+	}
+	return dnsTimeout
+}
+
+func (c *Client) writeTimeout() time.Duration {
+	if c.WriteTimeout != 0 {
+		return c.WriteTimeout
+	}
+	return dnsTimeout
+}
+
+func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
+	var co *Conn
 	if c.Net == "" {
-		co, err = DialTimeout("udp", a, timeout)
+		co, err = DialTimeout("udp", a, c.dialTimeout())
 	} else {
-		co, err = DialTimeout(c.Net, a, timeout)
+		co, err = DialTimeout(c.Net, a, c.dialTimeout())
 	}
 	if err != nil {
 		return nil, 0, err
 	}
-	timeout = dnsTimeout
-	if c.ReadTimeout != 0 {
-		timeout = c.ReadTimeout
-	}
-	co.SetReadDeadline(time.Now().Add(timeout))
-	timeout = dnsTimeout
-	if c.WriteTimeout != 0 {
-		timeout = c.WriteTimeout
-	}
-	co.SetWriteDeadline(time.Now().Add(timeout))
 	defer co.Close()
+
 	opt := m.IsEdns0()
 	// If EDNS0 is used use that for size.
 	if opt != nil && opt.UDPSize() >= MinMsgSize {
@@ -156,6 +164,10 @@ func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err erro
 	if opt == nil && c.UDPSize >= MinMsgSize {
 		co.UDPSize = c.UDPSize
 	}
+
+	co.SetReadDeadline(time.Now().Add(c.readTimeout()))
+	co.SetWriteDeadline(time.Now().Add(c.writeTimeout()))
+
 	co.TsigSecret = c.TsigSecret
 	if err = co.WriteMsg(m); err != nil {
 		return nil, 0, err