/
conn_handler.go
76 lines (66 loc) · 1.4 KB
/
conn_handler.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
package server
import (
"net"
"sync"
"sync/atomic"
)
type serverState int32
const (
stateDefault = int32(iota)
stateStopped
)
func (s *serverState) StopOnce() bool {
return atomic.CompareAndSwapInt32((*int32)(s), stateDefault, stateStopped)
}
func (s *serverState) Stopped() bool {
return atomic.LoadInt32((*int32)(s)) == stateStopped
}
type connHandler struct {
store map[net.Conn]struct{}
mu sync.Mutex
wg sync.WaitGroup
state serverState
}
func (s *connHandler) remove(conn net.Conn) {
s.mu.Lock()
delete(s.store, conn)
s.wg.Done()
s.mu.Unlock()
}
func (s *connHandler) handle(handler ConnectionHandler, conn net.Conn) {
defer s.remove(conn)
handler.HandleConnection(conn)
}
func (s *connHandler) Handle(handler ConnectionHandler, conn net.Conn) {
// fast exit: don't attempt to acquire the mutex or
// handle the conn if shutdown
if s.state.Stopped() {
conn.Close()
return
}
s.mu.Lock()
defer s.mu.Unlock()
// recheck the state now that we've locked the mutex
// as we may have been blocked on call to Shutdown()
if s.state.Stopped() {
conn.Close()
return
}
// lazily initialize the store
if s.store == nil {
s.store = make(map[net.Conn]struct{})
}
s.store[conn] = struct{}{}
s.wg.Add(1)
go s.handle(handler, conn)
}
func (s *connHandler) Shutdown() {
if s.state.StopOnce() {
s.mu.Lock()
for c := range s.store {
c.Close()
}
s.mu.Unlock()
s.wg.Wait()
}
}