Go语言中websocket的使用demo分享-编程学习网

服务端代码 main.go

package main

import (
   "errors"
   "fmt"
   "net/http"
   "sync"
   "time"

   "github.com/gorilla/websocket"
)

// http升级websocket协议的配置
var wsUpgrader = websocket.Upgrader{
   // 允许所有CORS跨域请求
   CheckOrigin: func(r *http.Request) bool {
      return true
   },
}

// 客户端读写消息
type wsMessage struct {
   messageType int
   data        []byte
}

// 客户端连接
type wsConnection struct {
   wsSocket *websocket.Conn // 底层websocket
   inChan   chan *wsMessage // 读队列
   outChan  chan *wsMessage // 写队列

   mutex     sync.Mutex // 避免重复关闭管道
   isClosed  bool
   closeChan chan byte // 关闭通知
}

func (wsConn *wsConnection) wsReadLoop() {
   for {
      // 读一个message
      msgType, data, err := wsConn.wsSocket.ReadMessage()
      if err != nil {
         goto error
      }
      req := &wsMessage{
         msgType,
         data,
      }
      // 放入请求队列
      select {
      case wsConn.inChan <- req:
      case <-wsConn.closeChan:
         goto closed
      }
   }
error:
   wsConn.wsClose()
closed:
}

func (wsConn *wsConnection) wsWriteLoop() {
   for {
      select {
      // 取一个应答
      case msg := <-wsConn.outChan:
         // 写给websocket
         if err := wsConn.wsSocket.WriteMessage(msg.messageType, msg.data); err != nil {
            goto error
         }
      case <-wsConn.closeChan:
         goto closed
      }
   }
error:
   wsConn.wsClose()
closed:
}

func (wsConn *wsConnection) procLoop() {
   // 启动一个gouroutine发送心跳
   go func() {
      for { //不断向客户端写数据，其实没有它也是一样的，客户端可以检测到断开
         time.Sleep(2 * time.Second)
         if err := wsConn.wsWrite(websocket.TextMessage, []byte("heartbeat from server")); err != nil {
            fmt.Println("heartbeat fail")
            wsConn.wsClose()
            break
         }
      }
   }()

   // 这是一个同步处理模型（只是一个例子），如果希望并行处理可以每个请求一个gorutine，注意控制并发goroutine的数量!!!
   for {
      msg, err := wsConn.wsRead()
      if err != nil {
         fmt.Println("read fail")
         break
      }
      fmt.Println(string(msg.data))
      err = wsConn.wsWrite(msg.messageType, msg.data) // 读到数据后，同步的去写数据。应该写成异步
      if err != nil {
         fmt.Println("write fail")
         break
      }
   }
}

func wsHandler(resp http.ResponseWriter, req *http.Request) {
   // 应答客户端告知升级连接为websocket
   wsSocket, err := wsUpgrader.Upgrade(resp, req, nil)
   if err != nil {
      return
   }
   wsConn := &wsConnection{
      wsSocket:  wsSocket,
      inChan:    make(chan *wsMessage, 1000),
      outChan:   make(chan *wsMessage, 1000),
      closeChan: make(chan byte),
      isClosed:  false,
   }

   // 处理器
   go wsConn.procLoop()
   // 读协程
   go wsConn.wsReadLoop()
   // 写协程
   go wsConn.wsWriteLoop()
}

func (wsConn *wsConnection) wsWrite(messageType int, data []byte) error {
   select {
   case wsConn.outChan <- &wsMessage{messageType, data}:
   case <-wsConn.closeChan:
      return errors.New("websocket closed")
   }
   return nil
}

func (wsConn *wsConnection) wsRead() (*wsMessage, error) {
   select {
   case msg := <-wsConn.inChan:
      return msg, nil
   case <-wsConn.closeChan:
   }
   return nil, errors.New("websocket closed")
}

func (wsConn *wsConnection) wsClose() {
   wsConn.wsSocket.Close()

   wsConn.mutex.Lock()
   defer wsConn.mutex.Unlock()
   if !wsConn.isClosed {
      wsConn.isClosed = true
      close(wsConn.closeChan)
   }
}

func main() {
   http.HandleFunc("/ws", wsHandler)
   http.ListenAndServe("0.0.0.0:7777", nil)
}

前端代码 client.html

<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <script>
        window.addEventListener("load", function(evt) {
            const output = document.getElementById("output");
            const input = document.getElementById("input");
            let ws;
            const print = function(message) {
                const d = document.createElement("div");
                d.innerHTML = message;
                output.appendChild(d);
            };
            document.getElementById("open").onclick = function(evt) {
                if (ws) {
                    return false;
                }
                ws = new WebSocket("ws://localhost:7777/ws");
                ws.onopen = function(evt) {
                    print("OPEN");
                }
                ws.onclose = function(evt) {
                    print("CLOSE");
                    ws = null;
                }
                ws.onmessage = function(evt) {
                    print("RESPONSE: " + evt.data);
                }
                ws.onerror = function(evt) {
                    print("ERROR: " + evt.data);
                }
                return false;
            };
            document.getElementById("send").onclick = function(evt) {
                if (!ws) {
                    return false;
                }
                print("SEND: " + input.value);
                ws.send(input.value);
                return false;
            };
            document.getElementById("close").onclick = function(evt) {
                if (!ws) {
                    return false;
                }
                ws.close();
                return false;
            };
        });
    </script>
</head>
<body>
<table>
    <tr><td valign="top" width="50%">
        <p>Click "Open" to create a connection to the server,
            "Send" to send a message to the server and "Close" to close the connection.
            You can change the message and send multiple times.
        </p>
            <form>
                <button id="open">Open</button>
                <button id="close">Close</button>
            <input id="input" type="text" value="Hello world!">
            <button id="send">Send</button>
            </form>
    </td><td valign="top" width="50%">
        <div id="output"></div>
    </td></tr></table>
</body>
</html>

到此这篇关于Go语言中websocket的使用demo分享的文章就介绍到这了,更多相关Go语言 websocket内容请搜索编程网以前的文章或继续浏览下面的相关文章希望大家以后多多支持编程网！