forked from wolfeidau/gbc
/
connection.go
159 lines (129 loc) · 3.87 KB
/
connection.go
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
package gbc
import (
"bufio"
"io"
"net"
"sync"
)
// noLimit is an effective infinite upper bound for io.LimitedReader
const (
defaultBufferSize = 1024
)
var (
// value used when the next listener is allocated
bufferSize = defaultBufferSize
)
// SetBufferSize assigs the buffer size used when Listen is next called, this must be done before that happens of course
func SetBufferSize(size int) {
bufferSize = size
}
// BeforeAccept callback which is invoked on each buffered connection before being provided to Accept, returning an error
// will trigger closing of the connection
type BeforeAccept func(bconn *BufferedConn) error
// BufferedConnListener wraps a net Listener and provides BufferedConn rather than net.Conn via accept callback
type BufferedConnListener struct {
net.Listener
bufioReaderPool sync.Pool
bufioWriterPool sync.Pool
bsize int
before BeforeAccept
}
// Listen announces on the local network address laddr. The network net must be a stream-oriented
// network: "tcp", "tcp4", "tcp6", "unix" or "unixpacket". See net.Dial for the syntax of laddr.
func Listen(network, laddr string) (*BufferedConnListener, error) {
ln, err := net.Listen(network, laddr)
if err != nil {
return nil, err
}
return &BufferedConnListener{Listener: ln, bsize: bufferSize}, nil
}
// Accept accept
func (bcl *BufferedConnListener) Accept() (c net.Conn, err error) {
c, err = bcl.Listener.Accept()
if err != nil {
return
}
c = newBufferedConn(c, bcl)
// cast the connection
bconn := c.(*BufferedConn)
if bcl.before != nil {
err = bcl.before(bconn)
}
if err != nil {
c.Close()
c = nil
}
return
}
// SetBeforeAccept assign a before accept function which can intercept, use and reject connections
func (bcl *BufferedConnListener) SetBeforeAccept(before BeforeAccept) {
bcl.before = before
}
// BufferedConn simple buffered connection
type BufferedConn struct {
net.Conn // the underlying net connection
bufwr *bufio.ReadWriter // buffered reading/writing from rwc
bcl *BufferedConnListener // the listener who is managing the buffer pools
}
func newBufferedConn(rwc net.Conn, bcl *BufferedConnListener) *BufferedConn {
c := &BufferedConn{Conn: rwc, bcl: bcl}
br := c.bcl.newBufioReader(c.Conn)
bw := c.bcl.newBufioWriter(c.Conn)
c.bufwr = bufio.NewReadWriter(br, bw)
return c
}
// Read read from the underlying buffered readwriter to avoid issues
func (c *BufferedConn) Read(b []byte) (int, error) {
return c.bufwr.Read(b)
}
// ReadWriter access the read writer for this connection
func (c *BufferedConn) ReadWriter() *bufio.ReadWriter {
return c.bufwr
}
func (c *BufferedConn) finalFlush() {
if c.bufwr != nil {
c.bufwr.Flush()
// Steal the bufio.Reader (~4KB worth of memory) and its associated
// reader for a future connection.
c.bcl.putBufioReader(c.bufwr.Reader)
// Steal the bufio.Writer (~4KB worth of memory) and its associated
// writer for a future connection.
c.bcl.putBufioWriter(c.bufwr.Writer)
c.bufwr = nil
}
}
// Close the connection.
func (c *BufferedConn) Close() (err error) {
log.Debugf("closing %s", c.RemoteAddr())
// panic("woops")
c.finalFlush()
if c.Conn != nil {
err = c.Conn.Close()
c.Conn = nil
}
return
}
func (bcl *BufferedConnListener) newBufioReader(r io.Reader) *bufio.Reader {
if v := bcl.bufioReaderPool.Get(); v != nil {
br := v.(*bufio.Reader)
br.Reset(r)
return br
}
return bufio.NewReaderSize(r, bcl.bsize)
}
func (bcl *BufferedConnListener) putBufioReader(br *bufio.Reader) {
br.Reset(nil)
bcl.bufioReaderPool.Put(br)
}
func (bcl *BufferedConnListener) newBufioWriter(w io.Writer) *bufio.Writer {
if v := bcl.bufioWriterPool.Get(); v != nil {
bw := v.(*bufio.Writer)
bw.Reset(w)
return bw
}
return bufio.NewWriterSize(w, bcl.bsize)
}
func (bcl *BufferedConnListener) putBufioWriter(bw *bufio.Writer) {
bw.Reset(nil)
bcl.bufioWriterPool.Put(bw)
}