golang的websocket封装

小小Desperado · · 798 次点击 · · 开始浏览    
这是一个创建于 的文章,其中的信息可能已经有所发展或是发生改变。

实现代码如下:


import (
    "errors"
    "github.com/gorilla/websocket"
    "sync"
)

type Connection struct {
    wsConn *websocket.Conn
    inChan chan []byte
    outChan chan []byte
    closeChan chan byte
    mutex sync.Mutex
    isClosed   bool
}

func InitConnection(wsConn *websocket.Conn)(conn *Connection ,err error){
    conn=&Connection{
        wsConn:wsConn,
        inChan:make(chan []byte,1000),
        outChan:make(chan []byte,1000),
        closeChan:make(chan byte,1),
    }
    go conn.readLoop()
    go conn.writeLoop()
    return
}

func (conn *Connection) ReadMessage()(data []byte,err error){
    select {
    case data=<-conn.inChan:
    case <-conn.closeChan:
        err=errors.New("connection is closed")
    }
    return
}

func (conn *Connection) WriteMessage(data []byte)(err error){
        select {
        case conn.outChan<-data:
        case <-conn.closeChan:
            err=errors.New("connection is closed")
        }
        return
}

func (conn *Connection) Close(){
    //线程安全的Close,可重入
    conn.wsConn.Close()
    conn.mutex.Lock()
    if !conn.isClosed {
        close(conn.closeChan)
        conn.isClosed=true
    }
    conn.mutex.Unlock()

}

func (conn *Connection) readLoop(){
    var(
        data []byte
        err error
    )
    for{
        if _,data,err=conn.wsConn.ReadMessage();err !=nil{
            goto ERR
        }
        select {
        case conn.inChan<-data:
        case <-conn.closeChan:
            goto ERR
        }

    }
    ERR:
        conn.Close()
}
func (conn *Connection) writeLoop(){
    var (
        data []byte
        err error
    )
    for{
        select {
        case data=<-conn.outChan:
        case <-conn.closeChan:
            goto ERR
        }
        if err=conn.wsConn.WriteMessage(websocket.TextMessage,data);err!=nil{
            goto ERR
        }
    }
    ERR:
        conn.Close()
}

首先定义一个结构体

type Connection struct {
    wsConn *websocket.Conn
    inChan chan []byte
    outChan chan []byte
    closeChan chan byte
    mutex sync.Mutex
    isClosed   bool
}

wsConn websocket的长链接的实体
inChan 读数据的channel
outChan 写数据的channel
closeChan 链接关闭的channel
mutex 互斥锁
isClosed 链接关闭标识符

func InitConnection(wsConn *websocket.Conn)(conn *Connection ,err error){
    conn=&Connection{
        wsConn:wsConn,
        inChan:make(chan []byte,1000),
        outChan:make(chan []byte,1000),
        closeChan:make(chan byte,1),
    }
    go conn.readLoop()
    go conn.writeLoop()
    return
}

初始化链接
readLoop 和writeLoop 循环从websocket中读取数据和写入数据

func (conn *Connection) ReadMessage()(data []byte,err error){
    select {
    case data=<-conn.inChan:
    case <-conn.closeChan:
        err=errors.New("connection is closed")
    }
    return
}

ReadMessage() 从inChan中读取数据

func (conn *Connection) WriteMessage(data []byte)(err error){
        select {
        case conn.outChan<-data:
        case <-conn.closeChan:
            err=errors.New("connection is closed")
        }
        return
}

WriteMessage(data []byte) 写入数据传递给outChan ,writeLoop 监听outChan并写入数据

func (conn *Connection) Close(){
    //线程安全的Close,可重入
    conn.wsConn.Close()
    conn.mutex.Lock()
    if !conn.isClosed {
        close(conn.closeChan)
        conn.isClosed=true
    }
    conn.mutex.Unlock()
}

mutex锁住关闭操作 ,避免重复循环关闭链接
关闭链接时,传递closeChan ,同时关闭readLoop 和writeLoop

func (conn *Connection) readLoop(){
    var(
        data []byte
        err error
    )
    for{
        if _,data,err=conn.wsConn.ReadMessage();err !=nil{
            goto ERR
        }
        select {
        case conn.inChan<-data:
        case <-conn.closeChan:
            goto ERR
        }

    }
    ERR:
        conn.Close()
}

func (conn *Connection) writeLoop(){
    var (
        data []byte
        err error
    )
    for{
        select {
        case data=<-conn.outChan:
        case <-conn.closeChan:
            goto ERR
        }
        if err=conn.wsConn.WriteMessage(websocket.TextMessage,data);err!=nil{
            goto ERR
        }
    }
    ERR:
        conn.Close()
}

通过outChan和inChan 传递信息,保证线程的安全。


有疑问加站长微信联系(非本文作者)

本文来自:简书

感谢作者:小小Desperado

查看原文:golang的websocket封装

入群交流(和以上内容无关):加入Go大咖交流群,或添加微信:liuxiaoyan-s 备注:入群;或加QQ群:692541889

798 次点击  
加入收藏 微博
添加一条新回复 (您需要 登录 后才能回复 没有账号 ?)
  • 请尽量让自己的回复能够对别人有帮助
  • 支持 Markdown 格式, **粗体**、~~删除线~~、`单行代码`
  • 支持 @ 本站用户;支持表情(输入 : 提示),见 Emoji cheat sheet
  • 图片支持拖拽、截图粘贴等方式上传