Golang 简易WS服务 - 客户端

客户端代码较为简单

type SocketClient struct {
    host        string
    isClosed    chan bool
    log         *logger.Log // 自定义的log包
}

type wsMessage struct {
    Type int         `json:"type"`
    Data interface{} `json:"data"`
}

func NewSocketClient(configPath, logPath, host string) *SocketClient {
    w := io.MultiWriter(os.Stdout)
    if logPath != "" {
        f, err := util.OpenFile(logPath)
        if err != nil {
            panic(err)
        }
        w = io.MultiWriter(os.Stdout, f)
    }

    return &SocketClient{
        host:        host,
        isClosed:    make(chan bool),
        log:         logger.New(w, logger.LINFO, log.LstdFlags|log.Lmsgprefix),
    }
}

func (s *SocketClient) Start() {
    // 设置在线
    ClientStatus = ClientStatusOnline

    // 连接服务器
    u := url.URL{Scheme: "ws", Host: s.host, Path: "/socket"}
    s.log.Info("正在连接到", u.String())

    conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        s.log.ErrorF("连接到服务器错误: %v", err)
    }
    defer conn.Close()

    s.log.Info("已连接到服务器")

    interrupt := make(chan os.Signal, 1)

    // 监听中断信号
    signal.Notify(interrupt, syscall.SIGKILL, syscall.SIGINT, syscall.SIGTERM, os.Kill)
    // 心跳包
    go s.pingHandler(conn)
    // 接收消息
    go s.receiveHandler(conn)

    for {
        select {
        case <-interrupt:
            s.log.Info("收到SIGINT中断信号，正在关闭ws连接。。。")
            _ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
            select {
            case <-s.isClosed:
                s.log.Info("WS链接已关闭，退出中。。。")
            case <-time.After(time.Duration(1) * time.Second):
                s.log.Info("关闭ws链接超时，退出中。。。")
            }
            return
        case <-s.isClosed:
            s.log.Info("ws链接已关闭，退出中。。。")
            return
        }
    }
}

func (s *SocketClient) pingHandler(conn *websocket.Conn) {
    ticker := time.NewTicker(pingTime)
    for {
        select {
        case <-ticker.C:
            status := strconv.FormatInt(int64(ClientStatus), 10)
            err := conn.WriteMessage(websocket.PingMessage, []byte(status))
            if err != nil {
                s.log.ErrorF("发送心跳包错误: %v", err)
                return
            }
        case <-s.isClosed:
            return
        }
    }
}

func (s *SocketClient) receiveHandler(ws *websocket.Conn) {
    defer close(s.isClosed)
    for {
        messageType, message, err := ws.ReadMessage()
        if err != nil {
            s.log.ErrorF("读取消息 %v", err)
            return
        }
        switch messageType {
        case websocket.TextMessage:
            var textMessage *wsMessage

            err = json.Unmarshal(message, &textMessage)
            if err != nil {
                s.log.ErrorF("消息格式错误: %v", err)
                break
            }

            switch textMessage.Type {
                    ...
            }
        }
    }
}

作者： Yuming