forked from v2fly/v2ray-core
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pool.go
125 lines (109 loc) · 2.27 KB
/
pool.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
package internal
import (
"net"
"sync"
"time"
v2net "v2ray.com/core/common/net"
"v2ray.com/core/common/signal"
)
type ConnectionRecyler interface {
Put(ConnectionId, net.Conn)
}
type ConnectionId struct {
Local v2net.Address
Remote v2net.Address
RemotePort v2net.Port
}
func NewConnectionId(source v2net.Address, dest v2net.Destination) ConnectionId {
return ConnectionId{
Local: source,
Remote: dest.Address,
RemotePort: dest.Port,
}
}
type ExpiringConnection struct {
conn net.Conn
expire time.Time
}
func (o *ExpiringConnection) Expired() bool {
return o.expire.Before(time.Now())
}
type Pool struct {
sync.Mutex
connsByDest map[ConnectionId][]*ExpiringConnection
cleanupOnce signal.Once
}
func NewConnectionPool() *Pool {
return &Pool{
connsByDest: make(map[ConnectionId][]*ExpiringConnection),
}
}
func (o *Pool) Get(id ConnectionId) net.Conn {
o.Lock()
defer o.Unlock()
list, found := o.connsByDest[id]
if !found {
return nil
}
connIdx := -1
for idx, conn := range list {
if !conn.Expired() {
connIdx = idx
break
}
}
if connIdx == -1 {
return nil
}
listLen := len(list)
conn := list[connIdx]
if connIdx != listLen-1 {
list[connIdx] = list[listLen-1]
}
list = list[:listLen-1]
o.connsByDest[id] = list
return conn.conn
}
func (o *Pool) Cleanup() {
defer o.cleanupOnce.Reset()
for len(o.connsByDest) > 0 {
time.Sleep(time.Second * 5)
expiredConns := make([]net.Conn, 0, 16)
o.Lock()
for dest, list := range o.connsByDest {
validConns := make([]*ExpiringConnection, 0, len(list))
for _, conn := range list {
if conn.Expired() {
expiredConns = append(expiredConns, conn.conn)
} else {
validConns = append(validConns, conn)
}
}
if len(validConns) != len(list) {
o.connsByDest[dest] = validConns
}
}
o.Unlock()
for _, conn := range expiredConns {
conn.Close()
}
}
}
func (o *Pool) Put(id ConnectionId, conn net.Conn) {
expiringConn := &ExpiringConnection{
conn: conn,
expire: time.Now().Add(time.Second * 4),
}
o.Lock()
defer o.Unlock()
list, found := o.connsByDest[id]
if !found {
list = []*ExpiringConnection{expiringConn}
} else {
list = append(list, expiringConn)
}
o.connsByDest[id] = list
o.cleanupOnce.Do(func() {
go o.Cleanup()
})
}