forked from lesismal/nbio
-
Notifications
You must be signed in to change notification settings - Fork 0
/
engine.go
361 lines (298 loc) · 8.19 KB
/
engine.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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
// Copyright 2020 lesismal. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
package nbio
import (
"context"
"net"
"runtime"
"sync"
"time"
"unsafe"
"github.com/soar4/nbio/logging"
"github.com/soar4/nbio/timer"
)
const (
// DefaultReadBufferSize .
DefaultReadBufferSize = 1024 * 64
// DefaultMaxWriteBufferSize .
DefaultMaxWriteBufferSize = 1024 * 1024
// DefaultMaxConnReadTimesPerEventLoop .
DefaultMaxConnReadTimesPerEventLoop = 3
// DefaultUDPReadTimeout .
DefaultUDPReadTimeout = 120 * time.Second
)
var (
// MaxOpenFiles .
MaxOpenFiles = 1024 * 1024 * 2
)
// Config Of Engine.
type Config struct {
// Name describes your gopher name for logging, it's set to "NB" by default.
Name string
// Network is the listening protocol, used with Addrs together.
// tcp* supported only by now, there's no plan for other protocol such as udp,
// because it's too easy to write udp server/client.
Network string
// Addrs is the listening addr list for a nbio server.
// if it is empty, no listener created, then the Engine is used for client by default.
Addrs []string
// NPoller represents poller goroutine num, it's set to runtime.NumCPU() by default.
NPoller int
// ReadBufferSize represents buffer size for reading, it's set to 64k by default.
ReadBufferSize int
// MaxWriteBufferSize represents max write buffer size for Conn, it's set to 1m by default.
// if the connection's Send-Q is full and the data cached by nbio is
// more than MaxWriteBufferSize, the connection would be closed by nbio.
MaxWriteBufferSize int
// MaxConnReadTimesPerEventLoop represents max read times in one poller loop for one fd
MaxConnReadTimesPerEventLoop int
// LockListener represents listener's goroutine to lock thread or not, it's set to false by default.
LockListener bool
// LockPoller represents poller's goroutine to lock thread or not, it's set to false by default.
LockPoller bool
// EpollMod sets the epoll mod, EPOLLLT by default.
EpollMod uint32
// EPOLLONESHOT .
EPOLLONESHOT uint32
// UDPReadTimeout sets the timeout for udp sessions.
UDPReadTimeout time.Duration
// TimerExecute sets the executor for timer callbacks.
TimerExecute func(f func())
// Listen is used to create listener for Engine.
Listen func(network, addr string) (net.Listener, error)
// ListenUDP is used to create udp listener for Engine.
ListenUDP func(network string, laddr *net.UDPAddr) (*net.UDPConn, error)
}
// Gopher keeps old type to compatible with new name Engine.
type Gopher = Engine
func NewGopher(conf Config) *Gopher {
return NewEngine(conf)
}
// Engine is a manager of poller.
type Engine struct {
*timer.Timer
sync.WaitGroup
Name string
Execute func(f func())
TimerExecute func(f func())
mux sync.Mutex
wgConn sync.WaitGroup
network string
addrs []string
listen func(network, addr string) (net.Listener, error)
listenUDP func(network string, laddr *net.UDPAddr) (*net.UDPConn, error)
pollerNum int
readBufferSize int
maxWriteBufferSize int
maxConnReadTimesPerEventLoop int
udpReadTimeout time.Duration
epollMod uint32
epollOneshot uint32
lockListener bool
lockPoller bool
connsStd map[*Conn]struct{}
connsUnix []*Conn
listeners []*poller
pollers []*poller
onOpen func(c *Conn)
onClose func(c *Conn, err error)
onRead func(c *Conn)
onData func(c *Conn, data []byte)
onReadBufferAlloc func(c *Conn) []byte
onReadBufferFree func(c *Conn, buffer []byte)
// onWriteBufferFree func(c *Conn, buffer []byte)
beforeRead func(c *Conn)
afterRead func(c *Conn)
beforeWrite func(c *Conn)
onStop func()
}
// Stop closes listeners/pollers/conns/timer.
func (g *Engine) Stop() {
for _, l := range g.listeners {
l.stop()
}
g.mux.Lock()
conns := g.connsStd
g.connsStd = map[*Conn]struct{}{}
connsUnix := g.connsUnix
g.mux.Unlock()
g.wgConn.Done()
for c := range conns {
if c != nil {
cc := c
g.Async(func() {
cc.Close()
})
}
}
for _, c := range connsUnix {
if c != nil {
cc := c
g.Async(func() {
cc.Close()
})
}
}
g.wgConn.Wait()
time.Sleep(time.Second / 5)
g.onStop()
g.Timer.Stop()
for i := 0; i < g.pollerNum; i++ {
g.pollers[i].stop()
}
g.Wait()
logging.Info("NBIO[%v] stop", g.Name)
}
// Shutdown stops Engine gracefully with context.
func (g *Engine) Shutdown(ctx context.Context) error {
ch := make(chan struct{})
go func() {
g.Stop()
close(ch)
}()
select {
case <-ch:
case <-ctx.Done():
return ctx.Err()
}
return nil
}
// AddConn adds conn to a poller.
func (g *Engine) AddConn(conn net.Conn) (*Conn, error) {
c, err := NBConn(conn)
if err != nil {
return nil, err
}
p := g.pollers[c.Hash()%len(g.pollers)]
p.addConn(c)
return c, nil
}
// OnOpen registers callback for new connection.
func (g *Engine) OnOpen(h func(c *Conn)) {
if h == nil {
panic("invalid nil handler")
}
g.onOpen = func(c *Conn) {
g.wgConn.Add(1)
h(c)
}
}
// OnClose registers callback for disconnected.
func (g *Engine) OnClose(h func(c *Conn, err error)) {
if h == nil {
panic("invalid nil handler")
}
g.onClose = func(c *Conn, err error) {
g.Async(func() {
defer g.wgConn.Done()
h(c, err)
})
}
}
// OnRead registers callback for reading event.
func (g *Engine) OnRead(h func(c *Conn)) {
g.onRead = h
}
// OnData registers callback for data.
func (g *Engine) OnData(h func(c *Conn, data []byte)) {
if h == nil {
panic("invalid nil handler")
}
g.onData = h
}
// OnReadBufferAlloc registers callback for memory allocating.
func (g *Engine) OnReadBufferAlloc(h func(c *Conn) []byte) {
if h == nil {
panic("invalid nil handler")
}
g.onReadBufferAlloc = h
}
// OnReadBufferFree registers callback for memory release.
func (g *Engine) OnReadBufferFree(h func(c *Conn, b []byte)) {
if h == nil {
panic("invalid nil handler")
}
g.onReadBufferFree = h
}
// OnWriteBufferRelease registers callback for write buffer memory release.
// func (g *Engine) OnWriteBufferRelease(h func(c *Conn, b []byte)) {
// if h == nil {
// panic("invalid nil handler")
// }
// g.onWriteBufferFree = h
// }
// BeforeRead registers callback before syscall.Read
// the handler would be called on windows.
func (g *Engine) BeforeRead(h func(c *Conn)) {
if h == nil {
panic("invalid nil handler")
}
g.beforeRead = h
}
// AfterRead registers callback after syscall.Read
// the handler would be called on *nix.
func (g *Engine) AfterRead(h func(c *Conn)) {
if h == nil {
panic("invalid nil handler")
}
g.afterRead = h
}
// BeforeWrite registers callback befor syscall.Write and syscall.Writev
// the handler would be called on windows.
func (g *Engine) BeforeWrite(h func(c *Conn)) {
if h == nil {
panic("invalid nil handler")
}
g.beforeWrite = h
}
// OnStop registers callback before Engine is stopped.
func (g *Engine) OnStop(h func()) {
if h == nil {
panic("invalid nil handler")
}
g.onStop = h
}
// PollerBuffer returns Poller's buffer by Conn, can be used on linux/bsd.
func (g *Engine) PollerBuffer(c *Conn) []byte {
return c.p.ReadBuffer
}
func (g *Engine) initHandlers() {
g.wgConn.Add(1)
g.OnOpen(func(c *Conn) {})
g.OnClose(func(c *Conn, err error) {})
// g.OnRead(func(c *Conn, b []byte) ([]byte, error) {
// n, err := c.Read(b)
// if n > 0 {
// return b[:n], err
// }
// return nil, err
// })
g.OnData(func(c *Conn, data []byte) {})
g.OnReadBufferAlloc(g.PollerBuffer)
g.OnReadBufferFree(func(c *Conn, buffer []byte) {})
// g.OnWriteBufferRelease(func(c *Conn, buffer []byte) {})
g.BeforeRead(func(c *Conn) {})
g.AfterRead(func(c *Conn) {})
g.BeforeWrite(func(c *Conn) {})
g.OnStop(func() {})
if g.Execute == nil {
g.Execute = func(f func()) {
defer func() {
if err := recover(); err != nil {
const size = 64 << 10
buf := make([]byte, size)
buf = buf[:runtime.Stack(buf, false)]
logging.Error("execute failed: %v\n%v\n", err, *(*string)(unsafe.Pointer(&buf)))
}
}()
f()
}
}
}
func (g *Engine) borrow(c *Conn) []byte {
return g.onReadBufferAlloc(c)
}
func (g *Engine) payback(c *Conn, buffer []byte) {
g.onReadBufferFree(c, buffer)
}