创建一个websocket的服务端

package smileimport ( "errors" "log" "net/http" "sync" "time" "github.com/gorilla/websocket")const ( // 允许等待的写入时间 writeWait = 10 * time.Second // Time allowed to read the next pong message from the peer. pongWait = 60 * time.Second // Send pings to peer with this period. Must be less than pongWait. pingPeriod = (pongWait * 9) / 10 // Maximum message size allowed from peer. maxMessageSize = 512)// 最大的连接ID，每次连接都加1 处理var maxConnId int64// 客户端读写消息type wsMessage struct { // websocket.TextMessage 消息类型 messageType int data []byte}// ws 的所有连接// 用于广播var wsConnAll map[int64]*wsConnectionvar upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, // 允许所有的CORS 跨域请求，正式环境可以关闭 CheckOrigin: func(r *http.Request) bool { return true },}// 客户端连接type wsConnection struct { wsSocket *websocket.Conn // 底层websocket inChan chan *wsMessage // 读队列 outChan chan *wsMessage // 写队列 mutex sync.Mutex // 避免重复关闭管道,加锁处理 isClosed bool closeChan chan byte // 关闭通知 id int64}func wsHandler(resp http.ResponseWriter, req *http.Request) { // 应答客户端告知升级连接为websocket wsSocket, err := upgrader.Upgrade(resp, req, nil) if err != nil { log.Println("升级为websocket失败", err.Error()) return } maxConnId++ // TODO 如果要控制连接数可以计算，wsConnAll长度 // 连接数保持一定数量，超过的部分不提供服务 wsConn := &wsConnection{ wsSocket: wsSocket, inChan: make(chan *wsMessage, 1000), outChan: make(chan *wsMessage, 1000), closeChan: make(chan byte), isClosed: false, id: maxConnId, } wsConnAll[maxConnId] = wsConn log.Println("当前在线人数", len(wsConnAll)) // 处理器,发送定时信息，避免意外关闭 go wsConn.processLoop() // 读协程 go wsConn.wsReadLoop() // 写协程 go wsConn.wsWriteLoop()}// 处理队列中的消息func (wsConn *wsConnection) processLoop() { // 处理消息队列中的消息 // 获取到消息队列中的消息，处理完成后，发送消息给客户端 for { msg, err := wsConn.wsRead() if err != nil { log.Println("获取消息出现错误", err.Error()) break } log.Println("接收到消息", string(msg.data)) // 修改以下内容把客户端传递的消息传递给处理程序 err = wsConn.wsWrite(msg.messageType, msg.data) if err != nil { log.Println("发送消息给客户端出现错误", err.Error()) break } }}// 处理消息队列中的消息func (wsConn *wsConnection) wsReadLoop() { // 设置消息的最大长度 wsConn.wsSocket.SetReadLimit(maxMessageSize) wsConn.wsSocket.SetReadDeadline(time.Now().Add(pongWait)) for { // 读一个message msgType, data, err := wsConn.wsSocket.ReadMessage() if err != nil { websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) log.Println("消息读取出现错误", err.Error()) wsConn.close() return } req := &wsMessage{ msgType, data, } // 放入请求队列,消息入栈 select { case wsConn.inChan <- req: case <-wsConn.closeChan: return } }}// 发送消息给客户端func (wsConn *wsConnection) wsWriteLoop() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() }() for { select { // 取一个应答 case msg := <-wsConn.outChan: // 写给websocket if err := wsConn.wsSocket.WriteMessage(msg.messageType, msg.data); err != nil { log.Println("发送消息给客户端发生错误", err.Error()) // 切断服务 wsConn.close() return } case <-wsConn.closeChan: // 获取到关闭通知 return case <-ticker.C: // 出现超时情况 wsConn.wsSocket.SetWriteDeadline(time.Now().Add(writeWait)) if err := wsConn.wsSocket.WriteMessage(websocket.PingMessage, nil); err != nil { return } } }}// 写入消息到队列中func (wsConn *wsConnection) wsWrite(messageType int, data []byte) error { select { case wsConn.outChan <- &wsMessage{messageType, data}: case <-wsConn.closeChan: return errors.New("连接已经关闭") } return nil}// 读取消息队列中的消息func (wsConn *wsConnection) wsRead() (*wsMessage, error) { select { case msg := <-wsConn.inChan: // 获取到消息队列中的消息 return msg, nil case <-wsConn.closeChan: } return nil, errors.New("连接已经关闭")}// 关闭连接func (wsConn *wsConnection) close() { log.Println("关闭连接被调用了") wsConn.wsSocket.Close() wsConn.mutex.Lock() defer wsConn.mutex.Unlock() if wsConn.isClosed == false { wsConn.isClosed = true // 删除这个连接的变量 delete(wsConnAll, wsConn.id) close(wsConn.closeChan) }}// 启动程序func StartWebsocket(addrPort string) { wsConnAll = make(map[int64]*wsConnection) http.HandleFunc("/ws", wsHandler) http.ListenAndServe(addrPort, nil)}

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持。