blackbird/waf/internal/logclient/logclient.go

// Package logclient 提供异步批量日志推送客户端。
//
// 通过 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 载荷）
package logclient

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"time"
)

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

// SetLevel 设置最低日志级别（debug|info|warn|error，大小写不敏感）。
func SetLevel(level string) {
	if n, ok := levelMap[level]; ok {
		minLevel = n
	}
}

// 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 设定的最低级别的日志将被静默丢弃。
func Push(service, level, message string, extra map[string]any) {
	if global == nil {
		return
	}
	n, ok := levelMap[level]
	if !ok || n < minLevel {
		return
	}
	global.push(Entry{
		Service: service,
		Level:   level,
		Message: message,
		Extra:   extra,
	})
}

// Debug 推送 DEBUG 级别日志。
func Debug(service, message string, extra map[string]any) {
	Push(service, "DEBUG", message, extra)
}

// Info 推送 INFO 级别日志。
func Info(service, message string, extra map[string]any) {
	Push(service, "INFO", message, extra)
}

// Warning 推送 WARNING 级别日志。
func Warning(service, message string, extra map[string]any) {
	Push(service, "WARNING", message, extra)
}

// Error 推送 ERROR 级别日志。
func Error(service, message string, extra map[string]any) {
	Push(service, "ERROR", message, extra)
}

// Flush 同步 flush 队列中剩余日志（用于进程退出前）。
func Flush() {
	if global == nil {
		return
	}
	global.flush()
}

func (c *Client) push(e Entry) {
	select {
	case c.ch <- e:
	default:
		// 队列满，静默丢弃（fail-safe）
	}
}

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