[network] Go 长连接服务的三种读写并发模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
type Conn struct {
    ws *websocket.Conn
    mu sync.Mutex // 每连接一把写锁
}

// 任意 goroutine 都能调用，靠锁保证同一连接不并发写
func (c *Conn) Write(data []byte) error {
    c.mu.Lock()
    defer c.mu.Unlock()
    c.ws.SetWriteDeadline(time.Now().Add(5 * time.Second))
    return c.ws.WriteMessage(websocket.TextMessage, data)
}

// 单读协程读 goroutine
func (c *Conn) readLoop() {
    for {
        _, msg, err := c.ws.ReadMessage()
        if err != nil {
            return
        }
        c.handleMessage(msg)
    }
}

func (c *Conn) handleMessage(msg []byte) {
    if string(msg) == "ping" {
        c.Write([]byte("pong"))
    }
}

func (c *Conn) serve() {
    go c.readLoop()

    // 与 readLoop 的 pong 并发写
    for i := 0; i < 3; i++ {
        go func(id int) {
            c.Write([]byte(fmt.Sprintf("push from %d", id)))
        }(i)
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
type Conn struct {
    ws        *websocket.Conn
    send      chan []byte   // 带缓冲发送队列，所有要发的消息先进这里
    done      chan struct{} // 关闭信号：广播"这条连接要收摊了"
    closeOnce sync.Once     // 保证 done 只关一次
}

// 关闭只关 done、永不关 send：任何 goroutine 的 Push 都不会
// "向已关闭 channel 发送"而 panic。sync.Once 保证幂等。
func (c *Conn) close() {
    c.closeOnce.Do(func() { close(c.done) })
}

// 非阻塞投递；满了 = 客户端太慢 -> 踢掉（背压），返回是否投递成功
func (c *Conn) Push(data []byte) bool {
    select {
    case c.send <- data: // 有缓冲位，投递成功
        return true
    case <-c.done: // 已关闭，别再投
        return false
    default: // 队列满 = 慢客户端，踢掉，不拖累别人
        c.close()
        return false
    }
}

// 读协程：只读，把要回的内容丢进队列，不直接写
func (c *Conn) readPump() {
    defer c.close() // 读出错/对端关闭 -> 通知 writePump 退出
    for {
        _, msg, err := c.ws.ReadMessage()
        if err != nil {
            return
        }
        c.handleMessage(msg)
    }
}

// 处理一条业务消息：解析 + 决定回什么。读协程只「投递」、从不直接写 ws，
// 回包统一走 Push 进 send 队列，由唯一的 writePump 串行落地，保持单 writer 无锁。
func (c *Conn) handleMessage(msg []byte) {
    if string(msg) == "ping" {
        c.Push([]byte("pong"))
    }
}

// 写协程：本连接唯一 writer，无需锁；顺带统一管心跳和写超时
func (c *Conn) writePump() {
    tick := time.NewTicker(30 * time.Second)
    defer func() {
        tick.Stop()
        c.ws.Close()
    }()
    for {
        select {
        case data := <-c.send:
            c.ws.SetWriteDeadline(time.Now().Add(5 * time.Second))
            if c.ws.WriteMessage(websocket.TextMessage, data) != nil {
                return
            }
        case <-tick.C:
            if c.ws.WriteMessage(websocket.PingMessage, nil) != nil {
                return
            }
        case <-c.done:
            c.drain() // 退出前把队列残余尽量发完，不丢已投递的数据
            return
        }
    }
}

// 连接关闭后，非阻塞地把 send 队列残余消息发完再退出。
// send 永不被 close，故用 select-default 取到就写、取空即返回。
func (c *Conn) drain() {
    for {
        select {
        case data := <-c.send:
            c.ws.SetWriteDeadline(time.Now().Add(5 * time.Second))
            if c.ws.WriteMessage(websocket.TextMessage, data) != nil {
                return
            }
        default:
            return
        }
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
// 每连接的上下文，挂在 gnet.Conn 上
type ConnContext struct {
    ConnId     string
    ClientID   string
    Addr       string
    StartTime  time.Time
    IsConnOpen bool
}

type Server struct {
    gnet.BuiltinEventEngine
    eng gnet.Engine
}

// 新连接：建上下文 + 注册进连接管理器
func (s *Server) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
    ctx := &ConnContext{
        ConnId:     genUUID(),
        Addr:       c.RemoteAddr().String(),
        StartTime:  time.Now(),
        IsConnOpen: true,
    }
    c.SetContext(ctx)                            // 状态挂到连接上
    clientMgr.Add(ctx.ConnId, c)                 // 注册，后面推送靠它找到 c
    return nil, gnet.None
}

// 可读：拆包 -> 处理 -> 回包，全程不阻塞
func (s *Server) OnTraffic(c gnet.Conn) gnet.Action {
    ctx, ok := c.Context().(*ConnContext)
    if !ok || !ctx.IsConnOpen {
        return gnet.Close
    }

    packet, err := readFullPacket(c)             // 自己粘包拆包，凑不齐一个整包就返回
    if err != nil {
        return gnet.Close
    }
    if packet == nil {
        return gnet.None                         // 半个包，等下次 OnTraffic
    }

    rsp := s.process(ctx, packet)                // 业务处理（注意：别在这里阻塞！）
    c.AsyncWrite(rsp, nil)                        // 异步回包，框架保证并发安全
    return gnet.None
}

// 断开：注销 + 清理，状态从 context 里捞
func (s *Server) OnClose(c gnet.Conn, err error) gnet.Action {
    ctx, ok := c.Context().(*ConnContext)
    if ok && ctx != nil {
        ctx.IsConnOpen = false
        clientMgr.Remove(ctx.ConnId)             // 先注销，别让推送再找到它
        c.SetContext(nil)
    }
    return gnet.Close
}