websocket小案例

浏览器客户端：

    
    
    





    

    websocket测试程序，消息又客户端发送到server然后原封不动的返回，server使用go实现
    

    

    

    
    
    
    

---

版本一：

package main

import (
    "net/http"
    "github.com/gorilla/websocket"
)

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

func wsHandler(writer http.ResponseWriter, request *http.Request) {
    var (
        conn *websocket.Conn
        err error
        //msgType int
        data []byte
    )
    //完成握手应答
    if conn, err = upgrader.Upgrade(writer, request, nil); err != nil {
        return
    }

    //数据收发
    for {
        //数据类型有text、binary，此处选text
        if _, data, err = conn.ReadMessage(); err != nil {
            goto ERR
        }
        if err = conn.WriteMessage(websocket.TextMessage, data); err != nil {
            goto ERR
        }
    }
ERR:
    conn.Close()

}

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

版本一未做优化

版本二：

connection.go：

package impl

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

type Connection struct {
    wsConn *websocket.Conn   // 底层websocket
    inChan chan []byte      // 读队列
    outChan chan []byte     //  写队列
    closeChan chan byte    // 关闭通知
    isClosed bool
    mutex sync.Mutex       // 避免重复关闭管道
}

//封装websocket长连接
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() {
    // wsConn.Close是线程安全的，可重入的（可以多次关闭）
    conn.wsConn.Close()
    //一个chan只能关闭一次（所以要保证这行代码只执行一次）
    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
        }
        //阻塞在这里，等待inChan有空闲位置
        select {
        case conn.inChan <- data:
        case <- conn.closeChan: //当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()
}

server.go：

package main

import (
    "net/http"
    "github.com/gorilla/websocket"
    "./impl"
    "time"
)

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

func wsHandler(writer http.ResponseWriter, request *http.Request) {
    var (
        wsConn *websocket.Conn
        err error
        //msgType int
        data []byte
        conn *impl.Connection
    )
    //完成握手应答
    if wsConn, err = upgrader.Upgrade(writer, request, nil); err != nil {
        return
    }

    if conn, err = impl.InitConnection(wsConn); err != nil {
        goto ERR
    }

    // 不停发送心跳信息
    go func() {
        var (
            err error
        )
        for {
            if err = conn.WriteMessage([]byte("heartbeat")); err != nil {
                return
            }
            time.Sleep(time.Second)
        }
    }()

    for {
        if data, err = conn.ReadMessage(); err != nil {
            goto ERR
        }
        if err = conn.WriteMessage(data); err != nil {
            goto ERR
        }
    }

ERR:
    // 关闭连接的操作
    conn.Close()

    /*//数据收发
    for {
        //数据类型有text、binary，此处选text
        if _, data, err = conn.ReadMessage(); err != nil {
            goto ERR
        }
        if err = conn.WriteMessage(websocket.TextMessage, data); err != nil {
            goto ERR
        }
    }
ERR:
    conn.Close()*/

}

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

做了封装，线程安全