feat(gomsg): pymsg/libmsg 协议的 Go 实现，waf logclient 抽取接入
- 新增 gomsg/：帧协议（4B 大端长度 + UTF-8，64MiB 上限）、MsgClient/MsgServer、 RpcClient/RpcServer（_route 分发、错误脱敏、panic 恢复、json.Number 保精度）、 AsyncQueue[T]、LogClient（集中式日志，级别过滤 + fail-safe），线协议与 Python/C++ 字节级兼容 - waf/internal/logclient 改为 gomsg 薄适配层（336→20 行），main/engine/proxy 零改动 - gomsg 26 项单测全绿（-race）；waf 编译通过
This commit is contained in:
qing committed
1 parent
0ed03ac095
commit
c4154d4db9
16 files changed
+1624
-317
No files matched your search
@@ -0,0 +1,73 @@
|
||||
# gomsg — pymsg / libmsg 协议的 Go 实现
|
||||
|
||||
Go 版消息通信库，对应 Python [pymsg](../pymsg/) 与 C++ [libmsg](../libmsg/)，线协议字节级兼容：
|
||||
|
||||
```
|
||||
[4 字节 big-endian uint32 长度][UTF-8 载荷] （单条上限 64 MiB）
|
||||
```
|
||||
|
||||
## 模块组成
|
||||
|
||||
| 文件 | 内容 | 对应 |
|
||||
|------|------|------|
|
||||
| `frame.go` | 帧编解码：`Encode` / `WriteFrame` / `ReadFrame`（长度上限防护） | pymsg `send`/`recv` |
|
||||
| `client.go` | `MsgClient`：复用长连接 + 请求-响应 + 失败重连一次 | pymsg `MsgClient` |
|
||||
| `server.go` | `MsgServer`：accept 循环 + 每连接 goroutine 串行处理 | pymsg `MsgServer` |
|
||||
| `rpc.go` | `RpcClient` / `RpcServer` / `RpcError`：`_route` 分发、错误脱敏、panic 恢复 | pymsg `rpc.py` |
|
||||
| `asyncqueue.go` | `AsyncQueue[T]`：批量 + 定时 flush，满队列丢新数据（fail-safe） | libmsg `AsyncQueue` |
|
||||
| `logclient.go` | `LogClient`：异步批量推送 db2 Logs 表（级别过滤 + fail-safe） | waf `internal/logclient`、`log_client.py`、liblogc |
|
||||
|
||||
## 使用
|
||||
|
||||
```go
|
||||
// RPC 调用（如调 db2 execute 路由）
|
||||
rpc := gomsg.NewRpcClient("127.0.0.1:8005", 5*time.Second)
|
||||
data, err := rpc.Call("execute", map[string]any{
|
||||
"action": "select", "table": "Users",
|
||||
})
|
||||
// data 为 any：数字是 json.Number（保 int64 精度），需用 .Int64()/.Float64() 转换
|
||||
|
||||
// 提供 RPC 服务
|
||||
srv := gomsg.NewRpcServer("127.0.0.1:8009")
|
||||
srv.AddRoute("health", func(req map[string]any) (any, error) { return map[string]any{"ok": true}, nil })
|
||||
if err := srv.Start(); err != nil { ... }
|
||||
defer srv.Stop()
|
||||
|
||||
// 集中式日志（写入 db2 Logs 表）
|
||||
lc := gomsg.NewLogClient("127.0.0.1:8005")
|
||||
lc.SetLevel("info")
|
||||
lc.Info("backend", "服务启动", nil)
|
||||
defer lc.Flush() // 进程退出前 flush
|
||||
```
|
||||
|
||||
## 与 Python / C++ 的兼容约定
|
||||
|
||||
- **请求**：`{"_route": "<路由名>", ...业务字段}`（`_route` 在分发前被弹出，handler 看不到）
|
||||
- **响应**：`{"success": true, "data": ...}` / `{"success": false, "error": "..."}`
|
||||
- **未知路由**：`{"success": false, "error": "未知路由: <route>"}`
|
||||
- **非法 JSON**：`{"success": false, "error": "invalid JSON request"}`（不回显输入片段）
|
||||
- **错误脱敏**：异常信息中的绝对路径替换为 `<redacted>`（对齐 Python `_sanitize_error`）
|
||||
- **handler panic**：自动恢复并包装为 `success=false`（对应 Python 的 `except Exception`）
|
||||
|
||||
## 安全性
|
||||
|
||||
- 帧长度上限 64 MiB 在读写两侧强制（恶意长度前缀不会导致内存耗尽）
|
||||
- 请求 JSON 解析：Go `encoding/json` 内置 10000 层嵌套上限；数字经 `UseNumber` 以字符串保存，
|
||||
超大数字不会触发 CPU 密集型解析；输入大小受帧上限约束
|
||||
- 异常信息脱敏，不透传绝对路径等内部细节
|
||||
- 数字一律 `json.Number` 保留原始精度，避免 int64 ID 经 float64 丢失
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
go test ./... -count=1 -race
|
||||
```
|
||||
|
||||
覆盖：帧往返/超限/截断、Client-Server 回环、自动重连、并发客户端、RPC 成功/失败/脱敏/
|
||||
未知路由/非法 JSON/panic 恢复/超时、AsyncQueue 批量/定时/关停 drain/满队列丢弃、
|
||||
LogClient 对假 db2 服务的推送与级别过滤、db2 不可用时 fail-safe。
|
||||
|
||||
## 来源
|
||||
|
||||
`MsgClient` 连接复用与 `LogClient` 逻辑抽取自 [waf/internal/logclient](../waf/internal/logclient/)；
|
||||
API 与协议约定对齐 [pymsg](../pymsg/)（Python）与 [libmsg](../libmsg/)（C++）。
|
||||
@@ -0,0 +1,150 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AsyncQueueConfig 异步批量队列配置（语义对齐 C++ libmsg.AsyncQueueConfig）。
|
||||
type AsyncQueueConfig struct {
|
||||
BatchSize int // 批量发送阈值，达到即 flush；必须 > 0，默认 50
|
||||
FlushInterval time.Duration // 定时 flush 间隔；必须 > 0 且 <= 300s，默认 2s
|
||||
MaxQueueSize int // 队列容量；必须 >= 1，默认 10000；满时丢新数据（对齐 waf logclient 的 fail-safe 语义）
|
||||
}
|
||||
|
||||
func (c AsyncQueueConfig) withDefaults() AsyncQueueConfig {
|
||||
if c.BatchSize <= 0 {
|
||||
c.BatchSize = 50
|
||||
}
|
||||
if c.FlushInterval <= 0 || c.FlushInterval > 300*time.Second {
|
||||
c.FlushInterval = 2 * time.Second
|
||||
}
|
||||
if c.MaxQueueSize < 1 {
|
||||
c.MaxQueueSize = 10000
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// AsyncQueue 通用异步批量发送队列（对应 C++ libmsg.AsyncQueue）。
|
||||
// Push 非阻塞；队列满时静默丢弃并计数（fail-safe，不影响业务）。
|
||||
type AsyncQueue[T any] struct {
|
||||
cfg AsyncQueueConfig
|
||||
handler func([]T)
|
||||
ch chan T
|
||||
stop chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
mu sync.Mutex
|
||||
stopped bool
|
||||
dropped uint64
|
||||
flushed uint64
|
||||
}
|
||||
|
||||
// NewAsyncQueue 创建异步批量队列，handler 在单后台 goroutine 中按批调用。
|
||||
func NewAsyncQueue[T any](handler func([]T), cfg AsyncQueueConfig) *AsyncQueue[T] {
|
||||
if handler == nil {
|
||||
panic("gomsg: AsyncQueue handler must not be nil")
|
||||
}
|
||||
cfg = cfg.withDefaults()
|
||||
q := &AsyncQueue[T]{
|
||||
cfg: cfg,
|
||||
handler: handler,
|
||||
ch: make(chan T, cfg.MaxQueueSize),
|
||||
stop: make(chan struct{}),
|
||||
}
|
||||
q.wg.Add(1)
|
||||
go q.loop()
|
||||
return q
|
||||
}
|
||||
|
||||
// Push 入队一条数据；队列满时返回 false（该条被丢弃）。
|
||||
func (q *AsyncQueue[T]) Push(item T) bool {
|
||||
q.mu.Lock()
|
||||
if q.stopped {
|
||||
q.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
q.mu.Unlock()
|
||||
select {
|
||||
case q.ch <- item:
|
||||
return true
|
||||
default:
|
||||
q.mu.Lock()
|
||||
q.dropped++
|
||||
q.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown 停止队列，flush 剩余数据后退出后台 goroutine。
|
||||
func (q *AsyncQueue[T]) Shutdown() {
|
||||
q.mu.Lock()
|
||||
if q.stopped {
|
||||
q.mu.Unlock()
|
||||
return
|
||||
}
|
||||
q.stopped = true
|
||||
q.mu.Unlock()
|
||||
close(q.stop)
|
||||
q.wg.Wait()
|
||||
}
|
||||
|
||||
// Stats 返回丢弃与已发送批次数（供监控/测试）。
|
||||
func (q *AsyncQueue[T]) Stats() (dropped, flushed uint64) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
return q.dropped, q.flushed
|
||||
}
|
||||
|
||||
func (q *AsyncQueue[T]) loop() {
|
||||
defer q.wg.Done()
|
||||
batch := make([]T, 0, q.cfg.BatchSize)
|
||||
ticker := time.NewTicker(q.cfg.FlushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
// 非阻塞续取，直到凑满 batchSize 或暂时无数据
|
||||
for len(batch) < q.cfg.BatchSize {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
default:
|
||||
goto send
|
||||
}
|
||||
}
|
||||
send:
|
||||
if len(batch) >= q.cfg.BatchSize {
|
||||
q.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
q.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-q.stop:
|
||||
// 退出前 drain 剩余数据
|
||||
for {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
default:
|
||||
if len(batch) > 0 {
|
||||
q.send(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *AsyncQueue[T]) send(batch []T) {
|
||||
q.handler(batch)
|
||||
q.mu.Lock()
|
||||
q.flushed++
|
||||
q.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func collectBatches(batches *[][]int, mu *sync.Mutex) func([]int) {
|
||||
return func(b []int) {
|
||||
mu.Lock()
|
||||
*batches = append(*batches, append([]int(nil), b...))
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func TestAsyncQueueFlushOnBatchSize(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 5,
|
||||
FlushInterval: time.Hour, // 只测批量阈值，不测定时
|
||||
})
|
||||
for i := 1; i <= 5; i++ {
|
||||
if !q.Push(i) {
|
||||
t.Fatal("push should succeed")
|
||||
}
|
||||
}
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
mu.Lock()
|
||||
n := len(batches)
|
||||
mu.Unlock()
|
||||
if n >= 1 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("batch not flushed in time")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(batches[0]) != 5 || batches[0][4] != 5 {
|
||||
t.Fatalf("batch = %v", batches[0])
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueFlushOnInterval(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 100,
|
||||
FlushInterval: 100 * time.Millisecond,
|
||||
})
|
||||
if !q.Push(42) {
|
||||
t.Fatal("push should succeed")
|
||||
}
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
mu.Lock()
|
||||
n := len(batches)
|
||||
mu.Unlock()
|
||||
if n < 1 {
|
||||
t.Fatal("expected interval flush")
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueShutdownDrains(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 100,
|
||||
FlushInterval: time.Hour,
|
||||
})
|
||||
for i := 1; i <= 3; i++ {
|
||||
q.Push(i)
|
||||
}
|
||||
q.Shutdown()
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(batches) != 1 || len(batches[0]) != 3 {
|
||||
t.Fatalf("expected drained batch [1 2 3], got %v", batches)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAsyncQueueDropOnFull(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 1000,
|
||||
FlushInterval: time.Hour,
|
||||
MaxQueueSize: 4,
|
||||
})
|
||||
// 后台协程可能已取走部分，填满队列直到 Push 返回 false
|
||||
dropped := false
|
||||
for i := 0; i < 100; i++ {
|
||||
if !q.Push(i) {
|
||||
dropped = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !dropped {
|
||||
t.Fatal("expected drop when queue full")
|
||||
}
|
||||
d, _ := q.Stats()
|
||||
if d == 0 {
|
||||
t.Fatal("dropped counter should be > 0")
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueRejectAfterShutdown(t *testing.T) {
|
||||
q := NewAsyncQueue[int](func([]int) {}, AsyncQueueConfig{})
|
||||
q.Shutdown()
|
||||
if q.Push(1) {
|
||||
t.Fatal("push after shutdown should fail")
|
||||
}
|
||||
}
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DefaultTimeout 单次请求-响应交互的总超时。
|
||||
const DefaultTimeout = 5 * time.Second
|
||||
|
||||
// MsgClient pymsg 协议客户端：复用长连接 + 请求-响应调用。
|
||||
// 连接断开时自动重建一次（对应 Python pymsg.MsgClient / C++ libmsg.MsgClient）。
|
||||
type MsgClient struct {
|
||||
addr string
|
||||
timeout time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
conn net.Conn
|
||||
}
|
||||
|
||||
// NewMsgClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。
|
||||
func NewMsgClient(addr string, timeout time.Duration) *MsgClient {
|
||||
if timeout <= 0 {
|
||||
timeout = DefaultTimeout
|
||||
}
|
||||
return &MsgClient{addr: addr, timeout: timeout}
|
||||
}
|
||||
|
||||
// Call 发送一条请求并读取响应（复用长连接，失败时关闭重连一次）。
|
||||
// 返回响应载荷；超时 / 连接错误返回 error。
|
||||
func (c *MsgClient) Call(payload []byte) ([]byte, error) {
|
||||
resp, err := c.call(payload)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
// 旧连接可能已断开（服务端重启等），关闭后重连一次
|
||||
c.closeConn()
|
||||
return c.call(payload)
|
||||
}
|
||||
|
||||
// Send 发送一条请求，不等待响应（fire-and-forget）。失败时重连一次。
|
||||
func (c *MsgClient) Send(payload []byte) error {
|
||||
if err := c.sendOnly(payload); err != nil {
|
||||
c.closeConn()
|
||||
return c.sendOnly(payload)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close 关闭底层连接。
|
||||
func (c *MsgClient) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.closeLocked()
|
||||
}
|
||||
|
||||
func (c *MsgClient) call(payload []byte) ([]byte, error) {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = conn.SetDeadline(time.Now().Add(c.timeout))
|
||||
if err := WriteFrame(conn, payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ReadFrame(conn, 0)
|
||||
}
|
||||
|
||||
func (c *MsgClient) sendOnly(payload []byte) error {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(c.timeout))
|
||||
return WriteFrame(conn, payload)
|
||||
}
|
||||
|
||||
func (c *MsgClient) getConn() (net.Conn, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.conn != nil {
|
||||
return c.conn, nil
|
||||
}
|
||||
conn, err := net.DialTimeout("tcp", c.addr, c.timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.conn = conn
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (c *MsgClient) closeConn() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
_ = c.closeLocked()
|
||||
}
|
||||
|
||||
func (c *MsgClient) closeLocked() error {
|
||||
if c.conn != nil {
|
||||
err := c.conn.Close()
|
||||
c.conn = nil
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// startEchoServer 启动一个 MsgServer，把收到的载荷原样回复。
|
||||
func startEchoServer(t *testing.T) (*MsgServer, string) {
|
||||
t.Helper()
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
return append([]byte("echo:"), payload...)
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(s.Stop)
|
||||
return s, s.Addr().String()
|
||||
}
|
||||
|
||||
func TestMsgClientCall(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
resp, err := c.Call([]byte(`{"hello":"世界"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "echo:{\"hello\":\"世界\"}"
|
||||
if string(resp) != want {
|
||||
t.Fatalf("resp = %q, want %q", resp, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientReusesConn(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := c.Call([]byte("x")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if c.conn == nil {
|
||||
t.Fatal("expected connection reuse")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientReconnectAfterServerRestart(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(p []byte) []byte { return p })
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
addr := s.Addr().String()
|
||||
|
||||
c := NewMsgClient(addr, 2*time.Second)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("before")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 服务端重启：旧连接失效，Call 应自动重连成功
|
||||
s.Stop()
|
||||
// 重新监听同一地址：旧 listener 已释放端口
|
||||
s3 := NewMsgServer(addr, func(p []byte) []byte { return p })
|
||||
if err := s3.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s3.Stop()
|
||||
|
||||
resp, err := c.Call([]byte("after"))
|
||||
if err != nil {
|
||||
t.Fatalf("expected auto-reconnect, got %v", err)
|
||||
}
|
||||
if string(resp) != "after" {
|
||||
t.Fatalf("resp = %q", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientSendNoReply(t *testing.T) {
|
||||
got := make(chan []byte, 1)
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
got <- payload
|
||||
return nil // 不回复
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
if err := c.Send([]byte("fire-and-forget")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case p := <-got:
|
||||
if string(p) != "fire-and-forget" {
|
||||
t.Fatalf("payload = %q", p)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("server did not receive message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientTimeout(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
time.Sleep(500 * time.Millisecond) // 拖过客户端超时
|
||||
return nil
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 50*time.Millisecond)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("slow")); err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgServerStopClosesConns(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("ping")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgServerConcurrentClients(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
return bytes.ToUpper(payload)
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
errCh := make(chan error, 8)
|
||||
for i := 0; i < 8; i++ {
|
||||
go func(n int) {
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
for j := 0; j < 10; j++ {
|
||||
in := []byte{byte('a' + n)}
|
||||
resp, err := c.Call(in)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
if len(resp) != 1 || resp[0] != in[0]-'a'+'A' {
|
||||
errCh <- errors.New("unexpected reply")
|
||||
return
|
||||
}
|
||||
}
|
||||
errCh <- nil
|
||||
}(i)
|
||||
}
|
||||
for i := 0; i < 8; i++ {
|
||||
if err := <-errCh; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
// Package gomsg 是 pymsg（Python）/ libmsg（C++）协议的 Go 实现。
|
||||
//
|
||||
// 线协议（与 Python / C++ 字节级兼容）：
|
||||
//
|
||||
// [4 字节 big-endian uint32 长度][UTF-8 载荷]
|
||||
//
|
||||
// 消息大小上限 64 MiB（与 Python pymsg / C++ libmsg 对齐）。
|
||||
//
|
||||
// 覆盖功能：帧编解码（本文件）、MsgClient/MsgServer（client.go / server.go）、
|
||||
// JSON RPC 约定层（rpc.go）、通用异步批量队列（asyncqueue.go）与
|
||||
// 集中式日志客户端（logclient.go，对应各服务目录下的 log_client.py / liblogc）。
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
// MaxMsgSize 单条消息载荷大小上限，与 Python pymsg / C++ libmsg 保持一致。
|
||||
const MaxMsgSize = 64 * 1024 * 1024
|
||||
|
||||
var headerSize = 4
|
||||
|
||||
// Encode 将载荷封装为 pymsg 帧：4 字节大端长度前缀 + 载荷。
|
||||
func Encode(payload []byte) ([]byte, error) {
|
||||
if len(payload) > MaxMsgSize {
|
||||
return nil, fmt.Errorf("payload too large: %d bytes (max %d)", len(payload), MaxMsgSize)
|
||||
}
|
||||
frame := make([]byte, headerSize+len(payload))
|
||||
binary.BigEndian.PutUint32(frame, uint32(len(payload)))
|
||||
copy(frame[headerSize:], payload)
|
||||
return frame, nil
|
||||
}
|
||||
|
||||
// WriteFrame 将载荷以 pymsg 帧形式写入 w。
|
||||
func WriteFrame(w io.Writer, payload []byte) error {
|
||||
frame, err := Encode(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = w.Write(frame)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReadFrame 从 r 读取一条 pymsg 帧，返回载荷。
|
||||
// max 为单条消息大小上限（传 0 表示使用 MaxMsgSize）。
|
||||
// 长度前缀超过上限视为协议错误（防止恶意长度导致内存耗尽）。
|
||||
func ReadFrame(r io.Reader, max int) ([]byte, error) {
|
||||
if max <= 0 || max > MaxMsgSize {
|
||||
max = MaxMsgSize
|
||||
}
|
||||
header := make([]byte, headerSize)
|
||||
if _, err := io.ReadFull(r, header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
length := binary.BigEndian.Uint32(header)
|
||||
if length > uint32(max) {
|
||||
return nil, fmt.Errorf("message too large: %d bytes (max %d)", length, max)
|
||||
}
|
||||
payload := make([]byte, length)
|
||||
if _, err := io.ReadFull(r, payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEncodeRoundTrip(t *testing.T) {
|
||||
payload := []byte(`{"hello":"世界"}`)
|
||||
frame, err := Encode(payload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(frame) != 4+len(payload) {
|
||||
t.Fatalf("frame length = %d, want %d", len(frame), 4+len(payload))
|
||||
}
|
||||
if got := binary.BigEndian.Uint32(frame[:4]); got != uint32(len(payload)) {
|
||||
t.Fatalf("header = %d, want %d", got, len(payload))
|
||||
}
|
||||
got, err := ReadFrame(bytes.NewReader(frame), 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !bytes.Equal(got, payload) {
|
||||
t.Fatalf("payload mismatch: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeTooLarge(t *testing.T) {
|
||||
payload := make([]byte, MaxMsgSize+1)
|
||||
if _, err := Encode(payload); err == nil {
|
||||
t.Fatal("expected error for oversized payload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameTooLarge(t *testing.T) {
|
||||
header := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(header, MaxMsgSize+1)
|
||||
_, err := ReadFrame(bytes.NewReader(header), 0)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for oversized header")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameTruncated(t *testing.T) {
|
||||
frame, _ := Encode([]byte("abcdef"))
|
||||
if _, err := ReadFrame(bytes.NewReader(frame[:len(frame)-3]), 0); err == nil {
|
||||
t.Fatal("expected error for truncated payload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameMaxParam(t *testing.T) {
|
||||
frame, _ := Encode([]byte("hello"))
|
||||
if _, err := ReadFrame(bytes.NewReader(frame), 4); err == nil {
|
||||
t.Fatal("expected error when max < payload length")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameEOF(t *testing.T) {
|
||||
if _, err := ReadFrame(bytes.NewReader(nil), 0); !errors.Is(err, io.EOF) {
|
||||
t.Fatalf("expected EOF, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
module github.com/micro_servers/gomsg
|
||||
|
||||
go 1.22
|
||||
@@ -0,0 +1,173 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 集中式日志客户端：通过 pymsg/TCP 调用 db2（或 db）execute 路由，
|
||||
// 逐条 INSERT 到 Logs 表。对应各服务目录下的 log_client.py 与 C++ liblogc。
|
||||
//
|
||||
// 设计要点（与 waf internal/logclient / Python log_client.py 对齐）：
|
||||
// - 异步非阻塞：AsyncQueue 批量推送，不拖慢业务
|
||||
// - fail-safe：db 服务不可用时静默丢弃，绝不影响业务
|
||||
// - 缓冲批量：满 50 条或超时 2 秒自动 flush
|
||||
|
||||
const defaultDBAddr = "127.0.0.1:8005"
|
||||
|
||||
// 日志级别数值（越小越详细）。
|
||||
const (
|
||||
levelDebug = 0
|
||||
levelInfo = 1
|
||||
levelWarn = 2
|
||||
levelError = 3
|
||||
)
|
||||
|
||||
var levelMap = map[string]int{
|
||||
"debug": levelDebug, "DEBUG": levelDebug,
|
||||
"info": levelInfo, "INFO": levelInfo,
|
||||
"warn": levelWarn, "warning": levelWarn, "WARN": levelWarn, "WARNING": levelWarn,
|
||||
"error": levelError, "ERROR": levelError,
|
||||
}
|
||||
|
||||
// Entry 单条日志。
|
||||
type Entry struct {
|
||||
Service string
|
||||
Level string
|
||||
Message string
|
||||
TraceID string
|
||||
Extra map[string]any
|
||||
}
|
||||
|
||||
// LogClient 异步批量日志推送客户端。
|
||||
type LogClient struct {
|
||||
queue *AsyncQueue[Entry]
|
||||
client *MsgClient
|
||||
minLevel int
|
||||
}
|
||||
|
||||
// NewLogClient 创建日志客户端并启动后台批量推送。
|
||||
// dbAddr 为空时依次取环境变量 LOG_DB_ADDR、默认 127.0.0.1:8005；
|
||||
// 默认最低级别 info，可用 SetLevel 调整（环境变量 LOG_LEVEL 同样生效）。
|
||||
func NewLogClient(dbAddr string) *LogClient {
|
||||
if dbAddr == "" {
|
||||
dbAddr = os.Getenv("LOG_DB_ADDR")
|
||||
}
|
||||
if dbAddr == "" {
|
||||
dbAddr = defaultDBAddr
|
||||
}
|
||||
lc := &LogClient{
|
||||
client: NewMsgClient(dbAddr, 5*time.Second),
|
||||
minLevel: levelInfo,
|
||||
}
|
||||
lc.queue = NewAsyncQueue[Entry](lc.sendBatch, AsyncQueueConfig{})
|
||||
if lv := os.Getenv("LOG_LEVEL"); lv != "" {
|
||||
lc.SetLevel(lv)
|
||||
}
|
||||
return lc
|
||||
}
|
||||
|
||||
// SetLevel 设置最低日志级别（debug|info|warn|error，大小写不敏感）。
|
||||
func (c *LogClient) SetLevel(level string) {
|
||||
if n, ok := levelMap[level]; ok {
|
||||
c.minLevel = n
|
||||
}
|
||||
}
|
||||
|
||||
// Push 推送一条日志（非阻塞，队列满或低于最低级别时静默丢弃）。
|
||||
func (c *LogClient) Push(e Entry) bool {
|
||||
n, ok := levelMap[e.Level]
|
||||
if !ok || n < c.minLevel {
|
||||
return false
|
||||
}
|
||||
return c.queue.Push(e)
|
||||
}
|
||||
|
||||
// Debug 推送 DEBUG 级别日志。
|
||||
func (c *LogClient) Debug(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "DEBUG", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Info 推送 INFO 级别日志。
|
||||
func (c *LogClient) Info(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "INFO", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Warning 推送 WARNING 级别日志。
|
||||
func (c *LogClient) Warning(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "WARNING", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Error 推送 ERROR 级别日志。
|
||||
func (c *LogClient) Error(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "ERROR", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Flush 停止队列并 flush 剩余日志（用于进程退出前）。
|
||||
func (c *LogClient) Flush() {
|
||||
c.queue.Shutdown()
|
||||
}
|
||||
|
||||
// dbV2Request execute 路由的 INSERT 请求体（与 Python/C++ 各 log 客户端一致）。
|
||||
type dbV2Request struct {
|
||||
Route string `json:"_route"`
|
||||
Action string `json:"action"`
|
||||
Table string `json:"table"`
|
||||
Values map[string]any `json:"values"`
|
||||
}
|
||||
|
||||
// dbV2Response execute 路由的通用响应。
|
||||
type dbV2Response struct {
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func (c *LogClient) sendBatch(batch []Entry) {
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
for _, e := range batch {
|
||||
c.sendOne(e, now)
|
||||
}
|
||||
}
|
||||
|
||||
// sendOne 通过 pymsg TCP 协议发送单条日志到 db2（fail-safe：任何错误静默丢弃）。
|
||||
func (c *LogClient) sendOne(e Entry, ts string) {
|
||||
extraJSON := ""
|
||||
if e.Extra != nil {
|
||||
b, err := json.Marshal(e.Extra)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
extraJSON = string(b)
|
||||
}
|
||||
|
||||
values := map[string]any{
|
||||
"ts": ts,
|
||||
"service": e.Service,
|
||||
"level": e.Level,
|
||||
"message": e.Message,
|
||||
}
|
||||
if e.TraceID != "" {
|
||||
values["trace_id"] = e.TraceID
|
||||
}
|
||||
if extraJSON != "" {
|
||||
values["extra"] = extraJSON
|
||||
}
|
||||
|
||||
req := dbV2Request{Route: "execute", Action: "insert", Table: "Logs", Values: values}
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := c.client.Call(data)
|
||||
if err != nil {
|
||||
return // fail-safe
|
||||
}
|
||||
|
||||
var r dbV2Response
|
||||
if json.Unmarshal(resp, &r) != nil || !r.Success {
|
||||
fmt.Fprintf(os.Stderr, "[gomsg:logclient] INSERT failed: %s\n", r.Error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeDB2 模拟 db2 的 execute 路由：校验 Logs INSERT 请求格式并返回成功。
|
||||
type fakeDB2 struct {
|
||||
server *MsgServer
|
||||
mu sync.Mutex
|
||||
inserts []map[string]any
|
||||
failures int
|
||||
}
|
||||
|
||||
func newFakeDB2(t *testing.T) *fakeDB2 {
|
||||
f := &fakeDB2{}
|
||||
f.server = NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
var req map[string]any
|
||||
if err := json.Unmarshal(payload, &req); err != nil {
|
||||
return []byte(`{"success":false,"error":"invalid JSON request"}`)
|
||||
}
|
||||
// db2 真实语义：非法 JSON 返回 success=false；其余按通用响应
|
||||
if req["_route"] != "execute" {
|
||||
return []byte(`{"success":false,"error":"未知路由"}`)
|
||||
}
|
||||
f.mu.Lock()
|
||||
f.inserts = append(f.inserts, req)
|
||||
f.mu.Unlock()
|
||||
return []byte(`{"success":true,"data":{}}`)
|
||||
})
|
||||
if err := f.server.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(f.server.Stop)
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeDB2) count() int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return len(f.inserts)
|
||||
}
|
||||
|
||||
func TestLogClientPushesToDB2(t *testing.T) {
|
||||
db := newFakeDB2(t)
|
||||
lc := NewLogClient(db.server.Addr().String())
|
||||
defer lc.Flush()
|
||||
|
||||
lc.SetLevel("debug")
|
||||
lc.Debug("test-svc", "debug msg", nil)
|
||||
lc.Info("test-svc", "hello", map[string]any{"k": "v"})
|
||||
lc.Error("test-svc", "boom", nil)
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for db.count() < 3 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("logs not pushed in time")
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
first := db.inserts[0]
|
||||
if first["_route"] != "execute" {
|
||||
t.Fatalf("_route = %v", first["_route"])
|
||||
}
|
||||
vals := first["values"].(map[string]any)
|
||||
if vals["service"] != "test-svc" || vals["level"] != "DEBUG" || vals["message"] != "debug msg" {
|
||||
t.Fatalf("values = %v", vals)
|
||||
}
|
||||
// ts 必须是 ISO 8601 字符串
|
||||
if _, ok := vals["ts"].(string); !ok {
|
||||
t.Fatalf("ts missing: %v", vals)
|
||||
}
|
||||
// 第二条（Info）的 extra 应序列化为 JSON 字符串
|
||||
vals2 := db.inserts[1]["values"].(map[string]any)
|
||||
if vals2["extra"] != `{"k":"v"}` {
|
||||
t.Fatalf("extra = %v", vals2["extra"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogClientLevelFilter(t *testing.T) {
|
||||
db := newFakeDB2(t)
|
||||
lc := NewLogClient(db.server.Addr().String())
|
||||
defer lc.Flush()
|
||||
|
||||
lc.Info("svc", "visible", nil)
|
||||
lc.Debug("svc", "hidden", nil) // 默认 minLevel=info，debug 应被过滤
|
||||
lc.Push(Entry{Service: "svc", Level: "NOPE", Message: "hidden"})
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for db.count() < 1 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("expected exactly 1 log")
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond) // 确认没有多余的推送
|
||||
if db.count() != 1 {
|
||||
t.Fatalf("expected 1 insert, got %d", db.count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogClientFailSafeWhenDBDown(t *testing.T) {
|
||||
lc := NewLogClient("127.0.0.1:1") // 不可达地址
|
||||
defer lc.Flush()
|
||||
lc.Info("svc", "will be dropped", nil)
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
// 不 panic、不阻塞即视为通过（fail-safe 语义）
|
||||
}
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// JSON RPC 约定层（与 Python pymsg.rpc / C++ libmsg RpcClient/RpcServer 兼容）：
|
||||
//
|
||||
// 请求: {"_route": "<路由名>", ...业务字段}
|
||||
// 响应: {"success": true, "data": ...}
|
||||
// {"success": false, "error": "..."}
|
||||
//
|
||||
// 数字一律以 json.Number 保留原始精度（避免 int64 经 float64 丢失）。
|
||||
|
||||
// RpcError 远端返回 success=false 的业务错误，或响应不符合 RPC 约定。
|
||||
type RpcError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *RpcError) Error() string { return e.msg }
|
||||
|
||||
// RpcHandler 处理一条 RPC 请求，返回值放入 data，错误包装为 success=false。
|
||||
type RpcHandler func(req map[string]any) (any, error)
|
||||
|
||||
// RpcClient JSON RPC 客户端（复用 MsgClient 长连接）。
|
||||
type RpcClient struct {
|
||||
client *MsgClient
|
||||
}
|
||||
|
||||
// NewRpcClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。
|
||||
func NewRpcClient(addr string, timeout time.Duration) *RpcClient {
|
||||
return &RpcClient{client: NewMsgClient(addr, timeout)}
|
||||
}
|
||||
|
||||
// Call 调用指定路由，返回响应中的 data 字段（可能为 nil）。
|
||||
// success=false 或非法响应时返回 *RpcError。
|
||||
func (c *RpcClient) Call(route string, payload map[string]any) (any, error) {
|
||||
req := make(map[string]any, len(payload)+1)
|
||||
for k, v := range payload {
|
||||
req[k] = v
|
||||
}
|
||||
req["_route"] = route
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := c.client.Call(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var obj map[string]any
|
||||
dec := json.NewDecoder(bytes.NewReader(resp))
|
||||
dec.UseNumber()
|
||||
if err := dec.Decode(&obj); err != nil || obj == nil {
|
||||
return nil, &RpcError{msg: fmt.Sprintf("响应不是合法 JSON: %v", err)}
|
||||
}
|
||||
success, ok := obj["success"].(bool)
|
||||
if !ok {
|
||||
return nil, &RpcError{msg: fmt.Sprintf("非法 RPC 响应: %s", truncate(resp))}
|
||||
}
|
||||
if !success {
|
||||
errMsg, _ := obj["error"].(string)
|
||||
return nil, &RpcError{msg: errMsg}
|
||||
}
|
||||
return obj["data"], nil
|
||||
}
|
||||
|
||||
// Close 关闭底层连接。
|
||||
func (c *RpcClient) Close() error { return c.client.Close() }
|
||||
|
||||
// RpcServer JSON RPC 服务端：基于 MsgServer 的路由分发与响应包装。
|
||||
type RpcServer struct {
|
||||
*MsgServer
|
||||
routes map[string]RpcHandler
|
||||
mu sync.RWMutex
|
||||
// OnError 可选钩子：handler 抛错时调用（用于记日志）。
|
||||
OnError func(route string, err error)
|
||||
// JSONMarshal 响应序列化函数，默认 json.Marshal（对应 Python 的 json_default）。
|
||||
JSONMarshal func(v any) ([]byte, error)
|
||||
}
|
||||
|
||||
// NewRpcServer 创建 RPC 服务端（不监听，需调用 Start）。
|
||||
func NewRpcServer(addr string) *RpcServer {
|
||||
s := &RpcServer{routes: make(map[string]RpcHandler)}
|
||||
s.MsgServer = NewMsgServer(addr, s.dispatch)
|
||||
return s
|
||||
}
|
||||
|
||||
// AddRoute 注册路由处理函数。
|
||||
func (s *RpcServer) AddRoute(name string, fn RpcHandler) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.routes[name] = fn
|
||||
}
|
||||
|
||||
// dispatch 解析请求、查找路由、包装响应；任何输入错误都不向客户端回显内部细节。
|
||||
// handler 内 panic 会被恢复并包装为 success=false（对应 Python 的 except Exception）。
|
||||
func (s *RpcServer) dispatch(payload []byte) (reply []byte) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err := fmt.Errorf("handler panic: %v", r)
|
||||
if s.OnError != nil {
|
||||
s.OnError("", err)
|
||||
}
|
||||
reply = s.reply(false, nil, sanitizeError(err.Error()))
|
||||
}
|
||||
}()
|
||||
|
||||
req, err := decodeRequest(payload)
|
||||
if err != nil {
|
||||
return s.reply(false, nil, "invalid JSON request")
|
||||
}
|
||||
|
||||
routeVal, has := req["_route"]
|
||||
delete(req, "_route")
|
||||
route := ""
|
||||
switch r := routeVal.(type) {
|
||||
case string:
|
||||
route = r
|
||||
case nil:
|
||||
// 保持 Python 行为：缺省 _route 时按空路由查找（未知路由）
|
||||
default:
|
||||
if has {
|
||||
route = fmt.Sprintf("%v", r)
|
||||
}
|
||||
}
|
||||
|
||||
s.mu.RLock()
|
||||
handler := s.routes[route]
|
||||
s.mu.RUnlock()
|
||||
if handler == nil {
|
||||
return s.reply(false, nil, "未知路由: "+route)
|
||||
}
|
||||
|
||||
data, err := handler(req)
|
||||
if err != nil {
|
||||
if s.OnError != nil {
|
||||
s.OnError(route, err)
|
||||
}
|
||||
return s.reply(false, nil, sanitizeError(err.Error()))
|
||||
}
|
||||
return s.reply(true, data, "")
|
||||
}
|
||||
|
||||
func (s *RpcServer) reply(success bool, data any, errMsg string) []byte {
|
||||
resp := map[string]any{"success": success}
|
||||
if data != nil {
|
||||
resp["data"] = data
|
||||
}
|
||||
if errMsg != "" {
|
||||
resp["error"] = errMsg
|
||||
}
|
||||
marshal := s.JSONMarshal
|
||||
if marshal == nil {
|
||||
marshal = json.Marshal
|
||||
}
|
||||
out, err := marshal(resp)
|
||||
if err != nil {
|
||||
out, _ = json.Marshal(map[string]any{
|
||||
"success": false,
|
||||
"error": sanitizeError(err.Error()),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// decodeRequest 解析请求 JSON：UseNumber 保留整数精度；
|
||||
// Go 的 encoding/json 内置 10000 层嵌套上限，数字位数字符串化保存，
|
||||
// 输入大小又受帧协议 64 MiB 上限约束，天然规避深度/超大数字 DoS。
|
||||
func decodeRequest(payload []byte) (map[string]any, error) {
|
||||
dec := json.NewDecoder(bytes.NewReader(payload))
|
||||
dec.UseNumber()
|
||||
var req map[string]any
|
||||
if err := dec.Decode(&req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("empty request")
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// pathRe 匹配绝对路径（如 /opt/micro_servers/db/xxx.py），用于脱敏异常信息。
|
||||
var pathRe = regexp.MustCompile(`/[A-Za-z0-9_.\-]+(?:/[A-Za-z0-9_.\-]+)*`)
|
||||
|
||||
// sanitizeError 脱敏异常信息中的绝对路径等内部细节，保留业务错误语义
|
||||
// （对齐 Python pymsg._sanitize_error：前导字符为单词字符时不匹配）。
|
||||
func sanitizeError(msg string) string {
|
||||
if msg == "" {
|
||||
return msg
|
||||
}
|
||||
var b strings.Builder
|
||||
last := 0
|
||||
for _, idx := range pathRe.FindAllStringIndex(msg, -1) {
|
||||
start, end := idx[0], idx[1]
|
||||
if start > 0 && isWordChar(msg[start-1]) {
|
||||
continue
|
||||
}
|
||||
b.WriteString(msg[last:start])
|
||||
b.WriteString("<redacted>")
|
||||
last = end
|
||||
}
|
||||
b.WriteString(msg[last:])
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func isWordChar(c byte) bool {
|
||||
return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')
|
||||
}
|
||||
|
||||
func truncate(b []byte) string {
|
||||
const limit = 200
|
||||
if len(b) <= limit {
|
||||
return string(b)
|
||||
}
|
||||
return string(b[:limit]) + "..."
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// startRPCServer 启动带两路由的 RPC 服务端。
|
||||
func startRPCServer(t *testing.T) (*RpcServer, string) {
|
||||
t.Helper()
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("echo", func(req map[string]any) (any, error) {
|
||||
return req, nil
|
||||
})
|
||||
s.AddRoute("sum", func(req map[string]any) (any, error) {
|
||||
a := req["a"].(json.Number)
|
||||
b := req["b"].(json.Number)
|
||||
ai, _ := a.Int64()
|
||||
bi, _ := b.Int64()
|
||||
return ai + bi, nil
|
||||
})
|
||||
s.AddRoute("fail", func(req map[string]any) (any, error) {
|
||||
return nil, errors.New("boom: /opt/micro_servers/db2/xxx.py failed")
|
||||
})
|
||||
s.AddRoute("panic", func(req map[string]any) (any, error) {
|
||||
panic("panic in handler")
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(s.Stop)
|
||||
return s, s.Addr().String()
|
||||
}
|
||||
|
||||
func TestRpcClientCallSuccess(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
data, err := c.Call("echo", map[string]any{"hello": "world"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m := data.(map[string]any)
|
||||
if m["hello"] != "world" {
|
||||
t.Fatalf("data = %v", data)
|
||||
}
|
||||
if _, has := m["_route"]; has {
|
||||
t.Fatal("_route should be stripped before handler")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientNumberPrecision(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
data, err := c.Call("sum", map[string]any{"a": json.Number("9007199254740993"), "b": json.Number("1")})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := data.(json.Number).String(); got != "9007199254740994" {
|
||||
t.Fatalf("sum = %s, want exact int64 precision", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientBusinessError(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
_, err := c.Call("fail", nil)
|
||||
var rpcErr *RpcError
|
||||
if !errors.As(err, &rpcErr) {
|
||||
t.Fatalf("expected *RpcError, got %T: %v", err, err)
|
||||
}
|
||||
// 绝对路径应被脱敏
|
||||
if strings.Contains(rpcErr.Error(), "/opt/") {
|
||||
t.Fatalf("error not sanitized: %v", rpcErr)
|
||||
}
|
||||
if !strings.Contains(rpcErr.Error(), "<redacted>") {
|
||||
t.Fatalf("expected <redacted> in error: %v", rpcErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcUnknownRoute(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
_, err := c.Call("no_such_route", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "未知路由") {
|
||||
t.Fatalf("expected 未知路由 error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerInvalidJSON(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
resp, err := c.Call([]byte("{not json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var r map[string]any
|
||||
if json.Unmarshal(resp, &r) != nil {
|
||||
t.Fatalf("invalid response: %q", resp)
|
||||
}
|
||||
if r["success"] != false || !strings.Contains(r["error"].(string), "invalid JSON request") {
|
||||
t.Fatalf("unexpected response: %v", r)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerPanicHandler(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
// handler 内 panic 不应打崩服务，后续请求仍可用
|
||||
if _, err := c.Call("panic", nil); err == nil {
|
||||
t.Fatal("expected error from panicking handler")
|
||||
}
|
||||
if _, err := c.Call("echo", map[string]any{"ok": true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerOnErrorHook(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
var hookRoute string
|
||||
s.OnError = func(route string, err error) { hookRoute = route }
|
||||
s.AddRoute("fail", func(req map[string]any) (any, error) {
|
||||
return nil, errors.New("boom")
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewRpcClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
_, _ = c.Call("fail", nil)
|
||||
if hookRoute != "fail" {
|
||||
t.Fatalf("OnError hook route = %q", hookRoute)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerNonDictRequest(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("echo", func(req map[string]any) (any, error) { return req, nil })
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
resp, err := c.Call([]byte(`["not","a","dict"]`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var r map[string]any
|
||||
if json.Unmarshal(resp, &r) != nil || r["success"] != false {
|
||||
t.Fatalf("expected failure response, got %q", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeError(t *testing.T) {
|
||||
cases := []struct{ in, want string }{
|
||||
{"", ""},
|
||||
{"boom", "boom"},
|
||||
{"failed: /opt/micro_servers/db2/xxx.py", "failed: <redacted>"},
|
||||
{"no abs path here", "no abs path here"},
|
||||
{"abc/def 保持原样", "abc/def 保持原样"}, // 前导字符为单词字符时不匹配（对齐 Python lookbehind）
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := sanitizeError(c.in); got != c.want {
|
||||
t.Errorf("sanitizeError(%q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientTimeout(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("slow", func(req map[string]any) (any, error) {
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
return nil, nil
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewRpcClient(s.Addr().String(), 50*time.Millisecond)
|
||||
defer c.Close()
|
||||
if _, err := c.Call("slow", nil); err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
}
|
||||
+130
@@ -0,0 +1,130 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// MsgHandler 处理一条完整消息，返回回复载荷；返回 nil 表示不回复。
|
||||
type MsgHandler func(payload []byte) []byte
|
||||
|
||||
// MsgServer pymsg 协议服务端：accept 循环 + 每连接一个 goroutine 串行处理。
|
||||
// 对应 Python pymsg.MsgServer / C++ libmsg.MsgServer。
|
||||
type MsgServer struct {
|
||||
addr string
|
||||
handler MsgHandler
|
||||
// ReadTimeout 单条消息读取总超时；0 表示不设超时（默认，与 Python 行为一致）。
|
||||
ReadTimeout time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
ln net.Listener
|
||||
conns map[net.Conn]struct{}
|
||||
closed bool
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewMsgServer 创建服务端（不监听，需调用 Start）。
|
||||
func NewMsgServer(addr string, handler MsgHandler) *MsgServer {
|
||||
return &MsgServer{
|
||||
addr: addr,
|
||||
handler: handler,
|
||||
conns: make(map[net.Conn]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Addr 返回监听地址（Start 后可用，便于测试取随机端口）。
|
||||
func (s *MsgServer) Addr() net.Addr {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.ln == nil {
|
||||
return nil
|
||||
}
|
||||
return s.ln.Addr()
|
||||
}
|
||||
|
||||
// Start 开始监听并启动 accept 循环（非阻塞）。
|
||||
func (s *MsgServer) Start() error {
|
||||
ln, err := net.Listen("tcp", s.addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
_ = ln.Close()
|
||||
return nil
|
||||
}
|
||||
s.ln = ln
|
||||
s.mu.Unlock()
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.acceptLoop(ln)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 关闭监听与全部连接，等待处理中的请求结束。
|
||||
func (s *MsgServer) Stop() {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
if s.ln != nil {
|
||||
_ = s.ln.Close()
|
||||
}
|
||||
for conn := range s.conns {
|
||||
_ = conn.Close()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
s.wg.Wait()
|
||||
}
|
||||
|
||||
func (s *MsgServer) acceptLoop(ln net.Listener) {
|
||||
defer s.wg.Done()
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
return // 监听关闭
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
return
|
||||
}
|
||||
s.conns[conn] = struct{}{}
|
||||
s.mu.Unlock()
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MsgServer) handleConn(conn net.Conn) {
|
||||
defer s.wg.Done()
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(s.conns, conn)
|
||||
s.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
if s.ReadTimeout > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
|
||||
}
|
||||
payload, err := ReadFrame(conn, 0)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
reply := s.handler(payload)
|
||||
if reply == nil {
|
||||
continue
|
||||
}
|
||||
if err := WriteFrame(conn, reply); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -144,7 +144,7 @@ waf/
|
||||
│ ├── engine/
|
||||
│ │ └── engine.go # coraza 封装：加载/热加载/列举/单测/输入检测
|
||||
│ ├── logclient/
|
||||
│ │ └── logclient.go # Go 版 pymsg 日志客户端（异步批量推送）
|
||||
│ │ └── logclient.go # gomsg 日志客户端适配层（协议/队列实现见 ../gomsg）
|
||||
│ └── proxy/
|
||||
│ └── proxy.go # ReverseProxy + WAF 事务接入 + 安全响应头注入
|
||||
├── rules/
|
||||
|
||||
@@ -4,9 +4,12 @@ go 1.22.0
|
||||
|
||||
require (
|
||||
github.com/corazawaf/coraza/v3 v3.3.0
|
||||
github.com/micro_servers/gomsg v0.0.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
replace github.com/micro_servers/gomsg => ../gomsg
|
||||
|
||||
require (
|
||||
github.com/corazawaf/libinjection-go v0.2.2 // indirect
|
||||
github.com/magefile/mage v1.15.1-0.20241126214340-bdc92f694516 // indirect
|
||||
|
||||
@@ -1,142 +1,38 @@
|
||||
// Package logclient 提供异步批量日志推送客户端。
|
||||
// Package logclient 基于 gomsg 的集中式日志客户端（异步批量推送 db2 Logs 表）。
|
||||
//
|
||||
// 通过 pymsg TCP 协议调用 dbv2 execute 路由，逐条 INSERT 到 Logs 表。
|
||||
//
|
||||
// 设计要点（与 Python example_log_client.py 对齐）：
|
||||
// - 异步非阻塞：内存 channel + 后台 goroutine 批量推送，不拖慢业务
|
||||
// - fail-safe：db 服务不可用时静默丢弃，绝不影响业务
|
||||
// - 缓冲批量：队列满 50 条或超时 2 秒时自动 flush
|
||||
//
|
||||
// 协议：pymsg 长度前缀帧（4 字节 big-endian uint32 长度 + UTF-8 载荷）
|
||||
// 原实现（pymsg 帧协议 + 异步队列）已抽取至 ../gomsg，本包保留包级 API
|
||||
// 兼容层：main / engine / proxy 的调用点零改动。
|
||||
package logclient
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
import "github.com/micro_servers/gomsg"
|
||||
|
||||
var (
|
||||
lc *gomsg.LogClient
|
||||
pendingLevel string
|
||||
)
|
||||
|
||||
const (
|
||||
defaultDBAddr = "127.0.0.1:8005"
|
||||
batchSize = 50
|
||||
flushInterval = 2 * time.Second
|
||||
socketTimeout = 5 * time.Second
|
||||
// 消息大小安全上限（64 MiB），与 Python / C++ 端保持一致。
|
||||
maxMsgSize = 64 * 1024 * 1024
|
||||
)
|
||||
|
||||
// 日志级别数值（越小越详细）。
|
||||
const (
|
||||
levelDebug = 0
|
||||
levelInfo = 1
|
||||
levelWarn = 2
|
||||
levelError = 3
|
||||
)
|
||||
|
||||
var levelMap = map[string]int{
|
||||
"debug": levelDebug,
|
||||
"info": levelInfo,
|
||||
"warn": levelWarn,
|
||||
"warning": levelWarn,
|
||||
"error": levelError,
|
||||
"DEBUG": levelDebug,
|
||||
"INFO": levelInfo,
|
||||
"WARN": levelWarn,
|
||||
"WARNING": levelWarn,
|
||||
"ERROR": levelError,
|
||||
}
|
||||
|
||||
// dbAddr 通过环境变量 LOG_DB_ADDR 配置（默认 127.0.0.1:8005），支持非标准部署场景。
|
||||
var dbAddr = defaultDBAddr
|
||||
|
||||
// minLevel 控制最低输出级别，默认 info（debug 被过滤）。
|
||||
var minLevel = levelInfo
|
||||
|
||||
func init() {
|
||||
if v := os.Getenv("LOG_DB_ADDR"); v != "" {
|
||||
dbAddr = v
|
||||
}
|
||||
if v := os.Getenv("LOG_LEVEL"); v != "" {
|
||||
SetLevel(v)
|
||||
// Init 初始化全局日志客户端（读 LOG_DB_ADDR / LOG_LEVEL 环境变量）。
|
||||
func Init() {
|
||||
lc = gomsg.NewLogClient("")
|
||||
if pendingLevel != "" {
|
||||
lc.SetLevel(pendingLevel)
|
||||
}
|
||||
}
|
||||
|
||||
// SetLevel 设置最低日志级别（debug|info|warn|error，大小写不敏感）。
|
||||
func SetLevel(level string) {
|
||||
if n, ok := levelMap[level]; ok {
|
||||
minLevel = n
|
||||
pendingLevel = level
|
||||
if lc != nil {
|
||||
lc.SetLevel(level)
|
||||
}
|
||||
}
|
||||
|
||||
// Entry 单条日志。
|
||||
type Entry struct {
|
||||
Service string `json:"service"`
|
||||
Level string `json:"level"`
|
||||
Message string `json:"message"`
|
||||
TraceID string `json:"trace_id,omitempty"`
|
||||
Extra map[string]any `json:"extra,omitempty"`
|
||||
}
|
||||
|
||||
// dbV2Request execute 路由的 INSERT 请求体。
|
||||
type dbV2Request struct {
|
||||
Route string `json:"_route"`
|
||||
Action string `json:"action"`
|
||||
Table string `json:"table"`
|
||||
Values map[string]any `json:"values"`
|
||||
}
|
||||
|
||||
// dbV2Response execute 路由的通用响应。
|
||||
type dbV2Response struct {
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// Client 异步批量日志推送客户端。
|
||||
// send 仅在 flushLoop 单 goroutine 中调用，conn 复用无需额外加锁。
|
||||
type Client struct {
|
||||
ch chan Entry
|
||||
stop chan struct{}
|
||||
wg sync.WaitGroup
|
||||
stopped bool
|
||||
mu sync.Mutex
|
||||
conn net.Conn // 复用的 pymsg 长连接，失败时重建
|
||||
}
|
||||
|
||||
// global 全局单例。
|
||||
var global *Client
|
||||
|
||||
// Init 初始化全局日志客户端并启动后台 flush goroutine。
|
||||
func Init() {
|
||||
global = New()
|
||||
}
|
||||
|
||||
// New 创建客户端实例。
|
||||
func New() *Client {
|
||||
c := &Client{
|
||||
ch: make(chan Entry, 256),
|
||||
stop: make(chan struct{}),
|
||||
}
|
||||
c.wg.Add(1)
|
||||
go c.flushLoop()
|
||||
return c
|
||||
}
|
||||
|
||||
// Push 推送一条日志（非阻塞，队列满时丢弃）。
|
||||
// 低于 SetLevel 设定的最低级别的日志将被静默丢弃。
|
||||
// Push 推送一条日志（非阻塞，队列满或低于最低级别时静默丢弃）。
|
||||
func Push(service, level, message string, extra map[string]any) {
|
||||
if global == nil {
|
||||
if lc == nil {
|
||||
return
|
||||
}
|
||||
n, ok := levelMap[level]
|
||||
if !ok || n < minLevel {
|
||||
return
|
||||
}
|
||||
global.push(Entry{
|
||||
lc.Push(gomsg.Entry{
|
||||
Service: service,
|
||||
Level: level,
|
||||
Message: message,
|
||||
@@ -164,199 +60,9 @@ func Error(service, message string, extra map[string]any) {
|
||||
Push(service, "ERROR", message, extra)
|
||||
}
|
||||
|
||||
// Flush 同步 flush 队列中剩余日志（用于进程退出前）。
|
||||
// Flush 停止队列并 flush 剩余日志（用于进程退出前）。
|
||||
func Flush() {
|
||||
if global == nil {
|
||||
return
|
||||
}
|
||||
global.flush()
|
||||
}
|
||||
|
||||
func (c *Client) push(e Entry) {
|
||||
select {
|
||||
case c.ch <- e:
|
||||
default:
|
||||
// 队列满，静默丢弃（fail-safe）
|
||||
if lc != nil {
|
||||
lc.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) flushLoop() {
|
||||
defer c.wg.Done()
|
||||
batch := make([]Entry, 0, batchSize)
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case e := <-c.ch:
|
||||
batch = append(batch, e)
|
||||
// 非阻塞地继续取，直到满 batchSize
|
||||
for len(batch) < batchSize {
|
||||
select {
|
||||
case e2 := <-c.ch:
|
||||
batch = append(batch, e2)
|
||||
default:
|
||||
goto send
|
||||
}
|
||||
}
|
||||
send:
|
||||
if len(batch) >= batchSize {
|
||||
c.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
c.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-c.stop:
|
||||
// 退出前 drain channel
|
||||
for {
|
||||
select {
|
||||
case e := <-c.ch:
|
||||
batch = append(batch, e)
|
||||
default:
|
||||
if len(batch) > 0 {
|
||||
c.send(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) send(batch []Entry) {
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
for _, entry := range batch {
|
||||
c.sendOne(entry, now)
|
||||
}
|
||||
}
|
||||
|
||||
// sendOne 通过 pymsg TCP 协议发送单条日志到 dbv2。
|
||||
func (c *Client) sendOne(entry Entry, ts string) {
|
||||
extraJSON := ""
|
||||
if entry.Extra != nil {
|
||||
b, err := json.Marshal(entry.Extra)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
extraJSON = string(b)
|
||||
}
|
||||
|
||||
values := map[string]any{
|
||||
"ts": ts,
|
||||
"service": entry.Service,
|
||||
"level": entry.Level,
|
||||
"message": entry.Message,
|
||||
}
|
||||
if entry.TraceID != "" {
|
||||
values["trace_id"] = entry.TraceID
|
||||
}
|
||||
if extraJSON != "" {
|
||||
values["extra"] = extraJSON
|
||||
}
|
||||
|
||||
req := dbV2Request{
|
||||
Route: "execute",
|
||||
Action: "insert",
|
||||
Table: "Logs",
|
||||
Values: values,
|
||||
}
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// pymsg 协议：4 字节 big-endian 长度 + UTF-8 载荷（复用长连接，失败时重连一次）
|
||||
resp, err := c.pymsgCall(data)
|
||||
if err != nil {
|
||||
return // fail-safe
|
||||
}
|
||||
|
||||
var r dbV2Response
|
||||
if json.Unmarshal(resp, &r) != nil || !r.Success {
|
||||
fmt.Fprintf(os.Stderr, "[logclient] INSERT failed: %s\n", r.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// pymsgCall 复用长连接发送一条 pymsg 请求并读取响应。
|
||||
// 首次失败时关闭旧连接重建一次，避免高并发拦截场景下大量 TIME_WAIT 短连接。
|
||||
func (c *Client) pymsgCall(payload []byte) ([]byte, error) {
|
||||
resp, err := c.doPymsg(payload)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
// 旧连接可能已断开（db 重启等），关闭后重连一次
|
||||
c.closeConn()
|
||||
return c.doPymsg(payload)
|
||||
}
|
||||
|
||||
// doPymsg 在当前连接上执行一次请求-响应交互。
|
||||
func (c *Client) doPymsg(payload []byte) ([]byte, error) {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conn.SetDeadline(time.Now().Add(socketTimeout))
|
||||
|
||||
// 发送：4 字节长度前缀 + 载荷
|
||||
header := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(header, uint32(len(payload)))
|
||||
if _, err := conn.Write(header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := conn.Write(payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 接收：4 字节长度前缀 + 载荷
|
||||
if _, err := io.ReadFull(conn, header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
respLen := binary.BigEndian.Uint32(header)
|
||||
if respLen > maxMsgSize {
|
||||
return nil, fmt.Errorf("response too large: %d bytes (max %d)", respLen, maxMsgSize)
|
||||
}
|
||||
resp := make([]byte, respLen)
|
||||
if _, err := io.ReadFull(conn, resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// getConn 返回当前复用连接，不存在时建立。
|
||||
func (c *Client) getConn() (net.Conn, error) {
|
||||
if c.conn != nil {
|
||||
return c.conn, nil
|
||||
}
|
||||
conn, err := net.DialTimeout("tcp", dbAddr, socketTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.conn = conn
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// closeConn 关闭并置空复用连接。
|
||||
func (c *Client) closeConn() {
|
||||
if c.conn != nil {
|
||||
_ = c.conn.Close()
|
||||
c.conn = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) flush() {
|
||||
c.mu.Lock()
|
||||
if c.stopped {
|
||||
c.mu.Unlock()
|
||||
return
|
||||
}
|
||||
c.stopped = true
|
||||
c.mu.Unlock()
|
||||
|
||||
close(c.stop)
|
||||
c.wg.Wait()
|
||||
c.closeConn()
|
||||
}
|
||||
Reference in new issue
Block a user