tcjlj/OpenFlare
1
0
Fork 0
mirror of https://github.com/Rain-kl/OpenFlare.git synced 2026-09-29 05:56:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: refactor

Browse Source
This commit is contained in:
ryan
2026-06-20 09:46:24 +08:00
parent 080be1e03a
commit cac6e88bc0
8 changed files with 244 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/apps/risk_control/logics.go
+84 -65
View File
@@ -5,88 +5,107 @@ package risk_control
import (
"context"
"time"
"sync"
"github.com/Rain-kl/Wavelet/internal/config"
"github.com/Rain-kl/Wavelet/internal/db/batchwriter"
"github.com/Rain-kl/Wavelet/internal/lifecycle"
"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 logChan chan *analytics.UserAccessLog
const (
defaultQueueSize = 10000
maxBatchSize = 1000
flushInterval = 1 * time.Second
var (
logWriterMu sync.RWMutex
logWriter *batchwriter.Writer[*analytics.UserAccessLog]
)
// InitLogWriter 初始化日志写入通道和后台写入协程
// InitLogWriter initializes the ClickHouse access-log batch writer.
func InitLogWriter(ctx context.Context) {
if !config.Config.ClickHouse.Enabled {
return
}
logChan = make(chan *analytics.UserAccessLog, defaultQueueSize)
go startBatchWorker(context.WithoutCancel(ctx))
}
// IsBufferFull 检查当前本地缓冲队列是否已满
// 如果没有启用 ClickHouse，默认返回 false，不触发限流
func IsBufferFull() bool {
if !config.Config.ClickHouse.Enabled || logChan == nil {
return false
}
return len(logChan) >= cap(logChan)
}
// QueueAccessLog 异步非阻塞地将日志推入缓冲队列
func QueueAccessLog(logItem *analytics.UserAccessLog) {
if !config.Config.ClickHouse.Enabled || logChan == nil {
logWriterMu.Lock()
defer logWriterMu.Unlock()
if logWriter != nil {
return
}
select {
case logChan <- logItem:
default:
logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", logItem.Path)
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
lifecycle.OnShutdown("risk_control_log_writer", StopLogWriter)
}
// 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 startBatchWorker(ctx context.Context) {
ticker := time.NewTicker(flushInterval)
defer ticker.Stop()
var batch []*analytics.UserAccessLog
flush := func() {
if len(batch) == 0 {
return
}
items := make([]analytics.UserAccessLog, len(batch))
for i, item := range batch {
items[i] = *item
}
if err := analyticsrepo.BatchInsert(ctx, items); err != nil {
logger.ErrorF(ctx, "[RiskControl] Send ClickHouse batch failed: %v", err)
}
batch = nil
}
for {
select {
case item, ok := <-logChan:
if !ok {
flush()
return
}
batch = append(batch, item)
if len(batch) >= maxBatchSize {
flush()
}
case <-ticker.C:
flush()
}
}
}
func currentLogWriter() *batchwriter.Writer[*analytics.UserAccessLog] {
logWriterMu.RLock()
defer logWriterMu.RUnlock()
return logWriter
}
internal/apps/upload/filesrv/file_server.go
+6 -6
View File
@@ -25,7 +25,7 @@ import (
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/diskcache"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/pkg/logger"
"github.com/gin-gonic/gin"
"golang.org/x/sync/singleflight"
@@ -66,14 +66,14 @@ func ServeFileByID(c *gin.Context) {
upload, err := GetUploadRecordByID(c)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if _, ok := err.(*strconv.NumError); ok {
c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid upload ID"})
response.AbortBadRequest(c, "无效的上传ID")
return
}
c.AbortWithStatus(http.StatusInternalServerError)
response.AbortInternal(c, "服务器内部错误")
return
}
@@ -263,7 +263,7 @@ func ImageCompressionCacheKey(upload *model.Upload, quality string) string {
func serveOriginal(c *gin.Context, upload *model.Upload) {
obj, err := uploadstorage.OpenStoredObject(c.Request.Context(), upload)
if err != nil {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件未找到")
return
}
defer func() { _ = obj.Body.Close() }()
@@ -313,4 +313,4 @@ func CheckFileAccessPermission(c *gin.Context, upload *model.Upload) error {
}
}
return nil
}
}
internal/apps/upload/filesrv/file_server_test.go
+71 -41
View File
@@ -6,6 +6,7 @@ package filesrv
import (
"bytes"
"context"
"encoding/json"
"image"
"image/color"
@@ -13,6 +14,7 @@ import (
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/upload/cache"
@@ -20,12 +22,16 @@ import (
"github.com/Rain-kl/Wavelet/internal/apps/upload/util"
"github.com/Rain-kl/Wavelet/internal/common"
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/diskcache"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/storage"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
func TestServeFileByIDAccessControl(t *testing.T) {
@@ -33,8 +39,8 @@ func TestServeFileByIDAccessControl(t *testing.T) {
defer cleanup()
cache.ResetAccessCaches()
// Ensure uploads dir is cleaned up
defer func() { _ = os.RemoveAll("uploads") }()
tempDir := t.TempDir()
configureLocalStorageRoot(t, dbConn, tempDir)
// Create a user in DB
user := model.User{
@@ -60,33 +66,36 @@ func TestServeFileByIDAccessControl(t *testing.T) {
// Create two files: one in whitelist (avatar), one not in whitelist (attachment)
avatarFile := model.Upload{
ID: 8001,
UserID: user.ID,
FileName: "avatar.png",
FilePath: "uploads/avatar.png",
FileSize: 5,
MimeType: "image/png",
Extension: "png",
Type: "avatar",
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 8001,
UserID: user.ID,
FileName: "avatar.png",
FilePath: "avatar.png",
FileSize: 5,
MimeType: "image/png",
Extension: "png",
Type: "avatar",
Status: model.UploadStatusUsed,
AccessMode: 1,
}
attachmentFile := model.Upload{
ID: 8002,
UserID: user.ID,
FileName: "doc.pdf",
FilePath: "uploads/doc.pdf",
FileSize: 5,
MimeType: "application/pdf",
Extension: "pdf",
Type: "attachment",
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 8002,
UserID: user.ID,
FileName: "doc.pdf",
FilePath: "doc.pdf",
FileSize: 5,
MimeType: "application/pdf",
Extension: "pdf",
Type: "attachment",
Status: model.UploadStatusUsed,
AccessMode: 1,
}
_ = os.MkdirAll("uploads", 0755)
_ = os.WriteFile(avatarFile.FilePath, []byte("image"), 0644)
_ = os.WriteFile(attachmentFile.FilePath, []byte("bytes"), 0644)
if err := os.WriteFile(filepath.Join(tempDir, "avatar.png"), []byte("image"), 0644); err != nil {
t.Fatalf("failed to write avatar file: %v", err)
}
if err := os.WriteFile(filepath.Join(tempDir, "doc.pdf"), []byte("bytes"), 0644); err != nil {
t.Fatalf("failed to write attachment file: %v", err)
}
dbConn.Create(&avatarFile)
dbConn.Create(&attachmentFile)
@@ -159,15 +168,14 @@ func TestImageCompression(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
tempDir := t.TempDir()
configureLocalStorageRoot(t, dbConn, tempDir)
cache := diskcache.GetGlobalCache()
if err := cache.Clear(); err != nil {
t.Fatalf("failed to clear disk cache before test: %v", err)
}
// Ensure uploads dir is cleaned up
defer func() {
_ = os.RemoveAll("uploads")
}()
defer func() {
if err := cache.Clear(); err != nil {
t.Errorf("failed to clear disk cache after test: %v", err)
@@ -190,24 +198,23 @@ func TestImageCompression(t *testing.T) {
t.Fatalf("failed to encode test png: %v", err)
}
_ = os.MkdirAll("uploads", 0755)
filePath := "uploads/test_image.png"
filePath := filepath.Join(tempDir, "test_image.png")
if err := os.WriteFile(filePath, pngBuf.Bytes(), 0644); err != nil {
t.Fatalf("failed to write test png: %v", err)
}
// Save upload record to DB
uploadRecord := model.Upload{
ID: 3001,
UserID: user.ID,
FileName: "test_image.png",
FilePath: filePath,
FileSize: int64(pngBuf.Len()),
MimeType: "image/png",
Extension: "png",
Type: "avatar", // Whitelisted by default
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 3001,
UserID: user.ID,
FileName: "test_image.png",
FilePath: "test_image.png",
FileSize: int64(pngBuf.Len()),
MimeType: "image/png",
Extension: "png",
Type: "avatar", // Whitelisted by default
Status: model.UploadStatusUsed,
AccessMode: 1,
}
dbConn.Create(&uploadRecord)
@@ -344,3 +351,26 @@ func TestNormalizeImageQuality(t *testing.T) {
})
}
}
func configureLocalStorageRoot(t *testing.T, dbConn *gorm.DB, tempDir string) {
var sc model.SystemConfig
if err := dbConn.Where("key = ?", model.ConfigKeyStorageConfig).First(&sc).Error; err != nil {
t.Fatalf("failed to find storage config: %v", err)
}
var cfg storage.Config
if err := json.Unmarshal([]byte(sc.Value), &cfg); err != nil {
t.Fatalf("failed to unmarshal storage config: %v", err)
}
cfg.Local.Root = tempDir
newVal, err := json.Marshal(cfg)
if err != nil {
t.Fatalf("failed to marshal storage config: %v", err)
}
sc.Value = string(newVal)
if err := dbConn.Save(&sc).Error; err != nil {
t.Fatalf("failed to save storage config: %v", err)
}
_ = db.HSetJSON(context.Background(), repository.SystemConfigRedisHashKey, sc.Key, &sc)
repository.ResetSystemConfigRAMCacheForTest()
storage.ResetCache()
}
internal/apps/upload/handler/file_management.go
+6 -6
View File
@@ -112,7 +112,7 @@ func DeleteFile(c *gin.Context) {
if _, err := softDeleteUpload(ctx, uploadID); err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
response.AbortBadRequest(c, shared.ErrDeleteFileFailed)
@@ -233,11 +233,11 @@ func DeleteMyFile(c *gin.Context) {
if _, err := softDeleteOwnedUpload(ctx, currUser.ID, uploadID); err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if err == ingest.ErrForbidden {
c.AbortWithStatus(http.StatusForbidden)
response.AbortForbidden(c, "无权操作")
return
}
response.AbortBadRequest(c, shared.ErrDeleteFileFailed)
@@ -287,11 +287,11 @@ func UpdateMyFile(c *gin.Context) {
upload, err := updateOwnedUpload(ctx, currUser.ID, uploadID, updateMyUploadInput(req))
if err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if err == ingest.ErrForbidden {
c.AbortWithStatus(http.StatusForbidden)
response.AbortForbidden(c, "无权操作")
return
}
response.AbortBadRequest(c, "更新文件记录失败")
@@ -299,4 +299,4 @@ func UpdateMyFile(c *gin.Context) {
}
c.JSON(http.StatusOK, response.OK(upload))
}
}
internal/apps/upload/handler/routers.go
+1 -1
View File
@@ -170,7 +170,7 @@ func DownloadFile(c *gin.Context) {
upload, err := filesrv.GetUploadRecordByID(c)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if _, ok := err.(*strconv.NumError); ok {
internal/bootstrap/bootstrap.go
+7 -1
View File
@@ -12,6 +12,7 @@ import (
admin_push "github.com/Rain-kl/Wavelet/internal/apps/admin/push"
"github.com/Rain-kl/Wavelet/internal/apps/admin/push/custom_events"
"github.com/Rain-kl/Wavelet/internal/apps/risk_control"
"github.com/Rain-kl/Wavelet/internal/lifecycle"
taskhandlers "github.com/Rain-kl/Wavelet/internal/task/handlers"
"github.com/Rain-kl/Wavelet/pkg/logger"
)
@@ -87,7 +88,12 @@ func Init(ctx context.Context, opts Options) {
})
}
// Stop stops all batch writers and background resources.
func Stop(ctx context.Context) {
lifecycle.Stop(ctx)
}
// ResetInitRuntimeOnceForTest clears initRuntimeOnce so Init can run again in unit tests.
func ResetInitRuntimeOnceForTest() {
initRuntimeOnce = sync.Once{}
}
}
internal/lifecycle/lifecycle.go
+66
View File
@@ -0,0 +1,66 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package lifecycle manages global application and business component shutdown hooks.
package lifecycle
import (
"context"
"log"
"sync"
)
// ShutdownFunc defines the signature for a graceful shutdown callback.
type ShutdownFunc func(ctx context.Context) error
type hook struct {
name string
fn ShutdownFunc
}
var (
hooks []hook
mu sync.Mutex
)
// OnShutdown registers a callback to be run during graceful shutdown.
func OnShutdown(name string, fn ShutdownFunc) {
mu.Lock()
defer mu.Unlock()
hooks = append(hooks, hook{name: name, fn: fn})
}
// Stop executes all registered shutdown hooks concurrently and waits for completion or context timeout.
func Stop(ctx context.Context) {
mu.Lock()
localHooks := make([]hook, len(hooks))
copy(localHooks, hooks)
mu.Unlock()
var wg sync.WaitGroup
for _, h := range localHooks {
wg.Add(1)
go func(name string, fn ShutdownFunc) {
defer wg.Done()
log.Printf("[Lifecycle] stopping %s...\n", name)
if err := fn(ctx); err != nil {
log.Printf("[Lifecycle] stop %s failed: %v\n", name, err)
} else {
log.Printf("[Lifecycle] %s stopped successfully\n", name)
}
}(h.name, h.fn)
}
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
log.Println("[Lifecycle] all services stopped gracefully")
case <-ctx.Done():
log.Printf("[Lifecycle] shutdown timed out: %v\n", ctx.Err())
}
}
internal/router/router.go
+3
View File
@@ -16,6 +16,7 @@ import (
"time"
"github.com/Rain-kl/Wavelet/internal/apps/risk_control"
"github.com/Rain-kl/Wavelet/internal/bootstrap"
router_root "github.com/Rain-kl/Wavelet/internal/router/root"
v1 "github.com/Rain-kl/Wavelet/internal/router/v1"
@@ -99,9 +100,11 @@ func Serve() {
if err := srv.Shutdown(shutdownCtx); err != nil {
log.Printf("[API] server forced to shutdown: %v\n", err)
bootstrap.Stop(shutdownCtx)
cancel()
os.Exit(1)
}
bootstrap.Stop(shutdownCtx)
cancel()
log.Println("[API] server exited")