85 lines
1.8 KiB
Go
85 lines
1.8 KiB
Go
|
package websocket
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"log"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/gorilla/websocket"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type Client struct {
|
|||
|
|
ID string
|
|||
|
|
Conn *websocket.Conn
|
|||
|
|
Send chan []byte
|
|||
|
|
Context context.Context
|
|||
|
|
CancelFunc context.CancelFunc
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *Client) Read(hub *Hub) {
|
|||
|
|
defer func() {
|
|||
|
|
hub.Unregister <- c
|
|||
|
|
c.Conn.Close()
|
|||
|
|
log.Printf("Client %s disconnected\n", c.ID)
|
|||
|
|
c.CancelFunc()
|
|||
|
|
}()
|
|||
|
|
|
|||
|
|
c.Conn.SetReadLimit(512)
|
|||
|
|
c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|||
|
|
c.Conn.SetPongHandler(func(string) error {
|
|||
|
|
c.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|||
|
|
return nil
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
for {
|
|||
|
|
select {
|
|||
|
|
case <-c.Context.Done():
|
|||
|
|
return
|
|||
|
|
default:
|
|||
|
|
_, message, err := c.Conn.ReadMessage()
|
|||
|
|
if err != nil {
|
|||
|
|
log.Printf("Read error from client %s: %v", c.ID, err)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
log.Printf("Receive [%s]: %s", c.ID, message)
|
|||
|
|
// 这里你可以把消息发给 hub.Broadcast 或业务处理
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *Client) Write() {
|
|||
|
|
ticker := time.NewTicker(54 * time.Second) // 小于读超时，保证ping及时发
|
|||
|
|
defer func() {
|
|||
|
|
ticker.Stop()
|
|||
|
|
c.Conn.Close()
|
|||
|
|
c.CancelFunc()
|
|||
|
|
}()
|
|||
|
|
|
|||
|
|
for {
|
|||
|
|
select {
|
|||
|
|
case <-c.Context.Done():
|
|||
|
|
return
|
|||
|
|
case msg, ok := <-c.Send:
|
|||
|
|
if !ok {
|
|||
|
|
// 通道关闭，结束写入
|
|||
|
|
c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
|||
|
|
err := c.Conn.WriteMessage(websocket.TextMessage, msg)
|
|||
|
|
if err != nil {
|
|||
|
|
log.Printf("Write error to client %s: %v", c.ID, err)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
case <-ticker.C:
|
|||
|
|
// 发送 ping
|
|||
|
|
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
|||
|
|
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
|||
|
|
log.Printf("Ping error to client %s: %v", c.ID, err)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|