Go语言学习：搭建简易的websocket服务器

安装gorilla/websocket第三方开源库

go get -u github.com/gorilla/websocket

创建三个go文件，分别是mian.go server.go user.go

mian.go:程序入口，服务器启动

package main

import "log"

func main() {
	log.Println("start server")

	server.Start()
}

server.go：先建立一个http连接，然后使用gorilla/websocket第三方包中的功能将http连接升级成一个websocket长连接，这样就能实现客户端和服务端双工通信

package main

import (
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

type Server struct {
	onlineUserMap map[string]*User
	userMapLock   sync.RWMutex
}

var server = Server{
	onlineUserMap: make(map[string]*User),
}

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

func handler(w http.ResponseWriter, r *http.Request) {
	wsConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
	}

	NewUser(wsConn)
}

func (server *Server) Start() {
	http.HandleFunc("/", handler)
	http.ListenAndServe(":8888", nil)
}

func (server *Server) addOnlineUserMap(user *User) {
	server.userMapLock.Lock()
	server.onlineUserMap[user.addr] = user
	server.userMapLock.Unlock()
}

func (server *Server) deleteOnlineUserMap(user *User) {
	server.userMapLock.Lock()
	delete(server.onlineUserMap, user.addr)
	server.userMapLock.Unlock()
}

user.go：主要是通过websocket.Conn对象创建一个User对象，并且分别启动两个协程来收发消息

package main

import (
	"log"

	"github.com/gorilla/websocket"
)

type User struct {
	addr     string
	wsConn   *websocket.Conn
	sendChan chan []byte
}

func NewUser(conn *websocket.Conn) {
	var user = &User{
		addr:     conn.RemoteAddr().String(),
		wsConn:   conn,
		sendChan: make(chan []byte),
	}

	user.online()
	go user.recvMessage()
	go user.sendMessage()
}

func (user *User) recvMessage() {
	defer user.offline()
	for {
		_, p, err := user.wsConn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}

		user.sendChan <- p

		log.Printf("Recv [%s] msg:%s", user.addr, p)
	}
}

func (user *User) sendMessage() {
	defer user.offline()
	for {
		buf := <-user.sendChan

		err := user.wsConn.WriteMessage(1, buf)
		if err != nil {
			log.Println(err)
			return
		}

		log.Printf("Send [%s] msg:%s", user.addr, buf)
	}
}

func (user *User) online() {
	server.addOnlineUserMap(user)

	log.Printf("[%s] 上线了", user.addr)
}

func (user *User) offline() {
	user.wsConn.Close()
	server.deleteOnlineUserMap(user)

	log.Printf("[%s] 下线了", user.addr)
}

在vscode终端启动我们的websocket服务器

PS D:\LearningFolder\GoWorkSpace\pro\websocketserver> go build
PS D:\LearningFolder\GoWorkSpace\pro\websocketserver> .\websocketserver.exe
2022/08/14 13:46:05 start server

利用websocket在线测试,点击进入

 服务器日志log