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
|
||
}
|
||
}
|
||
}
|
||
}
|