websocket 应用
websocket-golang 应用
目录结构
client：客户端
service：服务端
memory：服务端管理客户端长链接信息
dimain：通用结构体定义
websocket_test：测试用例
service
package websocket
import (
"net/http"
"time"
"github.com/gorilla/websocket"
)
/*
* webSocket 服务端
*/
// 升级 http 请求为 websocket 请求
var (
upgrade = &websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
)
type Serve struct {
conn *websocket.Conn // ws 链接
remoteAddr string // 目标地址
identity string // 身份标识
updateTime int64 // 更新请求时间
handleMessage func(message *Message, serve *Serve) // 处理消息
}
// NewService 将 http 请求升级为 websocket 请求
func NewService(w http.ResponseWriter, r *http.Request, handleMessage func(message *Message, serve *Serve)) {
conn, err := upgrade.Upgrade(w, r, nil)
if err != nil {
return
}
// notice： 验证 identity 是否存在
identity := r.Header.Get("identify")
s := &Serve{
conn: conn,
remoteAddr: r.RemoteAddr,
identity: identity,
handleMessage: handleMessage,
}
// 更新时间
s.UpdateTime()
memory := GetClientsManager(nil)
if tmp, ok := memory.GetClient(identity); ok {
tmp.CloseAndDelMemory()
}
// 添加到 memory 进行管理
memory.AddClient(s.identity, s)
go s.readLoop()
return
}
// readLoop 读取消息
func (s *Serve) readLoop() {
for {
var tempMessage = &Message{}
// 检测到连接是否关闭
if s.conn == nil {
return
}
// 读取消息
err := s.conn.ReadJSON(tempMessage)
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
return
} else if err != nil {
return
}
// 更新时间
s.UpdateTime()
// 处理消息
s.handleMessage(tempMessage, s)
}
}
// Send 发送消息
func (s *Serve) Send(message interface{}) {
if err := s.conn.WriteJSON(message); err != nil {
}
}
// CloseAndDelMemory 关闭连接并从内存中删除
func (s *Serve) CloseAndDelMemory() {
s.Close()
GetClientsManager(nil).DelClient(s.identity)
}
// Close 关闭连接
func (s *Serve) Close() {
err := s.conn.Close()
if err != nil {
return
}
}
// UpdateTime 更新时间
func (s *Serve) UpdateTime() {
s.updateTime = time.Now().Unix()
}
// GetUpdateTime 获取更新时间
func (s *Serve) GetUpdateTime() int64 {
return s.updateTime
}
client
package websocket
import (
"net/http"
"net/url"
"github.com/gorilla/websocket"
)
/*
* websocket 客户端
*/
type Client struct {
conn *websocket.Conn
identify string
remoteAddr string
handleMessage func(*Message, *Client)
}
// NewClient 创建客户端
func NewClient(remoteIp, identify string, handleMessage func(*Message, *Client)) *Client {
u := url.URL{Scheme: "ws", Host: remoteIp, Path: "/ws"}
header := http.Header{
"identify": []string{identify},
}
conn, _, err := websocket.DefaultDialer.Dial(u.String(), header)
if err != nil {
return nil
}
c := &Client{
conn: conn,
identify: identify,
remoteAddr: remoteIp,
handleMessage: handleMessage,
}
// 开启读协程
go c.readLoop()
return c
}
// readLoop 读协程
func (cli *Client) readLoop() {
for {
var tempMessage = &Message{}
if cli.conn == nil {
return
}
err := cli.conn.ReadJSON(tempMessage)
if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
return
} else if err != nil {
return
}
// 处理消息
cli.handleMessage(tempMessage, cli)
}
}
// Close 关闭连接
func (cli *Client) Close() {
err := cli.conn.Close()
if err != nil {
return
}
}
// Send 发送消息
func (cli *Client) Send(msg interface{}) {
if cli.conn == nil {
return
}
err := cli.conn.WriteJSON(msg)
if err != nil {
return
}
return
}
memory
package websocket
import (
"sync"
"time"
)
/*
* 管理客户端长链接
*/
var cliManager = ClientsManager{
clientsMap: make(map[string]*Serve),
rwMutex: sync.RWMutex{},
maxTime: 30, // 默认设置最长无响应时长 30 秒
checkTime: 5, // 默认检测时间 5 秒
}
type ClientsManager struct {
clientsMap map[string]*Serve
rwMutex sync.RWMutex
once sync.Once
maxTime int64 // 最长无响应时长 单位 秒
checkTime time.Duration // 检测时间 单位 秒
}
// GetClientsManager 获取客户端管理器
func GetClientsManager(memoryCfg *MemoryConfig) *ClientsManager {
cliManager.once.Do(func() {
if memoryCfg != nil {
// TODO 设置配置，可以使用 opt 设计模式
cliManager.maxTime = memoryCfg.MaxTime
cliManager.checkTime = memoryCfg.CheckTime
}
go cliManager.isLive()
})
return &cliManager
}
// AddClient 添加客户端链接信息
func (m *ClientsManager) AddClient(identify string, client *Serve) {
m.rwMutex.Lock()
defer m.rwMutex.Unlock()
m.clientsMap[identify] = client
}
// GetClient 获取客户端链接信息
func (m *ClientsManager) GetClient(identify string) (client *Serve, ok bool) {
m.rwMutex.RLock()
defer m.rwMutex.RUnlock()
client, ok = m.clientsMap[identify]
return
}
// DelClient 删除客户端链接信息
func (m *ClientsManager) DelClient(identify string) {
m.rwMutex.Lock()
defer m.rwMutex.Unlock()
delete(m.clientsMap, identify)
}
// 检测客户端链接是否存活
func (m *ClientsManager) isLive() {
ticker := time.NewTicker(m.checkTime * time.Second)
for {
select {
case <-ticker.C:
m.checkClients()
}
}
}
// 处理未响应的客户端链接
func (m *ClientsManager) checkClients() {
m.rwMutex.Lock()
defer m.rwMutex.Unlock()
for k, v := range m.clientsMap {
if time.Now().Unix()-v.GetUpdateTime() > m.maxTime {
v.Close()
delete(m.clientsMap, k)
}
}
}
dimain
package websocket
import "time"
type Message struct {
Type int `json:"type,omitempty"` // 事件
Data interface{} `json:"data,omitempty"` // 消息体
}
type MemoryConfig struct {
MaxTime int64
CheckTime time.Duration
}
websocket_test
package websocket
import (
"log"
"net/http"
"testing"
"time"
"github.com/google/uuid"
)
// 服务端
func Test_Websocket(t *testing.T) {
http.HandleFunc("/ws", WebSocketService)
err := http.ListenAndServe(":8080", nil)
if err != nil {
t.Fatal(err)
}
t.Log("end")
}
// 客户端发送消息
func Test_WebsocketClient(t *testing.T) {
c := NewClient("127.0.0.1:8080", uuid.NewString(), handleMessageByClient)
for i := 0; i < 20; i++ {
msg := &Message{
Type: 1,
Data: i,
}
time.Sleep(time.Second * 5)
c.Send(msg)
}
// time.Sleep(time.Second * 10)
msg := &Message{
Type: 2,
Data: "end",
}
c.Send(msg)
}
// 客户端处理消息
func handleMessageByClient(message *Message, cli *Client) {
log.Println("handleMessage:", message)
if message.Type == 2 {
cli.Close()
}
}
// 服务端处理消息
func handleMessageByServe(message *Message, s *Serve) {
// 处理消息
if message.Type == 1 {
log.Println("handleMessage:", message)
msg := &Message{
Type: 1,
Data: "收到",
}
s.Send(msg)
} else {
msg := &Message{
Type: 2,
Data: "end",
}
s.Send(msg)
s.CloseAndDelMemory()
}
}
func WebSocketService(w http.ResponseWriter, r *http.Request) {
// 验证请求是否安全
identity := r.Header.Get("identify")
if identity == "" {
log.Printf("identity: 没有参数")
return
}
// 创建 websocket 服务端
NewService(w, r, handleMessageByServe)
}