perf: refactor

internal/apps/risk_control/logics.go

@@ -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
+var (
-
+	logWriterMu sync.RWMutex
-const (
+	logWriter   *batchwriter.Writer[*analytics.UserAccessLog]
-	defaultQueueSize = 10000
-	maxBatchSize     = 1000
-	flushInterval    = 1 * time.Second
 )
 
-// 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)
+	logWriterMu.Lock()
-	go startBatchWorker(context.WithoutCancel(ctx))
+	defer logWriterMu.Unlock()
-}
+	if logWriter != nil {
-
-// 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 {
 		return
 	}
 
-	select {
+	cfg := batchwriter.DefaultConfig()
-	case logChan <- logItem:
+	writer, err := batchwriter.New[*analytics.UserAccessLog](cfg, func(ctx context.Context, items []*analytics.UserAccessLog) error {
-	default:
+		rows := make([]analytics.UserAccessLog, 0, len(items))
-		logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", logItem.Path)
+		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) {
+func currentLogWriter() *batchwriter.Writer[*analytics.UserAccessLog] {
-	ticker := time.NewTicker(flushInterval)
+	logWriterMu.RLock()
-	defer ticker.Stop()
+	defer logWriterMu.RUnlock()
-
+	return logWriter
-	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()
-		}
-	}
-}

internal/apps/upload/filesrv/file_server.go

@@ -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

@@ -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
+	tempDir := t.TempDir()
-	defer func() { _ = os.RemoveAll("uploads") }()
+	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,
+		ID:         8001,
-		UserID:        user.ID,
+		UserID:     user.ID,
-		FileName:      "avatar.png",
+		FileName:   "avatar.png",
-		FilePath:      "uploads/avatar.png",
+		FilePath:   "avatar.png",
-		FileSize:      5,
+		FileSize:   5,
-		MimeType:      "image/png",
+		MimeType:   "image/png",
-		Extension:     "png",
+		Extension:  "png",
-		Type:          "avatar",
+		Type:       "avatar",
-		Status:        model.UploadStatusUsed,
+		Status:     model.UploadStatusUsed,
-		AccessMode:    1,
+		AccessMode: 1,
 	}
 	attachmentFile := model.Upload{
-		ID:            8002,
+		ID:         8002,
-		UserID:        user.ID,
+		UserID:     user.ID,
-		FileName:      "doc.pdf",
+		FileName:   "doc.pdf",
-		FilePath:      "uploads/doc.pdf",
+		FilePath:   "doc.pdf",
-		FileSize:      5,
+		FileSize:   5,
-		MimeType:      "application/pdf",
+		MimeType:   "application/pdf",
-		Extension:     "pdf",
+		Extension:  "pdf",
-		Type:          "attachment",
+		Type:       "attachment",
-		Status:        model.UploadStatusUsed,
+		Status:     model.UploadStatusUsed,
-		AccessMode:    1,
+		AccessMode: 1,
 	}
 
-	_ = os.MkdirAll("uploads", 0755)
+	if err := os.WriteFile(filepath.Join(tempDir, "avatar.png"), []byte("image"), 0644); err != nil {
-	_ = os.WriteFile(avatarFile.FilePath, []byte("image"), 0644)
+		t.Fatalf("failed to write avatar file: %v", err)
-	_ = os.WriteFile(attachmentFile.FilePath, []byte("bytes"), 0644)
+	}
+	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 := filepath.Join(tempDir, "test_image.png")
-	filePath := "uploads/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,
+		ID:         3001,
-		UserID:        user.ID,
+		UserID:     user.ID,
-		FileName:      "test_image.png",
+		FileName:   "test_image.png",
-		FilePath:      filePath,
+		FilePath:   "test_image.png",
-		FileSize:      int64(pngBuf.Len()),
+		FileSize:   int64(pngBuf.Len()),
-		MimeType:      "image/png",
+		MimeType:   "image/png",
-		Extension:     "png",
+		Extension:  "png",
-		Type:          "avatar", // Whitelisted by default
+		Type:       "avatar", // Whitelisted by default
-		Status:        model.UploadStatusUsed,
+		Status:     model.UploadStatusUsed,
-		AccessMode:    1,
+		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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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")