-
Notifications
You must be signed in to change notification settings - Fork 29
/
map.go
140 lines (114 loc) · 2.45 KB
/
map.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
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package conns
import (
"github.com/iwind/TeaGo/types"
"net"
"sync"
)
var SharedMap = NewMap()
type Map struct {
m map[string]map[string]net.Conn // ip => { network_port => Conn }
locker sync.RWMutex
}
func NewMap() *Map {
return &Map{
m: map[string]map[string]net.Conn{},
}
}
func (this *Map) Add(conn net.Conn) {
if conn == nil {
return
}
key, ip, ok := this.connAddr(conn)
if !ok {
return
}
this.locker.Lock()
defer this.locker.Unlock()
connMap, ok := this.m[ip]
if !ok {
this.m[ip] = map[string]net.Conn{key: conn}
} else {
connMap[key] = conn
}
}
func (this *Map) Remove(conn net.Conn) {
if conn == nil {
return
}
key, ip, ok := this.connAddr(conn)
if !ok {
return
}
this.locker.Lock()
defer this.locker.Unlock()
connMap, ok := this.m[ip]
if !ok {
return
}
delete(connMap, key)
if len(connMap) == 0 {
delete(this.m, ip)
}
}
func (this *Map) CountIPConns(ip string) int {
this.locker.RLock()
var l = len(this.m[ip])
this.locker.RUnlock()
return l
}
func (this *Map) CloseIPConns(ip string) {
var conns = []net.Conn{}
this.locker.RLock()
connMap, ok := this.m[ip]
// 复制,防止在Close时产生并发冲突
if ok {
for _, conn := range connMap {
conns = append(conns, conn)
}
}
// 需要在Close之前结束,防止死循环
this.locker.RUnlock()
if ok {
for _, conn := range conns {
// 设置Linger
lingerConn, isLingerConn := conn.(LingerConn)
if isLingerConn {
_ = lingerConn.SetLinger(0)
}
// 关闭
_ = conn.Close()
}
// 这里不需要从 m 中删除,因为关闭时会自然触发回调
}
}
func (this *Map) AllConns() []net.Conn {
this.locker.RLock()
defer this.locker.RUnlock()
var result = []net.Conn{}
for _, m := range this.m {
for _, connInfo := range m {
result = append(result, connInfo)
}
}
return result
}
func (this *Map) connAddr(conn net.Conn) (key string, ip string, ok bool) {
if conn == nil {
return
}
var addr = conn.RemoteAddr()
switch realAddr := addr.(type) {
case *net.TCPAddr:
return addr.Network() + types.String(realAddr.Port), realAddr.IP.String(), true
case *net.UDPAddr:
return addr.Network() + types.String(realAddr.Port), realAddr.IP.String(), true
default:
var s = addr.String()
host, port, err := net.SplitHostPort(s)
if err != nil {
return
}
return addr.Network() + port, host, true
}
}