From ba90e9f452741e91020d41157fdf7dcbc3e94b00 Mon Sep 17 00:00:00 2001 From: vxtls <187420201+vxtls@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:31:44 -0400 Subject: [PATCH] feat(cache): add configurable HybridCache policies - Add auto, memory, and disk cache policies with global configuration and per-instance overrides. - Select the backing store in auto mode using the complete workload memory ceiling while preserving runtime disk spill. - Enforce memory ceilings independently of page-aligned allocations and prevent strict memory mode from falling back to disk. - Route unknown-size streams through HybridCache and support sequential writes without unnecessary file preallocation. - Bound downloader cache decisions by its concurrent working set and validate cache policy configuration. - Add policy, backing-store selection, stream override, cleanup, and downloader ceiling tests. --- internal/bootstrap/config.go | 2 + internal/cache/policy.go | 62 +++++++++ internal/cache/policy_test.go | 82 ++++++++++++ internal/conf/config.go | 57 ++++---- internal/conf/config_test.go | 21 +++ internal/conf/var.go | 7 +- internal/hybrid_cache/hybrid_cache.go | 179 +++++++++++++++++++++---- internal/hybrid_cache/policy_test.go | 184 ++++++++++++++++++++++++++ internal/mem/utils.go | 18 ++- internal/model/obj.go | 3 + internal/net/request.go | 25 +++- internal/net/request_test.go | 24 ++++ internal/stream/stream.go | 87 ++++++------ internal/stream/stream_test.go | 94 +++++++++++++ internal/stream/util.go | 11 +- 15 files changed, 753 insertions(+), 103 deletions(-) create mode 100644 internal/cache/policy.go create mode 100644 internal/cache/policy_test.go create mode 100644 internal/conf/config_test.go create mode 100644 internal/hybrid_cache/policy_test.go diff --git a/internal/bootstrap/config.go b/internal/bootstrap/config.go index 8304468080..9a5c0d7fbd 100644 --- a/internal/bootstrap/config.go +++ b/internal/bootstrap/config.go @@ -96,6 +96,8 @@ func InitConfig() { if !conf.Conf.Force { confFromEnv() } + conf.CachePolicy = conf.Conf.CachePolicy + log.Infof("cache policy: %s", conf.CachePolicy) if conf.Conf.MaxConcurrency > math.MaxInt32 { net.DefaultConcurrencyLimit = &net.ConcurrencyLimit{Limit: math.MaxInt32} diff --git a/internal/cache/policy.go b/internal/cache/policy.go new file mode 100644 index 0000000000..5a10220e53 --- /dev/null +++ b/internal/cache/policy.go @@ -0,0 +1,62 @@ +package cache + +import ( + "fmt" + "strings" +) + +type Policy string + +const ( + // PolicyInherit is only used by per-instance overrides. It is not a + // user-selectable cache policy. + PolicyInherit Policy = "" + PolicyAuto Policy = "auto" + PolicyMemory Policy = "memory" + PolicyDisk Policy = "disk" +) + +func ParsePolicy(value string) (Policy, error) { + policy := Policy(strings.ToLower(strings.TrimSpace(value))) + if policy == PolicyInherit { + return PolicyAuto, nil + } + if !policy.IsConcrete() { + return PolicyInherit, fmt.Errorf("invalid cache policy %q: expected auto, memory, or disk", value) + } + return policy, nil +} + +func ResolvePolicy(override, fallback Policy) (Policy, error) { + policy := override + if policy == PolicyInherit { + policy = fallback + } + if !policy.IsConcrete() { + return PolicyInherit, fmt.Errorf("invalid cache policy %q: expected auto, memory, or disk", policy) + } + return policy, nil +} + +func (p Policy) IsConcrete() bool { + return p == PolicyAuto || p == PolicyMemory || p == PolicyDisk +} + +func (p Policy) MarshalText() ([]byte, error) { + if p == PolicyInherit { + return []byte{}, nil + } + if !p.IsConcrete() { + return nil, fmt.Errorf("invalid cache policy %q", p) + } + return []byte(p), nil +} + +func (p *Policy) UnmarshalText(text []byte) error { + policy, err := ParsePolicy(string(text)) + if err != nil { + return err + } + *p = policy + return nil +} diff --git a/internal/cache/policy_test.go b/internal/cache/policy_test.go new file mode 100644 index 0000000000..c982a47abe --- /dev/null +++ b/internal/cache/policy_test.go @@ -0,0 +1,82 @@ +package cache + +import ( + "encoding/json" + "testing" + + "github.com/caarlos0/env/v9" +) + +func TestParsePolicy(t *testing.T) { + tests := []struct { + input string + want Policy + }{ + {"", PolicyAuto}, + {"auto", PolicyAuto}, + {" MEMORY ", PolicyMemory}, + {"Disk", PolicyDisk}, + } + for _, tt := range tests { + got, err := ParsePolicy(tt.input) + if err != nil { + t.Fatalf("ParsePolicy(%q) error = %v", tt.input, err) + } + if got != tt.want { + t.Errorf("ParsePolicy(%q) = %q, want %q", tt.input, got, tt.want) + } + } + if _, err := ParsePolicy("hybrid"); err == nil { + t.Fatal("ParsePolicy() expected an error for an invalid policy") + } +} + +func TestPolicyEnvironment(t *testing.T) { + t.Setenv("OPENLIST_TEST_CACHE_POLICY", " Disk ") + var cfg struct { + Policy Policy `env:"CACHE_POLICY"` + } + if err := env.ParseWithOptions(&cfg, env.Options{Prefix: "OPENLIST_TEST_"}); err != nil { + t.Fatalf("env.ParseWithOptions() error = %v", err) + } + if cfg.Policy != PolicyDisk { + t.Fatalf("environment policy = %q, want disk", cfg.Policy) + } +} + +func TestResolvePolicy(t *testing.T) { + got, err := ResolvePolicy(PolicyInherit, PolicyDisk) + if err != nil || got != PolicyDisk { + t.Fatalf("ResolvePolicy(inherit, disk) = %q, %v", got, err) + } + got, err = ResolvePolicy(PolicyMemory, PolicyDisk) + if err != nil || got != PolicyMemory { + t.Fatalf("ResolvePolicy(memory, disk) = %q, %v", got, err) + } + if _, err := ResolvePolicy(PolicyInherit, Policy("invalid")); err == nil { + t.Fatal("ResolvePolicy() expected an error for an invalid fallback") + } +} + +func TestPolicyJSON(t *testing.T) { + type config struct { + Policy Policy `json:"cache_policy"` + } + var cfg config + if err := json.Unmarshal([]byte(`{"cache_policy":" MEMORY "}`), &cfg); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + if cfg.Policy != PolicyMemory { + t.Fatalf("json.Unmarshal() policy = %q, want memory", cfg.Policy) + } + b, err := json.Marshal(cfg) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + if string(b) != `{"cache_policy":"memory"}` { + t.Fatalf("json.Marshal() = %s", b) + } + if err := json.Unmarshal([]byte(`{"cache_policy":"invalid"}`), &cfg); err == nil { + t.Fatal("json.Unmarshal() expected an error for an invalid policy") + } +} diff --git a/internal/conf/config.go b/internal/conf/config.go index f8423d5908..5e72ce46d8 100644 --- a/internal/conf/config.go +++ b/internal/conf/config.go @@ -3,6 +3,7 @@ package conf import ( "path/filepath" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/pkg/utils/random" ) @@ -112,33 +113,34 @@ type MCP struct { } type Config struct { - Force bool `json:"force" env:"FORCE"` - SiteURL string `json:"site_url" env:"SITE_URL"` - Cdn string `json:"cdn" env:"CDN"` - JwtSecret string `json:"jwt_secret" env:"JWT_SECRET"` - TokenExpiresIn int `json:"token_expires_in" env:"TOKEN_EXPIRES_IN"` - Database Database `json:"database" envPrefix:"DB_"` - Meilisearch Meilisearch `json:"meilisearch" envPrefix:"MEILISEARCH_"` - Scheme Scheme `json:"scheme"` - TempDir string `json:"temp_dir" env:"TEMP_DIR"` - BleveDir string `json:"bleve_dir" env:"BLEVE_DIR"` - DistDir string `json:"dist_dir"` - Log LogConfig `json:"log" envPrefix:"LOG_"` - DelayedStart int `json:"delayed_start" env:"DELAYED_START"` - AutoMemoryLimit int `json:"auto_memory_limit" env:"AUTO_MEMORY_LIMIT"` - MinFreeMemory int `json:"min_free_memory" env:"MIN_FREE_MEMORY"` - MaxBlockLimit int `json:"max_block_limit" env:"MAX_BLOCK_LIMIT"` - MaxConnections int `json:"max_connections" env:"MAX_CONNECTIONS"` - MaxConcurrency int `json:"max_concurrency" env:"MAX_CONCURRENCY"` - TlsInsecureSkipVerify bool `json:"tls_insecure_skip_verify" env:"TLS_INSECURE_SKIP_VERIFY"` - Tasks TasksConfig `json:"tasks" envPrefix:"TASKS_"` - Cors Cors `json:"cors" envPrefix:"CORS_"` - S3 S3 `json:"s3" envPrefix:"S3_"` - FTP FTP `json:"ftp" envPrefix:"FTP_"` - SFTP SFTP `json:"sftp" envPrefix:"SFTP_"` - MCP MCP `json:"mcp" envPrefix:"MCP_"` - LastLaunchedVersion string `json:"last_launched_version"` - ProxyAddress string `json:"proxy_address" env:"PROXY_ADDRESS"` + Force bool `json:"force" env:"FORCE"` + SiteURL string `json:"site_url" env:"SITE_URL"` + Cdn string `json:"cdn" env:"CDN"` + JwtSecret string `json:"jwt_secret" env:"JWT_SECRET"` + TokenExpiresIn int `json:"token_expires_in" env:"TOKEN_EXPIRES_IN"` + Database Database `json:"database" envPrefix:"DB_"` + Meilisearch Meilisearch `json:"meilisearch" envPrefix:"MEILISEARCH_"` + Scheme Scheme `json:"scheme"` + TempDir string `json:"temp_dir" env:"TEMP_DIR"` + BleveDir string `json:"bleve_dir" env:"BLEVE_DIR"` + DistDir string `json:"dist_dir"` + Log LogConfig `json:"log" envPrefix:"LOG_"` + DelayedStart int `json:"delayed_start" env:"DELAYED_START"` + CachePolicy cache.Policy `json:"cache_policy" env:"CACHE_POLICY"` + AutoMemoryLimit int `json:"auto_memory_limit" env:"AUTO_MEMORY_LIMIT"` + MinFreeMemory int `json:"min_free_memory" env:"MIN_FREE_MEMORY"` + MaxBlockLimit int `json:"max_block_limit" env:"MAX_BLOCK_LIMIT"` + MaxConnections int `json:"max_connections" env:"MAX_CONNECTIONS"` + MaxConcurrency int `json:"max_concurrency" env:"MAX_CONCURRENCY"` + TlsInsecureSkipVerify bool `json:"tls_insecure_skip_verify" env:"TLS_INSECURE_SKIP_VERIFY"` + Tasks TasksConfig `json:"tasks" envPrefix:"TASKS_"` + Cors Cors `json:"cors" envPrefix:"CORS_"` + S3 S3 `json:"s3" envPrefix:"S3_"` + FTP FTP `json:"ftp" envPrefix:"FTP_"` + SFTP SFTP `json:"sftp" envPrefix:"SFTP_"` + MCP MCP `json:"mcp" envPrefix:"MCP_"` + LastLaunchedVersion string `json:"last_launched_version"` + ProxyAddress string `json:"proxy_address" env:"PROXY_ADDRESS"` } func DefaultConfig(dataDir string) *Config { @@ -185,6 +187,7 @@ func DefaultConfig(dataDir string) *Config { }, }, }, + CachePolicy: cache.PolicyAuto, AutoMemoryLimit: 4, MaxConnections: 0, MaxConcurrency: 64, diff --git a/internal/conf/config_test.go b/internal/conf/config_test.go new file mode 100644 index 0000000000..231e653bd8 --- /dev/null +++ b/internal/conf/config_test.go @@ -0,0 +1,21 @@ +package conf + +import ( + "encoding/json" + "testing" + + "github.com/OpenListTeam/OpenList/v4/internal/cache" +) + +func TestDefaultCachePolicy(t *testing.T) { + cfg := DefaultConfig(t.TempDir()) + if got := cfg.CachePolicy; got != cache.PolicyAuto { + t.Fatalf("default cache policy = %q, want auto", got) + } + if err := json.Unmarshal([]byte(`{}`), cfg); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + if got := cfg.CachePolicy; got != cache.PolicyAuto { + t.Fatalf("cache policy after loading an old config = %q, want auto", got) + } +} diff --git a/internal/conf/var.go b/internal/conf/var.go index 6b25bcfb5b..57da166213 100644 --- a/internal/conf/var.go +++ b/internal/conf/var.go @@ -4,6 +4,8 @@ import ( "net/url" "regexp" "sync" + + "github.com/OpenListTeam/OpenList/v4/internal/cache" ) var ( @@ -25,10 +27,11 @@ var FilenameCharMap = make(map[string]string) var PrivacyReg []*regexp.Regexp var ( + CachePolicy cache.Policy = cache.PolicyAuto // 在HybridCache中使用[]byte缓存数据流的限制，内存为Go自动管理，直到GC AutoMemoryLimit uint64 = 4 * 1024 * 1024 - // 最小空闲内存，当内存不足时，HybridCache会回退到文件缓存。 - // 如果为0，HybridCache会使用文件缓存，不占用内存。 + // 最小空闲内存，当内存不足时，auto策略会回退到文件缓存。 + // 如果为0，auto策略会使用文件缓存，不占用内存。 MinFreeMemory uint64 = 16 * 1024 * 1024 // 限制HybridCache手动管理内存单次的扩容大小，超过该阈值将分多次扩容。 // MinFreeMemory大于0时，也限制 Downloader 的PartSize diff --git a/internal/hybrid_cache/hybrid_cache.go b/internal/hybrid_cache/hybrid_cache.go index c69147937e..027e750950 100644 --- a/internal/hybrid_cache/hybrid_cache.go +++ b/internal/hybrid_cache/hybrid_cache.go @@ -2,9 +2,11 @@ package hybrid_cache import ( "errors" + "fmt" "io" "runtime" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/internal/conf" "github.com/OpenListTeam/OpenList/v4/internal/mem" "github.com/OpenListTeam/OpenList/v4/pkg/buffer" @@ -13,12 +15,15 @@ import ( // 线程不安全，单线程使用，或者外部加锁保护 type HybridCache struct { - blockSize uint64 - memoryStore mem.LinearMemory - memoryOffset uint64 - backingStore BackingStore - backingOffset uint64 - cleanup runtime.Cleanup + blockSize uint64 + memoryStore mem.LinearMemory + memoryOffset uint64 + backingStore BackingStore + backingOffset uint64 + cleanup runtime.Cleanup + spillOnMemoryFailure bool + memoryBacking bool + memoryCeiling int64 } // HybridCache本身是一个大的Block，支持分块成多个小的Block @@ -27,6 +32,9 @@ type HybridCache struct { func (hc *HybridCache) AllocBlock(size uint64) (buffer.Block, error) { retry: if hc.backingStore != nil { + if hc.memoryBacking && hc.exceedsMemoryCeiling(hc.backingOffset, size) { + return nil, mem.ErrNotEnoughMemory + } if err := hc.backingStore.GrowTo(int64(hc.backingOffset + size)); err != nil { return nil, err } @@ -38,12 +46,21 @@ retry: ) return fs, nil } - all, err := hc.memoryStore.Reallocate(hc.memoryOffset + size) + var all []byte + var err error + if hc.exceedsMemoryCeiling(hc.memoryOffset, size) { + err = mem.ErrNotEnoughMemory + } else { + all, err = hc.memoryStore.Reallocate(hc.memoryOffset + size) + } if err == nil { start := hc.memoryOffset hc.memoryOffset += size return buffer.NewByteBlock(all[start : start+size]), nil } + if !hc.spillOnMemoryFailure { + return nil, err + } if err2 := hc.initFileCache(); err2 != nil { return nil, errors.Join(err, err2) } @@ -53,6 +70,9 @@ retry: func (hc *HybridCache) allocWriteAtSeeker(size uint64) (buffer.WriteAtSeeker, error) { retry: if hc.backingStore != nil { + if hc.memoryBacking && hc.exceedsMemoryCeiling(hc.backingOffset, size) { + return nil, mem.ErrNotEnoughMemory + } if err := hc.backingStore.GrowTo(int64(hc.backingOffset + size)); err != nil { return nil, err } @@ -60,12 +80,21 @@ retry: hc.backingOffset += size return io.NewOffsetWriter(hc.backingStore, int64(base)), nil } - all, err := hc.memoryStore.Reallocate(hc.memoryOffset + size) + var all []byte + var err error + if hc.exceedsMemoryCeiling(hc.memoryOffset, size) { + err = mem.ErrNotEnoughMemory + } else { + all, err = hc.memoryStore.Reallocate(hc.memoryOffset + size) + } if err == nil { start := hc.memoryOffset hc.memoryOffset += size return io.NewOffsetWriter(buffer.NewByteBlock(all[start:start+size]), 0), nil } + if !hc.spillOnMemoryFailure { + return nil, err + } if err2 := hc.initFileCache(); err2 != nil { return nil, errors.Join(err, err2) } @@ -95,8 +124,20 @@ func (hc *HybridCache) RewindOneBlock() { hc.RewindBySize(hc.blockSize) } +func (hc *HybridCache) exceedsMemoryCeiling(offset, size uint64) bool { + if hc.memoryCeiling < 0 { + return false + } + ceiling := uint64(hc.memoryCeiling) + return offset > ceiling || size > ceiling-offset +} + func (hc *HybridCache) initFileCache() error { - file, err := NewFileStore(int64(hc.blockSize)) + initialSize := hc.blockSize + if hc.memoryCeiling < 0 { + initialSize = 0 + } + file, err := NewFileStore(int64(initialSize)) if err != nil { return err } @@ -187,6 +228,33 @@ func (hc *HybridCache) WriteAt(p []byte, off int64) (n int, err error) { return n + nn, err } +// Write appends p to the cache. It is not safe for concurrent use. +func (hc *HybridCache) Write(p []byte) (n int, err error) { + for len(p) > 0 { + chunkSize := len(p) + if hc.blockSize > 0 && uint64(chunkSize) > hc.blockSize { + chunkSize = int(hc.blockSize) + } + w, allocErr := hc.allocWriteAtSeeker(uint64(chunkSize)) + if allocErr != nil { + return n, allocErr + } + nn, writeErr := w.Write(p[:chunkSize]) + n += nn + if nn < chunkSize { + hc.RewindBySize(uint64(chunkSize - nn)) + if writeErr == nil { + writeErr = io.ErrShortWrite + } + } + if writeErr != nil { + return n, writeErr + } + p = p[chunkSize:] + } + return n, nil +} + func (hc *HybridCache) CopyFromN(src io.Reader, n int64) (written int64, err error) { limit := n for limit > 0 { @@ -208,32 +276,87 @@ func (hc *HybridCache) CopyFromN(src io.Reader, n int64) (written int64, err err return written, nil } -// HybridCache 线程不安全，单线程使用，或者外部加锁保护 -func NewHybridCache(blockSize, maxMemorySize uint64) (hc *HybridCache, err error) { - if conf.MinFreeMemory > 0 { - // 策略1: Go自动内存管理 - if maxMemorySize <= conf.AutoMemoryLimit { - return &HybridCache{backingStore: &BufferStore{}, blockSize: blockSize}, nil +type memoryCheck func(uint64) error + +func selectPolicy(requested cache.Policy, memoryCeiling int64, check memoryCheck) (cache.Policy, error) { + if !requested.IsConcrete() { + return cache.PolicyInherit, fmt.Errorf("invalid cache policy %q", requested) + } + switch requested { + case cache.PolicyMemory, cache.PolicyDisk: + return requested, nil + case cache.PolicyAuto: + if memoryCeiling < 0 { + return cache.PolicyDisk, nil + } + if memoryCeiling == 0 { + return cache.PolicyMemory, nil + } + if err := check(uint64(memoryCeiling)); err != nil { + return cache.PolicyDisk, nil + } + return cache.PolicyMemory, nil + default: + panic("unreachable") + } +} + +// SelectPolicy resolves auto to a concrete memory or disk policy for a cache +// whose maximum simultaneous memory footprint is memoryCeiling. A negative +// ceiling means that the upper bound is unknown. +func SelectPolicy(requested cache.Policy, memoryCeiling int64) (cache.Policy, error) { + return selectPolicy(requested, memoryCeiling, mem.MemoryGrowCheck) +} + +// NewHybridCache creates a non-thread-safe cache using the requested policy. +func NewHybridCache(blockSize uint64, memoryCeiling int64, requested cache.Policy) (hc *HybridCache, err error) { + if memoryCeiling < 0 && blockSize == 0 { + return nil, fmt.Errorf("block size must be positive when memory ceiling is unknown") + } + if memoryCeiling > 0 { + blockSize = min(blockSize, uint64(memoryCeiling)) + if blockSize == 0 { + return nil, fmt.Errorf("block size must be positive for a non-empty cache") } + } - // 策略2: 手动内存管理 - if maxMemorySize >= blockSize { - var m mem.LinearMemory - // 手动管理内存，Uinx Mmap 或者 Windows VirtualAlloc - if m, err = mem.NewGuardedMemory(blockSize, maxMemorySize); err == nil { - hc = &HybridCache{memoryStore: m, blockSize: blockSize} - } + selected, err := SelectPolicy(requested, memoryCeiling) + if err != nil { + return nil, err + } + hc = &HybridCache{blockSize: blockSize, memoryCeiling: memoryCeiling} + if selected == cache.PolicyDisk { + if err := hc.initFileCache(); err != nil { + return nil, err } + return hc, nil + } + + if memoryCeiling < 0 || uint64(memoryCeiling) <= conf.AutoMemoryLimit { + hc.backingStore = &BufferStore{} + hc.memoryBacking = true + return hc, nil } - // 策略3: 文件后备 - if hc == nil { - hc = &HybridCache{blockSize: blockSize} - // 文件 - if err2 := hc.initFileCache(); err2 != nil { - return nil, errors.Join(err, err2) + + if requested == cache.PolicyMemory { + hc.memoryStore, err = mem.NewManagedMemory(blockSize, uint64(memoryCeiling), nil) + if err != nil { + return nil, err } + return hc, nil + } + + hc.memoryStore, err = mem.NewGuardedMemory(blockSize, uint64(memoryCeiling)) + if err == nil { + hc.spillOnMemoryFailure = true + return hc, nil + } + + if fileErr := hc.initFileCache(); fileErr != nil { + return nil, errors.Join(err, fileErr) } return hc, nil } var _ buffer.Block = (*HybridCache)(nil) +var _ io.Writer = (*HybridCache)(nil) diff --git a/internal/hybrid_cache/policy_test.go b/internal/hybrid_cache/policy_test.go new file mode 100644 index 0000000000..5ae505d1c9 --- /dev/null +++ b/internal/hybrid_cache/policy_test.go @@ -0,0 +1,184 @@ +package hybrid_cache + +import ( + "errors" + "io" + "os" + "testing" + + "github.com/OpenListTeam/OpenList/v4/internal/cache" + "github.com/OpenListTeam/OpenList/v4/internal/conf" + "github.com/OpenListTeam/OpenList/v4/internal/mem" +) + +func TestSelectPolicy(t *testing.T) { + errNoMemory := errors.New("no memory") + tests := []struct { + name string + requested cache.Policy + ceiling int64 + checkErr error + want cache.Policy + checks int + }{ + {"explicit memory", cache.PolicyMemory, -1, errNoMemory, cache.PolicyMemory, 0}, + {"explicit disk", cache.PolicyDisk, 1024, nil, cache.PolicyDisk, 0}, + {"auto unknown", cache.PolicyAuto, -1, nil, cache.PolicyDisk, 0}, + {"auto empty", cache.PolicyAuto, 0, nil, cache.PolicyMemory, 0}, + {"auto admitted", cache.PolicyAuto, 1024, nil, cache.PolicyMemory, 1}, + {"auto rejected", cache.PolicyAuto, 1024, errNoMemory, cache.PolicyDisk, 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + checks := 0 + got, err := selectPolicy(tt.requested, tt.ceiling, func(size uint64) error { + checks++ + if size != uint64(tt.ceiling) { + t.Fatalf("memory check size = %d, want %d", size, tt.ceiling) + } + return tt.checkErr + }) + if err != nil { + t.Fatalf("selectPolicy() error = %v", err) + } + if got != tt.want || checks != tt.checks { + t.Fatalf("selectPolicy() = %q with %d checks, want %q with %d", got, checks, tt.want, tt.checks) + } + }) + } + if _, err := selectPolicy(cache.PolicyInherit, 1, func(uint64) error { return nil }); err == nil { + t.Fatal("selectPolicy() expected an error for inherit") + } +} + +func TestHybridCacheDiskPolicy(t *testing.T) { + withCacheConfig(t, 0) + hc, err := NewHybridCache(4, 8, cache.PolicyDisk) + if err != nil { + t.Fatalf("NewHybridCache() error = %v", err) + } + store, ok := hc.backingStore.(*singleFileStore) + if !ok || hc.memoryStore != nil { + t.Fatalf("disk policy initialized unexpected stores: memory=%T backing=%T", hc.memoryStore, hc.backingStore) + } + name := store.Name() + if err := hc.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if _, err := os.Stat(name); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("cache file still exists after Close(): %v", err) + } +} + +func TestHybridCacheAutoRejectsWholeCeiling(t *testing.T) { + withCacheConfig(t, 1024) + hc, err := NewHybridCache(4, 8, cache.PolicyAuto) + if err != nil { + t.Fatalf("NewHybridCache() error = %v", err) + } + t.Cleanup(func() { _ = hc.Close() }) + if hc.memoryStore != nil || hc.memoryBacking { + t.Fatalf("auto policy initialized memory stores: memory=%T backing=%T", hc.memoryStore, hc.backingStore) + } + if _, ok := hc.backingStore.(*singleFileStore); !ok { + t.Fatalf("auto policy backing = %T, want singleFileStore", hc.backingStore) + } +} + +func TestHybridCacheStrictMemoryDoesNotSpill(t *testing.T) { + withCacheConfig(t, 0) + hc, err := NewHybridCache(4, 8, cache.PolicyMemory) + if err != nil { + t.Fatalf("NewHybridCache() error = %v", err) + } + t.Cleanup(func() { _ = hc.Close() }) + if _, err := hc.AllocBlock(4); err != nil { + t.Fatalf("AllocBlock() error = %v", err) + } + if _, err := hc.AllocBlock(5); !errors.Is(err, mem.ErrNotEnoughMemory) { + t.Fatalf("AllocBlock() error = %v, want ErrNotEnoughMemory", err) + } + if hc.backingStore != nil { + t.Fatalf("strict memory policy spilled to %T", hc.backingStore) + } +} + +func TestHybridCacheUnknownMemory(t *testing.T) { + withCacheConfig(t, 0) + hc, err := NewHybridCache(3, -1, cache.PolicyMemory) + if err != nil { + t.Fatalf("NewHybridCache() error = %v", err) + } + t.Cleanup(func() { _ = hc.Close() }) + if _, ok := hc.backingStore.(*BufferStore); !ok { + t.Fatalf("unknown memory cache backing = %T, want BufferStore", hc.backingStore) + } + if _, err := hc.Write([]byte("abcdefg")); err != nil { + t.Fatalf("Write() error = %v", err) + } + got := make([]byte, 7) + if _, err := hc.ReadAt(got, 0); err != nil { + t.Fatalf("ReadAt() error = %v", err) + } + if string(got) != "abcdefg" || hc.Size() != 7 { + t.Fatalf("cache = %q size=%d", got, hc.Size()) + } +} + +func TestHybridCacheAutoSpillKeepsMemoryPrefix(t *testing.T) { + withCacheConfig(t, 0) + hc := &HybridCache{ + blockSize: 2, + memoryStore: &limitedMemory{buf: make([]byte, 0, 2)}, + spillOnMemoryFailure: true, + memoryCeiling: 2, + } + t.Cleanup(func() { _ = hc.Close() }) + if _, err := hc.Write([]byte("abcd")); err != nil { + t.Fatalf("Write() error = %v", err) + } + if hc.memoryOffset != 2 || hc.backingOffset != 2 { + t.Fatalf("offsets = memory:%d disk:%d, want 2 and 2", hc.memoryOffset, hc.backingOffset) + } + got := make([]byte, 4) + if _, err := hc.ReadAt(got, 0); err != nil { + t.Fatalf("ReadAt() error = %v", err) + } + if string(got) != "abcd" { + t.Fatalf("cache = %q, want abcd", got) + } +} + +type limitedMemory struct { + buf []byte +} + +func (m *limitedMemory) Reallocate(size uint64) ([]byte, error) { + if size > uint64(cap(m.buf)) { + return nil, mem.ErrNotEnoughMemory + } + m.buf = m.buf[:size] + return m.buf, nil +} + +func (m *limitedMemory) Free() error { + m.buf = nil + return nil +} + +func withCacheConfig(t *testing.T, autoMemoryLimit uint64) { + t.Helper() + oldConf := conf.Conf + oldLimit := conf.AutoMemoryLimit + oldMinFreeMemory := conf.MinFreeMemory + conf.Conf = &conf.Config{TempDir: t.TempDir()} + conf.AutoMemoryLimit = autoMemoryLimit + conf.MinFreeMemory = 0 + t.Cleanup(func() { + conf.Conf = oldConf + conf.AutoMemoryLimit = oldLimit + conf.MinFreeMemory = oldMinFreeMemory + }) +} + +var _ io.Writer = (*HybridCache)(nil) diff --git a/internal/mem/utils.go b/internal/mem/utils.go index ef82f99589..9f7bd82e90 100644 --- a/internal/mem/utils.go +++ b/internal/mem/utils.go @@ -45,8 +45,16 @@ func MemoryGrowCheck(growSize uint64) error { } func NewGuardedMemory(cap, max uint64) (m LinearMemory, err error) { - if err := MemoryGrowCheck(cap); err != nil { - return nil, err + return NewManagedMemory(cap, max, MemoryGrowCheck) +} + +// NewManagedMemory creates memory with panic recovery and lifecycle cleanup. +// A nil growCheck intentionally permits growth without an availability check. +func NewManagedMemory(cap, max uint64, growCheck GrowCheck) (m LinearMemory, err error) { + if growCheck != nil { + if err := growCheck(cap); err != nil { + return nil, err + } } defer func() { if r := recover(); r != nil { @@ -57,8 +65,10 @@ func NewGuardedMemory(cap, max uint64) (m LinearMemory, err error) { if err != nil { return nil, err } - if s, ok := m.(interface{ SetGrowCheck(GrowCheck) }); ok { - s.SetGrowCheck(MemoryGrowCheck) + if growCheck != nil { + if s, ok := m.(interface{ SetGrowCheck(GrowCheck) }); ok { + s.SetGrowCheck(growCheck) + } } gm := &guardedMemory{LinearMemory: m} gm.cleanup = runtime.AddCleanup(gm, func(m LinearMemory) { diff --git a/internal/model/obj.go b/internal/model/obj.go index 1269b5b797..ef7689d522 100644 --- a/internal/model/obj.go +++ b/internal/model/obj.go @@ -6,6 +6,7 @@ import ( "strings" "time" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/pkg/http_range" "github.com/OpenListTeam/OpenList/v4/pkg/utils" "github.com/dlclark/regexp2" @@ -47,6 +48,8 @@ type FileStreamer interface { IsForceStreamUpload() bool GetExist() Obj SetExist(Obj) + GetCachePolicy() cache.Policy + SetCachePolicy(cache.Policy) error // for a non-seekable Stream, RangeRead supports peeking some data, and CacheFullAndWriter still works RangeRead(http_range.Range) (io.Reader, error) // for a non-seekable Stream, if Read is called, this function won't work. diff --git a/internal/net/request.go b/internal/net/request.go index 0cfa7942ea..152bcc138d 100644 --- a/internal/net/request.go +++ b/internal/net/request.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "math" "math/rand/v2" "net/http" stdpath "path" @@ -13,6 +14,7 @@ import ( "sync/atomic" "time" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/internal/conf" "github.com/OpenListTeam/OpenList/v4/internal/errs" hcache "github.com/OpenListTeam/OpenList/v4/internal/hybrid_cache" @@ -40,6 +42,8 @@ var DefaultConcurrencyLimit *ConcurrencyLimit type Downloader struct { PartSize int + // CachePolicy overrides the global cache policy. The zero value inherits it. + CachePolicy cache.Policy // PartBodyMaxRetries is the number of retry attempts to make for failed part downloads. PartBodyMaxRetries int @@ -92,9 +96,20 @@ func (d Downloader) Download(ctx context.Context, p *HttpRequestParams) (readClo if conf.MinFreeMemory > 0 && impl.cfg.PartSize > int(conf.MaxBlockLimit) { impl.cfg.PartSize = int(conf.MaxBlockLimit) } + if impl.cfg.Concurrency <= 0 { + return nil, fmt.Errorf("download concurrency must be positive") + } + if impl.cfg.PartSize <= 0 { + return nil, fmt.Errorf("download part size must be positive") + } if impl.cfg.HttpClient == nil { impl.cfg.HttpClient = DefaultHttpRequestFunc } + policy, err := cache.ResolvePolicy(impl.cfg.CachePolicy, conf.CachePolicy) + if err != nil { + return nil, err + } + impl.cfg.CachePolicy = policy return impl.download() } @@ -197,8 +212,9 @@ func (d *downloader) download() (io.ReadCloser, error) { d.maxPos = d.params.Range.Start + d.params.Range.Length d.concurrency = d.cfg.Concurrency + memoryCeiling := downloaderMemoryCeiling(d.params.Range.Length, d.cfg.Concurrency, d.cfg.PartSize) var err error - d.hc, err = hcache.NewHybridCache(uint64(d.cfg.PartSize), uint64(d.params.Range.Length)) + d.hc, err = hcache.NewHybridCache(uint64(d.cfg.PartSize), memoryCeiling, d.cfg.CachePolicy) if err == nil { d.bufMap = make(map[int]*buffer.PipeBuffer, d.cfg.Concurrency) err = d.sendChunkTask(true) @@ -214,6 +230,13 @@ func (d *downloader) download() (io.ReadCloser, error) { return &multiReadCloser{d: d, curBuf: d.popBuf(0), maxPos: maxPart}, nil } +func downloaderMemoryCeiling(rangeLength int64, concurrency, partSize int) int64 { + if int64(concurrency) > math.MaxInt64/int64(partSize) { + return rangeLength + } + return min(rangeLength, int64(concurrency)*int64(partSize)) +} + func (d *downloader) sendChunkTask(newConcurrency bool) (err error) { d.mu.Lock() defer d.mu.Unlock() diff --git a/internal/net/request_test.go b/internal/net/request_test.go index 0fdc56eb33..afd1771a6f 100644 --- a/internal/net/request_test.go +++ b/internal/net/request_test.go @@ -13,6 +13,7 @@ import ( "testing" "time" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/pkg/http_range" "github.com/sirupsen/logrus" ) @@ -26,6 +27,27 @@ func containsString(slice []string, val string) bool { return false } +func TestDownloaderMemoryCeiling(t *testing.T) { + tests := []struct { + name string + rangeLength int64 + concurrency int + partSize int + want int64 + }{ + {"working set", 100 << 20, 2, 8 << 20, 16 << 20}, + {"range smaller than pool", 10, 4, 8, 10}, + {"multiplication overflow", 100, int(^uint(0) >> 1), 2, 100}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := downloaderMemoryCeiling(tt.rangeLength, tt.concurrency, tt.partSize); got != tt.want { + t.Fatalf("downloaderMemoryCeiling() = %d, want %d", got, tt.want) + } + }) + } +} + func TestDownloadOrder(t *testing.T) { buff := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15} downloader, invocations, ranges := newDownloadRangeClient(buff) @@ -33,6 +55,7 @@ func TestDownloadOrder(t *testing.T) { d := NewDownloader(func(d *Downloader) { d.Concurrency = con d.PartSize = partSize + d.CachePolicy = cache.PolicyMemory d.HttpClient = downloader.HttpRequest }) @@ -122,6 +145,7 @@ func TestHighConcurrency(t *testing.T) { d := NewDownloader(func(d *Downloader) { d.Concurrency = con d.PartSize = partSize + d.CachePolicy = cache.PolicyMemory d.HttpClient = downloader.HttpRequest d.ConcurrencyLimit = &ConcurrencyLimit{ Limit: concurrencyLimit, diff --git a/internal/stream/stream.go b/internal/stream/stream.go index b1e6fd5faf..189bd93508 100644 --- a/internal/stream/stream.go +++ b/internal/stream/stream.go @@ -1,15 +1,14 @@ package stream import ( - "bytes" "context" "errors" "fmt" "io" "math" - "os" "sync" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/internal/conf" hcache "github.com/OpenListTeam/OpenList/v4/internal/hybrid_cache" "github.com/OpenListTeam/OpenList/v4/internal/model" @@ -28,14 +27,16 @@ type FileStream struct { ForceStreamUpload bool Exist model.Obj //the file existed in the destination, we can reuse some info since we wil overwrite it utils.Closers - size int64 - oriReader io.Reader // the original reader, used for caching - hc *hcache.HybridCache - peek buffer.SizedReadAtSeeker + size int64 + sizeSet bool + cachePolicyOverride cache.Policy + oriReader io.Reader // the original reader, used for caching + hc *hcache.HybridCache + peek buffer.SizedReadAtSeeker } func (f *FileStream) GetSize() int64 { - if f.size > 0 { + if f.sizeSet { return f.size } return f.Obj.GetSize() @@ -60,6 +61,25 @@ func (f *FileStream) SetExist(obj model.Obj) { f.Exist = obj } +func (f *FileStream) GetCachePolicy() cache.Policy { + policy, err := cache.ResolvePolicy(f.cachePolicyOverride, conf.CachePolicy) + if err != nil { + panic(err) + } + return policy +} + +func (f *FileStream) SetCachePolicy(policy cache.Policy) error { + if policy != cache.PolicyInherit && !policy.IsConcrete() { + return fmt.Errorf("invalid cache policy %q", policy) + } + if f.peek != nil { + return errors.New("cache policy cannot be changed after cache initialization") + } + f.cachePolicyOverride = policy + return nil +} + // CacheFullAndWriter save all data into tmpFile or memory. // It's not thread-safe! func (f *FileStream) CacheFullAndWriter(up *model.UpdateProgress, writer io.Writer) (model.File, error) { @@ -109,36 +129,7 @@ func (f *FileStream) CacheFullAndWriter(up *model.UpdateProgress, writer io.Writ reader = io.TeeReader(reader, writer) } - // 如果文件大小未知，直接缓存到磁盘 - if f.GetSize() < 0 { - // 检查是否有数据 - buf := []byte{0} - n, err := io.ReadFull(reader, buf) - br := bytes.NewReader(buf[:n]) - if err == io.ErrUnexpectedEOF || err == io.EOF { - f.size = br.Size() - f.Reader = br - return br, nil - } else if err != nil { - return nil, err - } - tmpF, err := utils.CreateTempFile(io.MultiReader(br, reader), 0) - if err != nil { - return nil, err - } - f.Add(utils.CloseFunc(func() error { - return errors.Join(tmpF.Close(), os.RemoveAll(tmpF.Name())) - })) - stat, err := tmpF.Stat() - if err != nil { - return nil, err - } - f.size = stat.Size() - f.Reader = tmpF - return tmpF, nil - } - - if up != nil { + if up != nil && f.GetSize() >= 0 { cacheProgress := model.UpdateProgressWithRange(*up, 0, 50) *up = model.UpdateProgressWithRange(*up, 50, 100) size := f.GetSize() @@ -197,9 +188,16 @@ func (f *FileStream) RangeRead(httpRange http_range.Range) (io.Reader, error) { // 确保指定大小的数据被缓存 func (f *FileStream) ensureCache(size int64) (model.File, error) { if f.peek == nil { - blockSize := min(size, f.GetSize(), int64(conf.MaxBlockLimit)) + memoryCeiling := f.GetSize() + blockSize := int64(conf.MaxBlockLimit) + if memoryCeiling >= 0 { + blockSize = min(memoryCeiling, int64(conf.MaxBlockLimit)) + if size > 0 { + blockSize = min(blockSize, size) + } + } var err error - f.hc, err = hcache.NewHybridCache(uint64(blockSize), uint64(f.GetSize())) + f.hc, err = hcache.NewHybridCache(uint64(blockSize), memoryCeiling, f.GetCachePolicy()) if err != nil { return nil, err } @@ -208,6 +206,16 @@ func (f *FileStream) ensureCache(size int64) (model.File, error) { f.Reader = io.MultiReader(f.peek, f.oriReader) f.Add(f.hc) } + if size < 0 { + _, err := utils.CopyWithBuffer(f.hc, f.oriReader) + if err != nil { + return nil, err + } + f.size = f.peek.Size() + f.sizeSet = true + f.Reader = f.peek + return f.peek, nil + } size = size - f.peek.Size() if size <= 0 { return f.peek, nil @@ -264,6 +272,7 @@ func NewSeekableStream(fs *FileStream, link *model.Link) (*SeekableStream, error fs.Add(rc) } fs.size = size + fs.sizeSet = true fs.Add(link) return &SeekableStream{FileStream: fs, rangeReader: rr}, nil } diff --git a/internal/stream/stream_test.go b/internal/stream/stream_test.go index 1d8d002e2d..22e6b1866f 100644 --- a/internal/stream/stream_test.go +++ b/internal/stream/stream_test.go @@ -5,8 +5,10 @@ import ( "errors" "fmt" "io" + "os" "testing" + "github.com/OpenListTeam/OpenList/v4/internal/cache" "github.com/OpenListTeam/OpenList/v4/internal/conf" "github.com/OpenListTeam/OpenList/v4/internal/model" "github.com/OpenListTeam/OpenList/v4/internal/stream" @@ -14,6 +16,98 @@ import ( "github.com/OpenListTeam/OpenList/v4/pkg/utils" ) +func TestFileStreamCachePolicy(t *testing.T) { + oldConf := conf.Conf + oldPolicy := conf.CachePolicy + oldBlockLimit := conf.MaxBlockLimit + oldAutoMemoryLimit := conf.AutoMemoryLimit + t.Cleanup(func() { + conf.Conf = oldConf + conf.CachePolicy = oldPolicy + conf.MaxBlockLimit = oldBlockLimit + conf.AutoMemoryLimit = oldAutoMemoryLimit + }) + conf.MaxBlockLimit = 4 + conf.AutoMemoryLimit = 0 + + t.Run("inherit and override", func(t *testing.T) { + conf.CachePolicy = cache.PolicyDisk + f := &stream.FileStream{} + if got := f.GetCachePolicy(); got != cache.PolicyDisk { + t.Fatalf("GetCachePolicy() = %q, want disk", got) + } + if err := f.SetCachePolicy(cache.PolicyMemory); err != nil { + t.Fatalf("SetCachePolicy() error = %v", err) + } + if got := f.GetCachePolicy(); got != cache.PolicyMemory { + t.Fatalf("GetCachePolicy() = %q, want memory", got) + } + if err := f.SetCachePolicy(cache.PolicyInherit); err != nil { + t.Fatalf("SetCachePolicy(inherit) error = %v", err) + } + if got := f.GetCachePolicy(); got != cache.PolicyDisk { + t.Fatalf("GetCachePolicy() = %q after inherit, want disk", got) + } + }) + + for _, tt := range []struct { + policy cache.Policy + wantFile bool + }{ + {cache.PolicyAuto, true}, + {cache.PolicyDisk, true}, + {cache.PolicyMemory, false}, + } { + t.Run(string(tt.policy)+" unknown size", func(t *testing.T) { + tempDir := t.TempDir() + conf.Conf = &conf.Config{TempDir: tempDir} + conf.CachePolicy = cache.PolicyAuto + input := []byte("unknown-size-stream") + f := &stream.FileStream{ + Obj: &model.Object{Size: -1}, + Reader: io.NopCloser(bytes.NewReader(input)), + } + if err := f.SetCachePolicy(tt.policy); err != nil { + t.Fatalf("SetCachePolicy() error = %v", err) + } + cached, err := f.CacheFullAndWriter(nil, nil) + if err != nil { + t.Fatalf("CacheFullAndWriter() error = %v", err) + } + if f.GetSize() != int64(len(input)) { + t.Fatalf("GetSize() = %d, want %d", f.GetSize(), len(input)) + } + got, err := io.ReadAll(cached) + if err != nil { + t.Fatalf("ReadAll() error = %v", err) + } + if !bytes.Equal(got, input) { + t.Fatalf("cached content = %q, want %q", got, input) + } + entries, err := os.ReadDir(tempDir) + if err != nil { + t.Fatalf("ReadDir() error = %v", err) + } + if gotFile := len(entries) > 0; gotFile != tt.wantFile { + t.Fatalf("temporary file present = %v, want %v", gotFile, tt.wantFile) + } + if err := f.SetCachePolicy(cache.PolicyDisk); err == nil { + t.Fatal("SetCachePolicy() expected an error after cache initialization") + } + if err := f.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + entries, err = os.ReadDir(tempDir) + if err != nil { + t.Fatalf("ReadDir() after close error = %v", err) + } + if len(entries) != 0 { + t.Fatalf("temporary files remain after close: %v", entries) + } + }) + } +} + func TestRangeRead(t *testing.T) { type args struct { httpRange http_range.Range diff --git a/internal/stream/util.go b/internal/stream/util.go index 2947fcbc2b..adbfa428bb 100644 --- a/internal/stream/util.go +++ b/internal/stream/util.go @@ -186,9 +186,16 @@ func NewStreamSectionReader(file model.FileStreamer, sectionSize int, up *model. if file.GetFile() != nil { return &cachedSectionReader{file.GetFile()}, nil } + if sectionSize <= 0 { + return nil, fmt.Errorf("section size must be positive") + } - blockSize := min(uint64(sectionSize), uint64(file.GetSize()), conf.MaxBlockLimit) - hc, err := hcache.NewHybridCache(blockSize, uint64(file.GetSize())) + fileSize := file.GetSize() + blockSize := min(uint64(sectionSize), conf.MaxBlockLimit) + if fileSize >= 0 { + blockSize = min(blockSize, uint64(fileSize)) + } + hc, err := hcache.NewHybridCache(blockSize, fileSize, file.GetCachePolicy()) if err != nil { return nil, err }