mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-03 23:06:36 +08:00
ryan
2a0ebd16fa
- 修复 ClickHouse Batch Writer 与 RiskControl LogWriter 优雅停机，确保退出前队列数据正确刷盘。 - 修复 OpenFlare 系统配置参数全局变量并发读写的 Data Race 冲突，在读取配置时引入读锁保护。 - 修复 upload 模块的 API 错误响应格式，使用 response.Abort* 代替原始的 c.AbortWithStatus 与 c.JSON，保证全局信封格式统一。 - 修复 pkg/utils/network 的 isPrivateIPv4 以使用标准库 net.IP.IsPrivate() 校验，修复 format 的 Bytes2Size 边界，修改大小单位因子变量为只读常量。 - 重构 upload 文件服务与路由器测试，使用 t.TempDir() 代替硬编码的 uploads 相对路径写入和删除，解决测试目录文件污染问题。
109 lines
2.8 KiB
Go
109 lines
2.8 KiB
Go
// Copyright 2026 Arctel.net
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
package risk_control
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/Rain-kl/Wavelet/internal/config"
|
|
"github.com/Rain-kl/Wavelet/internal/db/batchwriter"
|
|
"github.com/Rain-kl/Wavelet/internal/model/analytics"
|
|
analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics"
|
|
"github.com/Rain-kl/Wavelet/pkg/logger"
|
|
)
|
|
|
|
var (
|
|
logWriterMu sync.RWMutex
|
|
logWriter *batchwriter.Writer[*analytics.UserAccessLog]
|
|
)
|
|
|
|
// InitLogWriter initializes the ClickHouse access-log batch writer.
|
|
func InitLogWriter(ctx context.Context) {
|
|
if !config.Config.ClickHouse.Enabled {
|
|
return
|
|
}
|
|
|
|
logWriterMu.Lock()
|
|
defer logWriterMu.Unlock()
|
|
if logWriter != nil {
|
|
return
|
|
}
|
|
|
|
cfg := batchwriter.DefaultConfig()
|
|
writer, err := batchwriter.New[*analytics.UserAccessLog](cfg, func(ctx context.Context, items []*analytics.UserAccessLog) error {
|
|
rows := make([]analytics.UserAccessLog, 0, len(items))
|
|
for _, item := range items {
|
|
if item == nil {
|
|
continue
|
|
}
|
|
rows = append(rows, *item)
|
|
}
|
|
return analyticsrepo.BatchInsert(ctx, rows)
|
|
},
|
|
batchwriter.WithDropHandler[*analytics.UserAccessLog](func(item *analytics.UserAccessLog) {
|
|
path := ""
|
|
if item != nil {
|
|
path = item.Path
|
|
}
|
|
logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", path)
|
|
}),
|
|
batchwriter.WithFlushErrorHandler[*analytics.UserAccessLog](func(ctx context.Context, batchSize int, err error) {
|
|
logger.ErrorF(ctx, "[RiskControl] Send ClickHouse batch failed (batch=%d): %v", batchSize, err)
|
|
}),
|
|
)
|
|
if err != nil {
|
|
logger.ErrorF(ctx, "[RiskControl] init log writer failed: %v", err)
|
|
return
|
|
}
|
|
|
|
writer.Start(ctx)
|
|
logWriter = writer
|
|
}
|
|
|
|
// StopLogWriter stops the ClickHouse access-log batch writer and drains pending logs.
|
|
func StopLogWriter(ctx context.Context) error {
|
|
writer := currentLogWriter()
|
|
if writer == nil {
|
|
return nil
|
|
}
|
|
return writer.Stop(ctx)
|
|
}
|
|
|
|
// IsBufferFull reports whether the access-log queue has no remaining capacity.
|
|
func IsBufferFull() bool {
|
|
writer := currentLogWriter()
|
|
if writer == nil {
|
|
return false
|
|
}
|
|
return writer.IsFull()
|
|
}
|
|
|
|
// QueueAccessLog enqueues an access log without blocking.
|
|
func QueueAccessLog(logItem *analytics.UserAccessLog) {
|
|
writer := currentLogWriter()
|
|
if writer == nil || logItem == nil {
|
|
return
|
|
}
|
|
writer.TryEnqueue(logItem)
|
|
}
|
|
|
|
// SetLogWriterForTest swaps the access-log writer for unit tests.
|
|
func SetLogWriterForTest(writer *batchwriter.Writer[*analytics.UserAccessLog]) func() {
|
|
logWriterMu.Lock()
|
|
previous := logWriter
|
|
logWriter = writer
|
|
logWriterMu.Unlock()
|
|
return func() {
|
|
logWriterMu.Lock()
|
|
logWriter = previous
|
|
logWriterMu.Unlock()
|
|
}
|
|
}
|
|
|
|
func currentLogWriter() *batchwriter.Writer[*analytics.UserAccessLog] {
|
|
logWriterMu.RLock()
|
|
defer logWriterMu.RUnlock()
|
|
return logWriter
|
|
} |