gin框架实现websocket并广播消息

详细实现方式以及代码下载请前往 https://www.passerma.com/article/82

一、项目预览

二、安装gin框架

go get github.com/gin-gonic/gin

三、安装gorilla/websocket库

go get github.com/gorilla/websocket

四、服务端实现

1.简单版本

新建 main.go 文件

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 <- c
		c.conn.Close()
	}()
	for {
		_, message, err := c.conn.ReadMessage()
		if err == nil {
			c.hub.broadcast <- 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 := <-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 := <-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 := <-h.broadcast:
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
		}
	}
}

// 链接实例
func newHub() *Hub {
	return &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 := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
	go client.writePump()
	go client.readPump()
	client.hub.register <- 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")
}

五、客户端实现

客户端使用原生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>

六、代码下载

详细实现方式以及代码下载请前往 https://www.passerma.com/article/82