zinx 增加 websocket 支持,实现socket,websocket 任意切换

501 阅读2分钟

要干啥

前段时间在研究go 语言的socket 框架,找到zinx ,看了下,是比较基础的框架,逻辑清晰,感觉还可以,最近在做websocket 相关的,看了下zinx 只能支持socket ,不支持websocket ,想着改改让zinx 两个都支持。

干活

源码下载下来,先分析下,项目结构

├─utils
│      globalobj.go
│
├─ziface
│      iconnection.go
│      iconnmanager.go
│      idatapack.go
│      imessage.go
│      imsghandler.go
│      ipacket.go
│      irequest.go
│      irouter.go
│      iserver.go
│
├─zlog
│  .....
├─znet
│      connection.go
│      connection_ws.go
│      connmanager.go
│      datapack.go
│      datapack_test.go
│      message.go
│      msghandler.go
│      options.go
│      request.go
│      router.go
│      server.go
│      server_test.go
│
└─ztimer
   .....

工程把消息路由,链接管理,消息收发,等相关的功能抽离出来,做成接口,然后在做实现类进行具体的实现,项目仅实现了socket相关的,所以接下来稍微改下就能支持websocket ,并通过配置文件控制

第一步 Ctrl+C+V

1,IConnection.go 里面有个方法 GetTCPConnection ,增加一个 GetWsConnection,可能用得着,找了下代码里面并没有用到这两个方法

GetTCPConnection() *net.TCPConn // 
GetWsConnection() *websocket.Conn // 

2, 把原有的zent/connection.go 复制衣分,重命名,为connection_ws.go,修改如下内容

// 1,把 Connection 修改为 ConnectionWs
type ConnectionWs struct {
    ....
    // 修改这里,把原来的 *net.TCPConn 改成 *websocket.Conn
    Conn *websocket.Conn
    ....
}
// 2,把所有方法前面的 *Connection 改成 *ConnectionWs

3,修改读写的方法 StartWriter , StartReader

//StartWriter 写消息Goroutine, 用户将数据发送给客户端
func (c *ConnectionWs) StartWriter() {
    fmt.Println("[Writer Goroutine is running]")
    defer fmt.Println(c.RemoteAddr().String(), "[conn Writer exit!]")
​
    for {
        select {
        case data := <-c.msgChan:
            // 有数据要写给客户端
            // 修改这里 把原来的 c.Conn.Write(data) 改成如下方式
            if err := c.Conn.WriteMessage(websocket.BinaryMessage,data); err != nil {
                fmt.Println("Send Data error:, ", err, " Conn Writer exit")
                return
            }
        case data, ok := <-c.msgBuffChan:
            if ok {
                // 有数据要写给客户端
                // 修改这里 把原来的 c.Conn.Write(data) 改成如下方式
                if err := c.Conn.WriteMessage(websocket.BinaryMessage,data); err != nil {
                    fmt.Println("Send Buff Data error:, ", err, " Conn Writer exit")
                    return
                }
            } else {
                fmt.Println("msgBuffChan is Closed")
                break
            }
        case <-c.ctx.Done():
            return
        }
    }
}
​
//StartReader 读消息Goroutine,用于从客户端中读取数据
func (c *ConnectionWs) StartReader() {
    fmt.Println("[Reader Goroutine is running]")
    defer fmt.Println(c.RemoteAddr().String(), "[conn Reader exit!]")
    defer c.Stop()
    
    for {
        select {
        case <-c.ctx.Done():
            return
        default:
            headData := make([]byte, c.TCPServer.Packet().GetHeadLen())
            // 写的消息的方式按照如下方式修改 ,增加这段代码
            _,reader,err:=c.Conn.NextReader()  
            if err!=nil{
                fmt.Println("read msg error ", err)
                return
            }
            
            // 修改这里 原来的 io.ReadFull(conn , headData) 改成如下
            if _, err := io.ReadFull(reader, headData); err != nil {
                fmt.Println("read msg head error ", err)
                return
            }
​
            msg, err := c.TCPServer.Packet().Unpack(headData)
            if err != nil {
                fmt.Println("unpack error ", err)
                return
            }
​
            var data []byte
            if msg.GetDataLen() > 0 {
                data = make([]byte, msg.GetDataLen())
                // 修改这里 原来的 io.ReadFull(conn , headData) 改成如下
                if _, err := io.ReadFull(reader, data); err != nil {
                    fmt.Println("read msg data error ", err)
                    return
                }
            }
            msg.SetData(data)
            req := Request{
                conn: c,
                msg:  msg,
            }
​
            if utils.GlobalObject.WorkerPoolSize > 0 {
                c.MsgHandler.SendMsgToTaskQueue(&req)
            } else {
                go c.MsgHandler.DoMsgHandler(&req)
            }
        }
    }
}

4,IServer 接口增加

//定义服务接口
type IServer interface {
    // .. 增加如下接口
    StartWs()       //启动服务器方法
    // ....
}

5 Server.go 里面实现StartWs()

//Start 开启网络服务
func (s *Server) StartWs() {
    fmt.Printf("[START] WsServer name: %s,listenner at IP: %s, Port %d is starting\n", s.Name, s.IP, s.Port)
​
    go func() {
        //0 启动worker工作池机制
        s.msgHandler.StartWorkerPool()
​
        //1 获取一个TCP的Addr
        addr, err := net.ResolveTCPAddr(s.IPVersion, fmt.Sprintf("%s:%d", s.IP, s.Port))
        if err != nil {
            fmt.Println("resolve tcp addr err: ", err)
            return
        }
​
        http.HandleFunc("/", s.wsHandler)
        //2 监听服务器地址
        fmt.Println("ws 启动地址 ",addr.String())
        cID = 0
        err = http.ListenAndServe(addr.String(), nil)
        if err != nil {
            panic(err)
        }
    }()
​
}
​
// ws 处理方法
func (s *Server)wsHandler(w http.ResponseWriter, r *http.Request)  {
    conn, err := upGrader.Upgrade(w, r, nil)
    if err != nil {
        log.Print("upgrade:", err)
        return
    }
    //defer conn.Close()
    fmt.Println("Get conn remote addr = ", conn.RemoteAddr().String())
​
    // 有一个id 生成的方法
​
    cID = 0
    //3.2 设置服务器最大连接控制,如果超过最大连接,那么则关闭此新的连接
    if s.ConnMgr.Len() >= utils.GlobalObject.MaxConn {
        conn.Close()
        return
    }
​
    //3.3 处理该新连接请求的 业务 方法, 此时应该有 handler 和 绑conn是定的
    dealConn := NewConnectionWs(s, conn, cID, s.msgHandler)
    cID++
​
    //3.4 启动当前链接的处理业务
    go dealConn.Start()
}

6, Server 结构体中增加 字段 IsWs ,在 Serve 方法中做判断用

type Server struct {
    // 
    IsWs bool
}
​
//Serve 运行服务
func (s *Server) Serve() {
    if s.IsWs{
        s.StartWs()
    }else{
        s.Start()
    }
    //TODO Server.Serve() 是否在启动服务的时候 还要处理其他的事情呢 可以在这里添加
    //阻塞,否则主Go退出, listenner的go将会退出
    select {}
}

最后在使用的时候只需要配置上 IsWs 属性就可以使用 websocket了

第二步 结尾

如此,不需要修改其他的协议逻辑,就可以实现在 socket 与 websocket , 之间任意切换