Golang

关注公众号 jb51net

关闭
首页 > 脚本专栏 > Golang > Go语言 websocket

Go语言中websocket的使用demo分享

作者:用户6512516549724

WebSocket是一种在单个TCP连接上进行全双工通信的协议。这篇文章主要和大家分享了一个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内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!

您可能感兴趣的文章:
阅读全文