optimize

quota.go

@@ -0,0 +1,112 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"sync"
+)
+
+// -----------------------------------------------------------------------------
+// 配额 (总目录字节数限制, 默认 10GB, 0=不限)
+// -----------------------------------------------------------------------------
+
+var (
+	quotaMu   sync.Mutex
+	quotaUsed int64 // 当前目录已用字节数
+)
+
+// initQuota 启动时扫描目录, 初始化配额计数
+func initQuota() {
+	if *quota <= 0 {
+		log.Printf("[quota] disabled")
+		return
+	}
+	rescanQuota()
+	log.Printf("[quota] %s / %s", formatSize(quotaUsed), formatSize(*quota))
+}
+
+// rescanQuota 重新扫描目录并重置计数. 调用频率低 (启动 + 清理器每次跑完), 用来修正漂移.
+func rescanQuota() {
+	quotaMu.Lock()
+	defer quotaMu.Unlock()
+	var total int64
+	_ = filepath.Walk(*uploadDir, func(_ string, info os.FileInfo, err error) error {
+		if err != nil || info.IsDir() {
+			return nil
+		}
+		total += info.Size()
+		return nil
+	})
+	quotaUsed = total
+}
+
+// quotaSnapshot 原子返回 used/cap
+func quotaSnapshot() (used, cap_ int64) {
+	quotaMu.Lock()
+	defer quotaMu.Unlock()
+	return quotaUsed, *quota
+}
+
+// quotaRemaining 返回剩余可上传字节数. quota<=0 时返回 math.MaxInt64.
+func quotaRemaining() int64 {
+	quotaMu.Lock()
+	defer quotaMu.Unlock()
+	if *quota <= 0 {
+		return 1 << 62
+	}
+	return *quota - quotaUsed
+}
+
+// quotaReserve 原子地预留 n 字节配额: 超限返回 false, 否则扣减并返回 true.
+func quotaReserve(n int64) bool {
+	quotaMu.Lock()
+	defer quotaMu.Unlock()
+	if *quota > 0 && quotaUsed+n > *quota {
+		return false
+	}
+	quotaUsed += n
+	return true
+}
+
+func quotaSub(n int64) {
+	if n <= 0 {
+		return
+	}
+	quotaMu.Lock()
+	quotaUsed -= n
+	if quotaUsed < 0 { // 漂移防护
+		quotaUsed = 0
+	}
+	quotaMu.Unlock()
+}
+
+// formatSize 字节 → 人类可读
+func formatSize(b int64) string {
+	const (
+		KB = 1 << 10
+		MB = 1 << 20
+		GB = 1 << 30
+		TB = 1 << 40
+	)
+	switch {
+	case b >= TB:
+		return fmt.Sprintf("%.2f TB", float64(b)/TB)
+	case b >= GB:
+		return fmt.Sprintf("%.2f GB", float64(b)/GB)
+	case b >= MB:
+		return fmt.Sprintf("%.2f MB", float64(b)/MB)
+	case b >= KB:
+		return fmt.Sprintf("%.2f KB", float64(b)/KB)
+	default:
+		return fmt.Sprintf("%d B", b)
+	}
+}
+
+// quotaUsedSnapshot 不持锁读, 仅用于 cleaner 自身的日志 (允许轻微漂移)
+func quotaUsedSnapshot() int64 {
+	quotaMu.Lock()
+	defer quotaMu.Unlock()
+	return quotaUsed
+}