net: add Dialer.KeepAlive option

LGTM=rsc
R=rsc
CC=golang-codereviews
https://golang.org/cl/68380043

Author: bradfitz

src/pkg/net/dial.go

@@ -44,6 +44,12 @@ type Dialer struct {
 	// destination is a host name that has multiple address family
 	// DNS records.
 	DualStack bool
+
+	// KeepAlive specifies the keep-alive period for an active
+	// network connection.
+	// If zero, keep-alives are not enabled. Network protocols
+	// that do not support keep-alives ignore this field.
+	KeepAlive time.Duration
 }
 
 // Return either now+Timeout or Deadline, whichever comes first.
@@ -162,9 +168,19 @@ func (d *Dialer) Dial(network, address string) (Conn, error) {
 			return dialMulti(network, address, d.LocalAddr, ras, deadline)
 		}
 	}
-	return dial(network, ra.toAddr(), dialer, d.deadline())
+	c, err := dial(network, ra.toAddr(), dialer, d.deadline())
+	if d.KeepAlive > 0 && err == nil {
+		if tc, ok := c.(*TCPConn); ok {
+			tc.SetKeepAlive(true)
+			tc.SetKeepAlivePeriod(d.KeepAlive)
+			testHookSetKeepAlive()
+		}
+	}
+	return c, err
 }
 
+var testHookSetKeepAlive = func() {} // changed by dial_test.go
+
 // dialMulti attempts to establish connections to each destination of
 // the list of addresses. It will return the first established
 // connection and close the other connections. Otherwise it returns

src/pkg/net/dial_test.go

@@ -555,3 +555,36 @@ func TestDialDualStackLocalhost(t *testing.T) {
 		}
 	}
 }
+
+func TestDialerKeepAlive(t *testing.T) {
+	ln := newLocalListener(t)
+	defer ln.Close()
+	defer func() {
+		testHookSetKeepAlive = func() {}
+	}()
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			c.Close()
+		}
+	}()
+	for _, keepAlive := range []bool{false, true} {
+		got := false
+		testHookSetKeepAlive = func() { got = true }
+		var d Dialer
+		if keepAlive {
+			d.KeepAlive = 30 * time.Second
+		}
+		c, err := d.Dial("tcp", ln.Addr().String())
+		if err != nil {
+			t.Fatal(err)
+		}
+		c.Close()
+		if got != keepAlive {
+			t.Errorf("Dialer.KeepAlive = %v: SetKeepAlive called = %v, want %v", d.KeepAlive, got, !got)
+		}
+	}
+}