feat(config): add cache configuration options for TTL, num counters, and max cost

This commit is contained in:
krau
2025-06-18 10:51:03 +08:00
parent f5e33472eb
commit 758564d436
2 changed files with 22 additions and 11 deletions

View File

@@ -2,19 +2,19 @@ package cache
import ( import (
"fmt" "fmt"
"time"
"github.com/charmbracelet/log" "github.com/charmbracelet/log"
"github.com/dgraph-io/ristretto/v2" "github.com/dgraph-io/ristretto/v2"
"github.com/krau/SaveAny-Bot/config"
) )
var cache *ristretto.Cache[string, any] var cache *ristretto.Cache[string, any]
// TODO: maybe we should use simple ttl cache instead of ristretto...
func init() { func init() {
c, err := ristretto.NewCache(&ristretto.Config[string, any]{ c, err := ristretto.NewCache(&ristretto.Config[string, any]{
NumCounters: 1e5, NumCounters: config.Cfg.Cache.NumCounters,
MaxCost: 1e6, // 1000000 / 112 ≈ 8928 MaxCost: config.Cfg.Cache.MaxCost,
BufferItems: 64, BufferItems: 64,
OnReject: func(item *ristretto.Item[any]) { OnReject: func(item *ristretto.Item[any]) {
log.Warnf("Cache item rejected: key=%d, value=%v", item.Key, item.Value) log.Warnf("Cache item rejected: key=%d, value=%v", item.Key, item.Value)
@@ -27,7 +27,7 @@ func init() {
} }
func Set(key string, value any) error { func Set(key string, value any) error {
ok := cache.Set(key, value, 0) ok := cache.SetWithTTL(key, value, 0, time.Duration(config.Cfg.Cache.TTL)*time.Second)
if !ok { if !ok {
return fmt.Errorf("failed to set value in cache") return fmt.Errorf("failed to set value in cache")
} }

View File

@@ -15,12 +15,13 @@ import (
) )
type Config struct { type Config struct {
Lang string `toml:"lang" mapstructure:"lang" json:"lang"` Lang string `toml:"lang" mapstructure:"lang" json:"lang"`
Workers int `toml:"workers" mapstructure:"workers"` Workers int `toml:"workers" mapstructure:"workers"`
Retry int `toml:"retry" mapstructure:"retry"` Retry int `toml:"retry" mapstructure:"retry"`
NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"` NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"`
Threads int `toml:"threads" mapstructure:"threads" json:"threads"` Threads int `toml:"threads" mapstructure:"threads" json:"threads"`
Stream bool `toml:"stream" mapstructure:"stream" json:"stream"` Stream bool `toml:"stream" mapstructure:"stream" json:"stream"`
Cache cacheConfig `toml:"cache" mapstructure:"cache" json:"cache"`
Users []userConfig `toml:"users" mapstructure:"users" json:"users"` Users []userConfig `toml:"users" mapstructure:"users" json:"users"`
@@ -30,6 +31,12 @@ type Config struct {
Storages []storage.StorageConfig `toml:"-" mapstructure:"-" json:"storages"` Storages []storage.StorageConfig `toml:"-" mapstructure:"-" json:"storages"`
} }
type cacheConfig struct {
TTL int64 `toml:"ttl" mapstructure:"ttl" json:"ttl"`
NumCounters int64 `toml:"num_counters" mapstructure:"num_counters" json:"num_counters"`
MaxCost int64 `toml:"max_cost" mapstructure:"max_cost" json:"max_cost"`
}
type tempConfig struct { type tempConfig struct {
BasePath string `toml:"base_path" mapstructure:"base_path" json:"base_path"` BasePath string `toml:"base_path" mapstructure:"base_path" json:"base_path"`
CacheTTL int64 `toml:"cache_ttl" mapstructure:"cache_ttl" json:"cache_ttl"` CacheTTL int64 `toml:"cache_ttl" mapstructure:"cache_ttl" json:"cache_ttl"`
@@ -87,6 +94,10 @@ func Init(ctx context.Context) error {
viper.SetDefault("retry", 3) viper.SetDefault("retry", 3)
viper.SetDefault("threads", 4) viper.SetDefault("threads", 4)
viper.SetDefault("cache.ttl", 86400)
viper.SetDefault("cache.num_counters", 1e5)
viper.SetDefault("cache.max_cost", 1e6)
viper.SetDefault("telegram.app_id", 1025907) viper.SetDefault("telegram.app_id", 1025907)
viper.SetDefault("telegram.app_hash", "452b0359b988148995f22ff0f4229750") viper.SetDefault("telegram.app_hash", "452b0359b988148995f22ff0f4229750")
viper.SetDefault("telegram.timeout", 60) viper.SetDefault("telegram.timeout", 60)