net/http: ensure ConnState for StateNew fires before Server.Serve returns

The addition of Server.ConnState provides all the necessary
hooks to stop a Server gracefully, but StateNew previously
could fire concurrently with Serve exiting (as it does when
its net.Listener is closed). This previously meant one
couldn't use a WaitGroup incremented in the StateNew hook
along with calling Wait after Serve. Now you can.

Update #4674

LGTM=bradfitz
R=bradfitz
CC=golang-codereviews
https://golang.org/cl/70410044

Author: rcrowley

src/pkg/net/http/serve_test.go

@@ -2372,6 +2372,27 @@ func TestServerKeepAlivesEnabled(t *testing.T) {
 	}
 }
 
+func TestServerConnStateNew(t *testing.T) {
+	sawNew := false // if the test is buggy, we'll race on this variable.
+	srv := &Server{
+		ConnState: func(c net.Conn, state ConnState) {
+			if state == StateNew {
+				sawNew = true // testing that this write isn't racy
+			}
+		},
+		Handler: HandlerFunc(func(w ResponseWriter, r *Request) {}), // irrelevant
+	}
+	srv.Serve(&oneConnListener{
+		conn: &rwTestConn{
+			Reader: strings.NewReader("GET / HTTP/1.1\r\nHost: foo\r\n\r\n"),
+			Writer: ioutil.Discard,
+		},
+	})
+	if !sawNew { // testing that this read isn't racy
+		t.Error("StateNew not seen")
+	}
+}
+
 func BenchmarkClientServer(b *testing.B) {
 	b.ReportAllocs()
 	b.StopTimer()

src/pkg/net/http/server.go

@@ -1090,7 +1090,6 @@ func (c *conn) setState(nc net.Conn, state ConnState) {
 // Serve a new connection.
 func (c *conn) serve() {
 	origConn := c.rwc // copy it before it's set nil on Close or Hijack
-	c.setState(origConn, StateNew)
 	defer func() {
 		if err := recover(); err != nil {
 			const size = 64 << 10
@@ -1722,6 +1721,7 @@ func (srv *Server) Serve(l net.Listener) error {
 		if err != nil {
 			continue
 		}
+		c.setState(c.rwc, StateNew) // before Serve can return
 		go c.serve()
 	}
 }