142 lines
4.5 KiB
Go
142 lines
4.5 KiB
Go
package transport
|
||
|
||
import (
|
||
"fmt"
|
||
"net"
|
||
"sync"
|
||
"syscall"
|
||
|
||
"omnisocketgo/cmd/internal/latencylog"
|
||
"omnisocketgo/cmd/internal/protocol"
|
||
)
|
||
|
||
// UDPConn 是对 UDP 连接的轻量封装。
|
||
// server 侧：共享同一个 net.UDPConn，Send 时通过 peerAddr 指定对端地址。
|
||
// peer 侧：独立的 net.UDPConn，已通过 Dial 连接到 server，Send 直接写即可。
|
||
type UDPConn struct {
|
||
conn *net.UDPConn
|
||
peerAddr *net.UDPAddr // server 侧为对端地址；peer 侧为 nil（连接模式下直接 Write）
|
||
raw syscall.RawConn // 底层 syscall 句柄，用于 Linux socket timestamping
|
||
|
||
logger latencylog.Logger
|
||
txTimestampDebugLogger TXTimestampDebugLogger
|
||
nodeRole string // 日志中记录的节点角色，例如 "server" 或 "peer"
|
||
nodeID string // 日志中记录的节点 ID
|
||
writeMu sync.Mutex // 保护 Send 的互斥锁
|
||
closeOnce sync.Once
|
||
closeErr error
|
||
}
|
||
|
||
// UDPOption 用于为 UDPConn 注入可选行为。
|
||
type UDPOption func(*UDPConn)
|
||
|
||
// WithUDPLogger 为 UDP 连接注入业务消息日志上下文。
|
||
func WithUDPLogger(logger latencylog.Logger, nodeRole, nodeID string) UDPOption {
|
||
return func(conn *UDPConn) {
|
||
conn.logger = logger
|
||
conn.nodeRole = nodeRole
|
||
conn.nodeID = nodeID
|
||
}
|
||
}
|
||
|
||
// WithUDPTXTimestampDebugLogger 为 UDP 连接注入可选的 TX errqueue 调试日志器。
|
||
func WithUDPTXTimestampDebugLogger(logger TXTimestampDebugLogger) UDPOption {
|
||
return func(conn *UDPConn) {
|
||
conn.txTimestampDebugLogger = logger
|
||
}
|
||
}
|
||
|
||
// NewUDPConn 创建 UDP transport 连接封装。
|
||
// peerAddr 为 nil 时表示 peer 侧已连接模式（conn 已 Dial 到 server）。
|
||
// peerAddr 非 nil 时表示 server 侧，Send 时需要指定目标地址。
|
||
func NewUDPConn(conn *net.UDPConn, peerAddr *net.UDPAddr, opts ...UDPOption) (*UDPConn, error) {
|
||
udpConn := &UDPConn{
|
||
conn: conn,
|
||
peerAddr: peerAddr,
|
||
logger: latencylog.NoopLogger{},
|
||
}
|
||
|
||
for _, opt := range opts {
|
||
opt(udpConn)
|
||
}
|
||
|
||
if udpConn.logger == nil {
|
||
udpConn.logger = latencylog.NoopLogger{}
|
||
}
|
||
|
||
if err := udpConn.initUDPLinuxTimestamping(); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return udpConn, nil
|
||
}
|
||
|
||
// Send 将一条协议消息编码为 UDP 数据报并发送。
|
||
// 多个 goroutine 可以并发调用，内部会串行化写入。
|
||
func (c *UDPConn) Send(msg protocol.Message) error {
|
||
c.writeMu.Lock()
|
||
defer c.writeMu.Unlock()
|
||
latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffBegin, msg)
|
||
|
||
if err := c.sendMessageLinux(msg); err != nil {
|
||
return fmt.Errorf("transport: udp send message: %w", err)
|
||
}
|
||
|
||
latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffEnd, msg)
|
||
|
||
return nil
|
||
}
|
||
|
||
// SendTo 将一条协议消息编码为 UDP 数据报并发送到指定地址。
|
||
// 主要用于 server 侧向特定 peer 发送消息。
|
||
func (c *UDPConn) SendTo(msg protocol.Message, addr *net.UDPAddr) error {
|
||
c.writeMu.Lock()
|
||
defer c.writeMu.Unlock()
|
||
latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffBegin, msg)
|
||
|
||
if err := c.sendMessageToLinux(msg, addr); err != nil {
|
||
return fmt.Errorf("transport: udp send message to %s: %w", addr, err)
|
||
}
|
||
|
||
latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffEnd, msg)
|
||
|
||
return nil
|
||
}
|
||
|
||
// Receive 从 UDP 连接读取一条完整协议消息。
|
||
// 返回解码后的消息和来源地址（peer 侧来源地址始终为 server 地址）。
|
||
func (c *UDPConn) Receive() (protocol.Message, *net.UDPAddr, error) {
|
||
msg, addr, err := c.receiveMessageLinux()
|
||
if err != nil {
|
||
return protocol.Message{}, nil, fmt.Errorf("transport: udp receive message: %w", err)
|
||
}
|
||
|
||
return msg, addr, nil
|
||
}
|
||
|
||
// ReceiveLoop 持续从 UDP 连接读取消息并交给 handler 处理。
|
||
// handler 的第二个参数是消息来源地址。
|
||
func (c *UDPConn) ReceiveLoop(handler func(protocol.Message, *net.UDPAddr) error) error {
|
||
for {
|
||
msg, addr, err := c.Receive()
|
||
if err != nil {
|
||
return fmt.Errorf("transport: udp receive loop read: %w", err)
|
||
}
|
||
|
||
if err := handler(msg, addr); err != nil {
|
||
return fmt.Errorf("transport: udp receive loop handler: %w", err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Close 关闭底层 UDP 连接，保证重复调用安全。
|
||
// 注意：server 侧多个 UDPConn 共享同一个 net.UDPConn 时，
|
||
// 只应由 UDPHub 负责关闭底层连接，不应通过此方法关闭。
|
||
func (c *UDPConn) Close() error {
|
||
c.closeOnce.Do(func() {
|
||
c.closeErr = c.conn.Close()
|
||
})
|
||
|
||
return c.closeErr
|
||
}
|