113 lines
2.4 KiB
Go
113 lines
2.4 KiB
Go
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
|
|
}
|