forked from yangsongbai/infini-framework
/
hub.go
executable file
·141 lines (116 loc) · 3.27 KB
/
hub.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
// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package websocket
import (
"encoding/json"
"github.com/huminghe/infini-framework/core/global"
"github.com/huminghe/infini-framework/core/logger"
"github.com/huminghe/infini-framework/core/stats"
"strings"
"time"
)
// Hub maintains the set of active connections and broadcasts messages to the
// connections.
type Hub struct {
// Registered connections.
connections map[*WebsocketConnection]bool
// Inbound messages from the connections.
broadcast chan string
// Register requests from the connections.
register chan *WebsocketConnection
// Unregister requests from connections.
unregister chan *WebsocketConnection
// Command handlers
handlers map[string]WebsocketHandlerFunc
//Command usage tips
usage map[string]string
}
// WebsocketHandlerFunc define the func to handle websocket
type WebsocketHandlerFunc func(c *WebsocketConnection, array []string)
var h = Hub{
broadcast: make(chan string, 5),
register: make(chan *WebsocketConnection),
unregister: make(chan *WebsocketConnection),
connections: make(map[*WebsocketConnection]bool),
handlers: make(map[string]WebsocketHandlerFunc),
usage: make(map[string]string),
}
var runningHub = false
// Register command handlers
func (h *Hub) registerHandlers() {
HandleWebSocketCommand("HELP", "type `help` for more commands", helpCommand)
}
// InitWebSocket start websocket
func InitWebSocket() {
if !runningHub {
h.registerHandlers()
go h.runHub()
}
}
// HandleWebSocketCommand used to register command and handler
func HandleWebSocketCommand(cmd, usage string, handler func(c *WebsocketConnection, array []string)) {
cmd = strings.ToLower(strings.TrimSpace(cmd))
h.handlers[cmd] = WebsocketHandlerFunc(handler)
h.usage[cmd] = usage
}
func (h *Hub) runHub() {
//TODO error handler, parameter assertion
if global.Env().IsDebug {
go func() {
t := time.NewTicker(time.Duration(30) * time.Second)
for {
select {
case <-t.C:
h.broadcast <- "testing websocket broadcast"
}
}
}()
}
//handle connect, disconnect, broadcast
for {
select {
case c := <-h.register:
h.connections[c] = true
c.WritePrivateMessage(global.Env().GetWelcomeMessage())
js, _ := json.Marshal(logger.GetLoggingConfig())
c.WriteMessage(ConfigMessage, string(js))
case c := <-h.unregister:
if _, ok := h.connections[c]; ok {
delete(h.connections, c)
close(c.signalChannel)
}
case m := <-h.broadcast:
h.broadcastMessage(m)
}
}
}
func (h *Hub) broadcastMessage(msg string) {
if len(msg) <= 0 {
return
}
for c := range h.connections {
c.Broadcast(msg)
}
}
// BroadcastMessage send broadcast message to channel and record stats
func BroadcastMessage(msg string) {
select {
case h.broadcast <- msg:
stats.Increment("websocket", "sended")
default:
stats.Increment("websocket", "dropped")
}
}
func getHelpMessage() string {
//list all commands and usage
help := "COMMAND LIST\n"
for k, v := range h.usage {
help += (k + ", " + v + "\n")
}
return help
}
// Help command returns command help information
func helpCommand(c *WebsocketConnection, a []string) {
c.WritePrivateMessage(getHelpMessage())
}