gin框架实现websocket并广播消息

2023-02-01 23:08 go 194 0

一、项目预览

项目预览

二、安装gin框架

go get github.com/gin-gonic/gin

三、安装gorilla/websocket库

go get github.com/gorilla/websocket

四、服务端实现

新建 main.go 文件

1.简单版本

package main

import (
    "encoding/json"
    "net/http"

    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
)

// 连接成功结构体
type connectedDataS struct {
    Event            string `json:"_event"`
    HistoryClientNum int    `json:"historyClientNum"`
    OnlineClientNum  int    `json:"onlineClientNum"`
}

// 消息结构体
type messageDataS struct {
    Event   string `json:"_event"`
    Message string `json:"message"`
}

// ws upGrader
var upGrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

// 全部客户端
var clients = make(map[*websocket.Conn]bool)

// 存储数据
var (
    historyClientNum = 0
)

// 广播消息
func broadcastMsg(mt int, message []byte) {
    for client := range clients {
        client.WriteMessage(mt, message)
    }
}

// 存储当前客户端
func saveClient(client *websocket.Conn) {
    clients[client] = true
    // 更新历史客户端数
    historyClientNum++
    // 广播进入
    var data = connectedDataS{
        Event:            "connected",
        HistoryClientNum: historyClientNum,
        OnlineClientNum:  len(clients),
    }
    var dataJson, _ = json.Marshal(data)
    broadcastMsg(websocket.TextMessage, dataJson)
}

// 删除当前客户端
func delClient(client *websocket.Conn) {
    // 断开连接删除client
    delete(clients, client)
    // 广播离开
    var data = connectedDataS{
        Event:            "connected",
        HistoryClientNum: historyClientNum,
        OnlineClientNum:  len(clients),
    }
    var dataJson, _ = json.Marshal(data)
    broadcastMsg(websocket.TextMessage, dataJson)
}

// 发送消息
func sendMsg(mt int, message []byte) {
    var data = messageDataS{
        Event:   "message",
        Message: string(message),
    }
    dataJson, _ := json.Marshal(data)
    broadcastMsg(mt, dataJson)
}

// 创建客户端实例
func createWs(c *gin.Context) {
    // 创建ws连接客户端
    client, err := upGrader.Upgrade(c.Writer, c.Request, nil)
    if err != nil {
        return
    }
    // 连接成功后存储当前客户端
    saveClient(client)
    defer func() {
        client.Close()
        delClient(client)
    }()
    for {
        // 监听接受信息
        mt, message, err := client.ReadMessage()
        if err == nil {
            sendMsg(mt, message)
        } else {
            break
        }
    }
}

func main() {
    router := gin.Default()
    router.LoadHTMLFiles("index.html")
    router.GET("/", func(c *gin.Context) {
        c.HTML(http.StatusOK, "index.html", nil)
    })
    router.GET("/ws", createWs)
    router.Run(":9999")
}

2.正常版本

package main

import (
	"encoding/json"
	"log"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// 连接成功结构体
type connectedDataS struct {
	Event            string `json:"_event"`
	HistoryClientNum int    `json:"historyClientNum"`
	OnlineClientNum  int    `json:"onlineClientNum"`
}

// 消息结构体
type messageDataS struct {
	Event   string `json:"_event"`
	Message string `json:"message"`
}

// 存储数据
var (
	historyClientNum = 0
)

// 广播消息
func broadcastMsg(hub *Hub, mt int, message []byte) {
	for client := range hub.clients {
		client.conn.WriteMessage(mt, message)
	}
}

// 客户端实例
type Client struct {
	hub  *Hub
	conn *websocket.Conn
	send chan []byte
}

// 读取数据通道
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister &lt;- c
		c.conn.Close()
	}()
	for {
		_, message, err := c.conn.ReadMessage()
		if err == nil {
			c.hub.broadcast &lt;- message
		} else {
			break
		}
	}
}

// 发送数据通道
func (c *Client) writePump() {
	defer func() {
		c.conn.Close()
	}()
	for message := range c.send {
		var data = messageDataS{
			Event:   "message",
			Message: string(message),
		}
		dataJson, _ := json.Marshal(data)
		c.conn.WriteMessage(websocket.TextMessage, dataJson)
	}
}

