137 lines
3.0 KiB
Go
137 lines
3.0 KiB
Go
package realtime
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
)
|
|
|
|
const sendBufferSize = 64
|
|
|
|
// client wraps a single WebSocket connection belonging to a user.
|
|
type client struct {
|
|
hub *Hub
|
|
conn *websocket.Conn
|
|
userID int64
|
|
send chan []byte
|
|
done chan struct{}
|
|
}
|
|
|
|
func newClient(hub *Hub, conn *websocket.Conn, userID int64) *client {
|
|
return &client{
|
|
hub: hub,
|
|
conn: conn,
|
|
userID: userID,
|
|
send: make(chan []byte, sendBufferSize),
|
|
done: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// enqueue marshals env and queues it for delivery; drops the message if the
|
|
// client cannot keep up.
|
|
func (c *client) enqueue(env envelope) {
|
|
payload, err := json.Marshal(env)
|
|
if err != nil {
|
|
return
|
|
}
|
|
c.enqueueRaw(payload)
|
|
}
|
|
|
|
func (c *client) enqueueRaw(payload []byte) {
|
|
select {
|
|
case <-c.done:
|
|
case c.send <- payload:
|
|
default:
|
|
c.closeSlow()
|
|
}
|
|
}
|
|
|
|
// closeSlow terminates a connection that cannot keep up with its send queue.
|
|
func (c *client) closeSlow() {
|
|
c.hub.mu.Lock()
|
|
defer c.hub.mu.Unlock()
|
|
if conns := c.hub.conns[c.userID]; conns != nil {
|
|
if _, ok := conns[c]; ok {
|
|
c.conn.Close(websocket.StatusPolicyViolation, "client too slow")
|
|
}
|
|
}
|
|
}
|
|
|
|
// writeLoop drains the send queue onto the connection.
|
|
func (c *client) writeLoop() {
|
|
ctx := context.Background()
|
|
for {
|
|
select {
|
|
case <-c.done:
|
|
return
|
|
case msg := <-c.send:
|
|
writeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
|
err := c.conn.Write(writeCtx, websocket.MessageText, msg)
|
|
cancel()
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// readLoop processes inbound client events until the connection closes.
|
|
func (c *client) readLoop() error {
|
|
ctx := context.Background()
|
|
for {
|
|
_, data, err := c.conn.Read(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var env envelope
|
|
if err := json.Unmarshal(data, &env); err != nil {
|
|
c.conn.Close(websocket.StatusInvalidFramePayloadData, "bad json")
|
|
return err
|
|
}
|
|
|
|
switch env.Type {
|
|
case EventTyping:
|
|
var msg clientTyping
|
|
if err := json.Unmarshal(env.Data, &msg); err != nil {
|
|
continue
|
|
}
|
|
if err := c.handleTyping(ctx, msg); err != nil {
|
|
return err
|
|
}
|
|
case EventRead:
|
|
var msg clientRead
|
|
if err := json.Unmarshal(env.Data, &msg); err != nil {
|
|
continue
|
|
}
|
|
if err := c.handleRead(ctx, msg); err != nil {
|
|
return err
|
|
}
|
|
default:
|
|
return websocket.CloseError{Code: websocket.StatusUnsupportedData, Reason: "unsupported event type"}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *client) handleTyping(ctx context.Context, msg clientTyping) error {
|
|
ok, err := c.hub.members.IsMember(ctx, msg.ConversationID, c.userID)
|
|
if err != nil || !ok {
|
|
return nil
|
|
}
|
|
c.hub.PublishTyping(msg.ConversationID, c.userID, msg.IsTyping)
|
|
return nil
|
|
}
|
|
|
|
func (c *client) handleRead(ctx context.Context, msg clientRead) error {
|
|
ok, err := c.hub.members.IsMember(ctx, msg.ConversationID, c.userID)
|
|
if err != nil || !ok {
|
|
return nil
|
|
}
|
|
if err := c.hub.members.UpdateLastRead(ctx, msg.ConversationID, c.userID); err != nil {
|
|
return err
|
|
}
|
|
c.hub.PublishRead(msg.ConversationID, c.userID, msg.MessageID)
|
|
return nil
|
|
}
|