【代码】站点在线人数统计

前言

通过WebSocket协议实现站点在线人数统计

客户端

HTML

<div>当前在线人数 <span id="current_visitors_count">0</span></div>
<script>
    // 建立连接
    const ws = new WebSocket("ws://127.0.0.1:8080");
    // 接受消息
    ws.addEventListener("message", function (event) {
        document.getElementById("current_visitors_count").innerText = event.data;
    });
</script>

服务端

Go

下载依赖

go get github.com/gin-gonic/gin
go get github.com/gorilla/websocket

源代码

client.go

package main

import (
	"bytes"
	"github.com/gorilla/websocket"
	"log"
	"time"
)

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 512
)

var (
	newline = []byte{'\n'}
	space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

type Client struct {
	Hub  *Hub
	Conn *websocket.Conn
	Send chan []byte
}

func (client *Client) readPump() {
	defer func() {
		client.Hub.Unregister <- client
		client.Conn.Close()
	}()
	client.Conn.SetReadLimit(maxMessageSize)
	client.Conn.SetReadDeadline(time.Now().Add(pongWait))
	client.Conn.SetPongHandler(func(string) error {
		client.Conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		_, message, err := client.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error: %v", err)
			}
			break
		}
		message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
		client.Hub.Broadcast <- len(client.Hub.Clients)
	}
}

func (client *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		client.Conn.Close()
	}()
	for {
		select {
		case message, ok := <-client.Send:
			client.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				client.Conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			w, err := client.Conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			n := len(client.Send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-client.Send)
			}
			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			client.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := client.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

hub.go

package main

import (
	"log"
	"strconv"
	"sync"
	"time"
)

type Hub struct {
	Clients    map[*Client]bool
	Broadcast  chan int
	Register   chan *Client
	Unregister chan *Client
}

func NewHub() *Hub {
	return &Hub{
		Clients:    make(map[*Client]bool),
		Broadcast:  make(chan int),
		Register:   make(chan *Client),
		Unregister: make(chan *Client),
	}
}

var once sync.Once
var singleton *Hub

func (hub *Hub) putDataInChan() {
	for {
		time.Sleep(1 * time.Second)
		hub.Broadcast <- 1
	}
}

func (hub *Hub) Run() {
	go hub.putDataInChan()
	for {
		select {
		case client := <-hub.Register:
			hub.Clients[client] = true
		case client := <-hub.Unregister:
			if _, ok := hub.Clients[client]; ok {
				delete(hub.Clients, client)
				close(client.Send)
			}
		case _ = <-hub.Broadcast:
			log.Println("当前在线人数:", len(hub.Clients))
			for client := range hub.Clients {
				var message = strconv.Itoa(len(hub.Clients))
				select {
				case client.Send <- []byte(message):
				default:
					close(client.Send)
					delete(hub.Clients, client)
				}
			}
		}
	}
}

controller.go

package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

var wsupgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func wshandler(hub *Hub, writer http.ResponseWriter, request *http.Request) {
	conn, err := wsupgrader.Upgrade(writer, request, nil)
	if err != nil {
		log.Panicln(err)
		return
	}
	client := &Client{Hub: hub, Conn: conn, Send: make(chan []byte, 256)}
	client.Hub.Register <- client
	go client.writePump()
	// 监听连接关闭事件
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				// 连接正常关闭或正在关闭
				fmt.Println("连接正常关闭:", err)
			} else {
				// 连接异常关闭
				fmt.Println("连接异常关闭:", err)
			}
			client.Hub.Unregister <- client
			break
		}
	}
}

func WebSocketController(context *gin.Context, hub *Hub) {
	wshandler(hub, context.Writer, context.Request)
}

main.go

package main

import "github.com/gin-gonic/gin"

func main() {

	app := gin.Default()

	hub := NewHub()

	go hub.Run()

	app.GET("/", func(context *gin.Context) {
		WebSocketController(context, hub)
	})

	app.Run(":8080")
}

完成

参考文献