code review 69260044: net/http: add optional Server.ConnState callback

src/pkg/net/http/serve_test.go

 // Copyright 2010 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // End-to-end serving tests
 
 package http_test
 
 import (
         "bufio"
         "bytes"
         "crypto/tls"
         "errors"
         "fmt"
         "io"
         "io/ioutil"
         "log"
         "net"
         . "net/http"
         "net/http/httptest"
         "net/http/httputil"
         "net/url"
         "os"
         "os/exec"
         "reflect"
         "runtime"
         "strconv"
         "strings"
+        "sync"
         "sync/atomic"
         "syscall"
         "testing"
         "time"
 )
 
 type dummyAddr string
 type oneConnListener struct {
         conn net.Conn
 }
@@ skipping 2197 matching lines @@
         res := ExportAppendTime(b[:0], t1)
         t2, err := ParseTime(string(res))
         if err != nil {
                 t.Fatalf("Error parsing time: %s", err)
         }
         if !t1.Equal(t2) {
                 t.Fatalf("Times differ; expected: %v, got %v (%s)", t1, t2, string(res))
         }
 }
 
+func TestServerConnState(t *testing.T) {
+        defer afterTest(t)
+        handler := map[string]func(w ResponseWriter, r *Request){
+                "/": func(w ResponseWriter, r *Request) {
+                        fmt.Fprintf(w, "Hello.")
+                },
+                "/close": func(w ResponseWriter, r *Request) {
+                        w.Header().Set("Connection", "close")
+                        fmt.Fprintf(w, "Hello.")
+                },
+                "/hijack": func(w ResponseWriter, r *Request) {
+                        c, _, _ := w.(Hijacker).Hijack()
+                        c.Write([]byte("HTTP/1.0 200 OK\r\nConnection: close\r\n\r\nHello."))
+                        c.Close()
+                },
+        }
+        ts := httptest.NewUnstartedServer(HandlerFunc(func(w ResponseWriter, r *Request) {
+                handler[r.URL.Path](w, r)
+        }))
+        defer ts.Close()
+
+        type connIDAndState struct {
+                connID int
+                state  ConnState
+        }
+        var mu sync.Mutex // guard stateLog and connID
+        var stateLog []connIDAndState
+        var connID = map[net.Conn]int{}
+
+        ts.Config.ConnState = func(c net.Conn, state ConnState) {
+                if c == nil {
+                        t.Error("nil conn seen in state %s", state)
+                        return
+                }
+                mu.Lock()
+                defer mu.Unlock()
+                id, ok := connID[c]
+                if !ok {
+                        id = len(connID) + 1
+                        connID[c] = id
+                }
+                stateLog = append(stateLog, connIDAndState{id, state})
+        }
+        ts.Start()
+
+        mustGet(t, ts.URL+"/")
+        mustGet(t, ts.URL+"/close")
+
+        mustGet(t, ts.URL+"/")
+        mustGet(t, ts.URL+"/", "Connection", "close")
+
+        mustGet(t, ts.URL+"/hijack")
+
+        want := []connIDAndState{
+                {1, StateNew},
+                {1, StateActive},
+                {1, StateIdle},
+                {1, StateActive},
+                {1, StateClosed},
+
+                {2, StateNew},
+                {2, StateActive},
+                {2, StateIdle},
+                {2, StateActive},
+                {2, StateClosed},
+
+                {3, StateNew},
+                {3, StateActive},
+                {3, StateHijacked},
+        }
+        logString := func(l []connIDAndState) string {
+                var b bytes.Buffer
+                for _, cs := range l {
+                        fmt.Fprintf(&b, "[%d %s] ", cs.connID, cs.state)
+                }
+                return b.String()
+        }
+
+        for i := 0; i < 5; i++ {
+                time.Sleep(time.Duration(i) * 50 * time.Millisecond)
+                mu.Lock()
+                match := reflect.DeepEqual(stateLog, want)
+                mu.Unlock()
+                if match {
+                        return
+                }
+        }
+
+        mu.Lock()
+        t.Errorf("Unexpected events.\nGot log: %s\n   Want: %s\n", logString(stateLog), logString(want))
+        mu.Unlock()
+}
+
+func mustGet(t *testing.T, url string, headers ...string) {
+        req, err := NewRequest("GET", url, nil)
+        if err != nil {
+                t.Fatal(err)
+        }
+        for len(headers) > 0 {
+                req.Header.Add(headers[0], headers[1])
+                headers = headers[2:]
+        }
+        res, err := DefaultClient.Do(req)
+        if err != nil {
+                t.Errorf("Error fetching %s: %v", url, err)
+                return
+        }
+        _, err = ioutil.ReadAll(res.Body)
+        defer res.Body.Close()
+        if err != nil {
+                t.Errorf("Error reading %s: %v", url, err)
+        }
+}
+
 func BenchmarkClientServer(b *testing.B) {
         b.ReportAllocs()
         b.StopTimer()
         ts := httptest.NewServer(HandlerFunc(func(rw ResponseWriter, r *Request) {
                 fmt.Fprintf(rw, "Hello world.\n")
         }))
         defer ts.Close()
         b.StartTimer()
 
         for i := 0; i < b.N; i++ {
@@ skipping 296 matching lines @@
                 closec: make(chan bool, 1),
         }
         ln := &oneConnListener{conn: conn}
         for i := 0; i < b.N; i++ {
                 conn.Reader = bytes.NewReader(req)
                 ln.conn = conn
                 Serve(ln, h)
                 <-conn.closec
         }
 }