forked from jumpserver/koko
/
broadcaster.go
62 lines (49 loc) · 1.38 KB
/
broadcaster.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
package httpd
import "github.com/meowgen/koko/pkg/logger"
const MaxMessageLen = 1024
func NewBroadcaster() *broadcaster {
return &broadcaster{
userConns: make(map[string]*UserWebsocket),
enterChannel: make(chan *UserWebsocket),
leavingChannel: make(chan *UserWebsocket),
messageChannel: make(chan *Message, MaxMessageLen),
checkChannel: make(chan string),
resultChannel: make(chan *UserWebsocket),
}
}
type broadcaster struct {
userConns map[string]*UserWebsocket
enterChannel chan *UserWebsocket
leavingChannel chan *UserWebsocket
messageChannel chan *Message
checkChannel chan string
resultChannel chan *UserWebsocket
}
func (b *broadcaster) Start() {
for {
select {
case conn := <-b.enterChannel:
b.userConns[conn.Uuid] = conn
logger.Infof("Ws[%s] enter", conn.Uuid)
case conn := <-b.leavingChannel:
delete(b.userConns, conn.Uuid)
logger.Infof("Ws[%s] leave", conn.Uuid)
case sid := <-b.checkChannel:
b.resultChannel <- b.userConns[sid]
case <-b.messageChannel:
}
}
}
func (b *broadcaster) EnterUserWebsocket(c *UserWebsocket) {
b.enterChannel <- c
}
func (b *broadcaster) LeaveUserWebsocket(c *UserWebsocket) {
b.leavingChannel <- c
}
func (b *broadcaster) Broadcast(msg *Message) {
b.messageChannel <- msg
}
func (b *broadcaster) GetUserWebsocket(sid string) *UserWebsocket {
b.checkChannel <- sid
return <-b.resultChannel
}