建立一個websocket的服務端

package smile

import (
  "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]*wsConnection

var 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,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,}
    // 放入請求佇列,訊息入棧
    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)
}

以上就是本文的全部內容，希望對大家的學習有所幫助，也希望大家多多支援我們。