Websocket

ws是基于 gorilla/websocket 封装的 websocket 库，支持自动客户端重新连接。

默认设置示例

服务端示例代码

package main

import (
    "context"
    "log"
    "github.com/go-dev-frame/sponge/pkg/ws"
    "github.com/gin-gonic/gin"
)

func main() {
    r := gin.Default()
    
    r.GET("/ws", func(c *gin.Context) {
        s := ws.NewServer(c.Writer, c.Request, loopReceiveMessage) // 默认设置
        err := s.Run(context.Background())
        if err != nil {
            log.Println("webSocket server error:", err)
        }
    })
    
    err := r.Run(":8080")
    if err != nil {
        panic(err)
    }
}

func loopReceiveMessage(ctx context.Context, conn *ws.Conn) {
    for {
        messageType, message, err := conn.ReadMessage()
        if err != nil {
            return
        }

        // 处理消息
        log.Println(messageType, message)
    }
}

客户端示例代码

package main

import (
    "strconv"
    "log"
    "time"
    "github.com/go-dev-frame/sponge/pkg/ws"
    "github.com/gorilla/websocket"
)

var wsURL = "ws://localhost:8080/ws"

func main() {
    c, err := ws.NewClient(wsURL) // 默认设置
    if err != nil {
        log.Println("connect error:", err)
        return
    }
    defer c.Close()

    go func() {
        for {
            _, message, err := c.GetConn().ReadMessage()
            if err != nil {
                log.Println("client read error:", err)
                return
            }
            log.Printf("client received: %s", message)
        }
    }()
    
    for i := 0; i < 5; i++ {
        data := "Hello, World " + strconv.Itoa(i)
        err = c.GetConn().WriteMessage(websocket.TextMessage, []byte(data))
        if err != nil {
            log.Println("write error:", err)
        }
        time.Sleep(100 * time.Millisecond)
    }
}

自定义设置示例

服务器端自定义设置，可设置ws.Upgrader、ws. WithNoClientPingData、ws.WithServerLogger、ws.WithResponseHeader等选项。

服务端示例代码

package main

import (
    "context"
    "net/http"
    "time"

    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
    "github.com/go-dev-frame/sponge/pkg/logger"
    "github.com/go-dev-frame/sponge/pkg/ws"
)

func main() {
    r := gin.Default()
    ug := &websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
    r.GET("/ws", func(c *gin.Context) {
        s := ws.NewServer(c.Writer, c.Request, loopReceiveMessage,
            ws.WithUpgrader(ug),
            ws.WithNoClientPingTimeout(time.Minute), // ping 消息超时时间，超时会主动断开客户端连接
            ws.WithServerLogger(logger.Get()),
        )
        err := s.Run(context.Background())
        if err != nil {
            logger.Warn("WebSocket server error:", logger.Err(err))
        }
    })

    err := r.Run(":8080")
    if err != nil {
        panic(err)
    }
}

func loopReceiveMessage(ctx context.Context, conn *ws.Conn) {
    for {
        messageType, message, err := conn.ReadMessage()
        if err != nil {
            logger.Warn("ReadMessage error", logger.Err(err))
            return
        }
        logger.Infof("server side received: %s", message)

        switch messageType {
        case websocket.TextMessage:
            err = conn.WriteMessage(messageType, message)
            if err != nil {
                logger.Warn("WriteMessage error", logger.Err(err))
                continue
            }

        case websocket.PingMessage:
            err = conn.WriteMessage(websocket.PongMessage, []byte("pong"))
            if err != nil {
                logger.Warn("Write pong message error:", logger.Err(err))
                continue
            }
        default:
            logger.Warnf("Unknown message type: %d", messageType)
        }
    }
}

客户端自定义设置，可设置ws.Dialer、ws.WithPing、ws.WithClientLogger、ws.WithRequestHeader等选项。

客户端示例代码

package main

import (
    "strconv"
    "time"

    "github.com/gorilla/websocket"
    "github.com/go-dev-frame/sponge/pkg/logger"
    "github.com/go-dev-frame/sponge/pkg/ws"
)

var wsURL = "ws://localhost:8080/ws"

func main() {
    c, err := ws.NewClient(wsURL,
        ws.WithPing(time.Second*20), // 设置的 ping 消息发送间隔
        ws.WithClientLogger(logger.Get()),
    )
    if err != nil {
        logger.Warn("connect error", logger.Err(err))
        return
    }
    defer c.Close()

    go clientLoopReadMessage(c)

    i := 0
    for {
        time.Sleep(time.Second * 3)
        i++
        data := "Hello, World " + strconv.Itoa(i)
        err = c.GetConn().WriteMessage(websocket.TextMessage, []byte(data))
        if err != nil {
            logger.Warn("WriteMessage error", logger.Err(err))
        }
    }
}

func clientLoopReadMessage(c *ws.Client) {
    for {
        select {
        case <-c.GetCtx().Done():
            return
        default:
            _, message, err := c.GetConn().ReadMessage()
            if err != nil {
                logger.Warn("ReadMessage error", logger.Err(err))
                time.Sleep(time.Second * 5)
                continue
            }
            logger.Infof("client side received: %s", message)
        }

    }
}