// ws hub
type Hub struct {
	clients    map[*Client]bool
	broadcast  chan []byte
	register   chan *Client
	unregister chan *Client
}

// 运行hub
func (h *Hub) run() {
	for {
		select {
		case client := &lt;-h.register:
			h.clients[client] = true
			// 更新历史客户端数
			historyClientNum++
			// 广播进入
			var data = connectedDataS{
				Event:            "connected",
				HistoryClientNum: historyClientNum,
				OnlineClientNum:  len(h.clients),
			}
			var dataJson, _ = json.Marshal(data)
			broadcastMsg(h, websocket.TextMessage, dataJson)
		case client := &lt;-h.unregister:
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.send)
				// 广播离开
				var data = connectedDataS{
					Event:            "connected",
					HistoryClientNum: historyClientNum,
					OnlineClientNum:  len(h.clients),
				}
				var dataJson, _ = json.Marshal(data)
				broadcastMsg(h, websocket.TextMessage, dataJson)
			}
		case message := &lt;-h.broadcast:
			for client := range h.clients {
				select {
				case client.send &lt;- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
		}
	}
}

// 链接实例
func newHub() *Hub {
	return &amp;Hub{
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[*Client]bool),
	}
}

// ws upGrader
var upGrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func createWs(c *gin.Context, hub *Hub) {
	conn, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println(err)
		return
	}
	client := &amp;Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
	go client.writePump()
	go client.readPump()
	client.hub.register &lt;- client
}

func main() {
	hub := newHub()
	go hub.run()
	router := gin.Default()
	router.LoadHTMLFiles("index.html")
	router.GET("/", func(c *gin.Context) {
		c.HTML(http.StatusOK, "index.html", nil)
	})
	router.GET("/ws", func(ctx *gin.Context) {
		createWs(ctx, hub)
	})
	router.Run(":9999")
}

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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

五、客户端实现

客户端使用原生WebSocket
新建 index.html 文件

<!DOCTYPE html>
<html lang="zh-cn">

<head>
  <meta charset="UTF-8">
  <title>Gin-WebSocket</title>
  <style>
    html,
    body {
      padding: 0;
      margin: 0;
      height: 100%;
      width: 100%;
      display: flex;
      flex-direction: column;
      justify-content: center;
      align-items: center;
    }

    #message {
      display: flex;
      flex-direction: column;
      justify-content: center;
      align-items: center;
    }
  </style>
</head>

<body>
  <span id="connected">WebSocket连接中...</span>
  <span id="onlineClientNum"></span>
  <span id="historyClientNum"></span>
  <div id="message"></div>
  <div><input id="input" /><button onclick="send()">发送</button></div>
  <script>
    const ws = new WebSocket("ws://localhost:9999/ws");
    // 连接成功时触发
    ws.onopen = () => {
      document.getElementById('connected').innerText = 'WebSocket连接成功'
    }
    // 接收到消息时触发  
    ws.onmessage = (e) => {
      const data = JSON.parse(e.data)
      switch (data._event) {
        case "connected":
          document.getElementById('historyClientNum').innerText = '历史：' + data.historyClientNum
          document.getElementById('onlineClientNum').innerText = '在线：' + data.onlineClientNum
          break;
        case "message":
          const span = document.createElement('span')
          span.innerText = data.message
          document.getElementById('message').appendChild(span)
          break;
        default:
          break;
      }
    };
    // 连接关闭时触发  
    ws.onclose = (e) => {
      document.getElementById('connected').innerText = 'WebSocket连接中...'
    };
    // 发送消息
    const send = () => {
      const value = document.getElementById('input').value
      if (value) {
        document.getElementById('input').value = ''
        ws.send(value)
      }
    }
  </script>
</body>

</html>

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

六、代码下载

gin-ws.zip

暂无评论，我会出手

说点什么

最多上传8张图片，仅支持jpg，png格式图片，单张大小5MB以内！

用户名：

passerma博客