1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
|
// +build !go1.3
package graceful
import (
"bufio"
"io"
"net"
"net/http"
"sync/atomic"
"github.com/zenazn/goji/graceful/listener"
)
// Middleware provides functionality similar to net/http.Server's
// SetKeepAlivesEnabled in Go 1.3, but in Go 1.2.
func middleware(h http.Handler) http.Handler {
if h == nil {
return nil
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, cn := w.(http.CloseNotifier)
_, fl := w.(http.Flusher)
_, hj := w.(http.Hijacker)
_, rf := w.(io.ReaderFrom)
bw := basicWriter{ResponseWriter: w}
if cn && fl && hj && rf {
h.ServeHTTP(&fancyWriter{bw}, r)
} else {
h.ServeHTTP(&bw, r)
}
if !bw.headerWritten {
bw.maybeClose()
}
})
}
type basicWriter struct {
http.ResponseWriter
headerWritten bool
}
func (b *basicWriter) maybeClose() {
b.headerWritten = true
if atomic.LoadInt32(&closing) != 0 {
b.ResponseWriter.Header().Set("Connection", "close")
}
}
func (b *basicWriter) WriteHeader(code int) {
b.maybeClose()
b.ResponseWriter.WriteHeader(code)
}
func (b *basicWriter) Write(buf []byte) (int, error) {
if !b.headerWritten {
b.maybeClose()
}
return b.ResponseWriter.Write(buf)
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
return b.ResponseWriter
}
// Optimize for the common case of a ResponseWriter that supports all three of
// CloseNotifier, Flusher, and Hijacker.
type fancyWriter struct {
basicWriter
}
func (f *fancyWriter) CloseNotify() <-chan bool {
cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
return cn.CloseNotify()
}
func (f *fancyWriter) Flush() {
fl := f.basicWriter.ResponseWriter.(http.Flusher)
fl.Flush()
}
func (f *fancyWriter) Hijack() (c net.Conn, b *bufio.ReadWriter, e error) {
hj := f.basicWriter.ResponseWriter.(http.Hijacker)
c, b, e = hj.Hijack()
if e == nil {
e = listener.Disown(c)
}
return
}
func (f *fancyWriter) ReadFrom(r io.Reader) (int64, error) {
rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
if !f.basicWriter.headerWritten {
f.basicWriter.maybeClose()
}
return rf.ReadFrom(r)
}
var _ http.CloseNotifier = &fancyWriter{}
var _ http.Flusher = &fancyWriter{}
var _ http.Hijacker = &fancyWriter{}
var _ io.ReaderFrom = &fancyWriter{}
|
