From 806623310fb352de0aa1739aa950af98f7fce2bd Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 01:09:04 +0800 Subject: [PATCH 01/11] perf: scale shared tenant backend cache --- cmd/drive9-server/main.go | 24 +++- cmd/drive9-server/main_test.go | 43 ++++++ pkg/backend/quota.go | 5 +- pkg/backend/quota_cache.go | 146 +++++++------------- pkg/backend/quota_cache_test.go | 189 +++++++++++++++----------- pkg/backend/quota_integration_test.go | 19 ++- pkg/tenant/pool.go | 7 + pkg/tenant/pool_test.go | 56 ++++++++ 8 files changed, 306 insertions(+), 183 deletions(-) diff --git a/cmd/drive9-server/main.go b/cmd/drive9-server/main.go index e556537f..586cbec1 100644 --- a/cmd/drive9-server/main.go +++ b/cmd/drive9-server/main.go @@ -43,8 +43,10 @@ import ( ) const ( - defaultListenAddr = ":9009" - defaultS3Dir = "s3" + defaultListenAddr = ":9009" + defaultS3Dir = "s3" + defaultTenantBackendCacheMaxTenants = 1024 + defaultSharedBackendCacheMaxTenants = 20480 ) type s3Config struct { @@ -341,7 +343,7 @@ func main() { S3SessionToken: s3cfg.SessionToken, S3EncryptionPolicy: s3cfg.EncryptionPolicy, BackendOptions: backendOptions, - MaxTenants: envInt("DRIVE9_POOL_MAX_TENANTS", 0), + MaxTenants: tenantBackendCacheMaxTenants(providerType), IdleTimeout: envDuration("DRIVE9_POOL_IDLE_TTL", 5*time.Minute), IdleReapInterval: envDuration("DRIVE9_POOL_IDLE_REAP_INTERVAL", 2*time.Minute), DisableDatabaseAutoEmbedding: disableDatabaseAutoEmbedding, @@ -621,8 +623,8 @@ environment: DRIVE9_LISTEN_ADDR serve listen address (default: :9009) DRIVE9_PUBLIC_URL externally reachable base URL for presigned URLs (required for remote clients) DRIVE9_META_DSN control-plane MySQL DSN (required) - DRIVE9_POOL_MAX_TENANTS max cached tenant user DB pools per pod (default: 1024) - DRIVE9_POOL_IDLE_TTL idle duration before a cached tenant backend is evicted (default: 5m, 0=disabled) + DRIVE9_POOL_MAX_TENANTS max cached tenant backends per pod (default: 20480 for tidb_cloud_native_shared, 1024 otherwise) + DRIVE9_POOL_IDLE_TTL idle duration before a standalone tenant backend is evicted (default: 5m, 0=disabled; shared tenants use capacity LRU only) DRIVE9_POOL_IDLE_REAP_INTERVAL how often the idle reaper scans (default: 2m) DRIVE9_META_DB_MAX_OPEN_CONNS max open connections for the per-pod meta DB pool (default: 100) DRIVE9_META_DB_MAX_IDLE_CONNS max idle connections for the per-pod meta DB pool (default: 20) @@ -1156,6 +1158,18 @@ func dbHealthProbeOptionsFromEnv() metrics.DBHealthProbeOptions { } } +func tenantBackendCacheMaxTenants(provider string) int { + fallback := defaultTenantBackendCacheMaxTenants + if provider == tenant.ProviderTiDBCloudNativeShared { + fallback = defaultSharedBackendCacheMaxTenants + } + maxTenants := envInt("DRIVE9_POOL_MAX_TENANTS", fallback) + if maxTenants <= 0 { + return fallback + } + return maxTenants +} + func envInt(key string, fallback int) int { raw := strings.TrimSpace(os.Getenv(key)) if raw == "" { diff --git a/cmd/drive9-server/main_test.go b/cmd/drive9-server/main_test.go index bcade259..99c6c0ae 100644 --- a/cmd/drive9-server/main_test.go +++ b/cmd/drive9-server/main_test.go @@ -9,6 +9,7 @@ import ( "github.com/mem9-ai/drive9/pkg/backend" "github.com/mem9-ai/drive9/pkg/meta" "github.com/mem9-ai/drive9/pkg/server" + "github.com/mem9-ai/drive9/pkg/tenant" ) func TestVersionTextUsesDrive9ServerComponent(t *testing.T) { @@ -18,6 +19,48 @@ func TestVersionTextUsesDrive9ServerComponent(t *testing.T) { } } +func TestTenantBackendCacheMaxTenantsUsesProviderDefault(t *testing.T) { + const key = "DRIVE9_POOL_MAX_TENANTS" + restore := snapshotEnv(t, []string{key}) + t.Cleanup(func() { restoreEnv(t, restore) }) + unsetEnv(t, []string{key}) + + tests := []struct { + provider string + want int + }{ + {provider: tenant.ProviderTiDBCloudNativeShared, want: 20480}, + {provider: tenant.ProviderTiDBCloudNative, want: 1024}, + {provider: tenant.ProviderTiDBZero, want: 1024}, + {provider: tenant.ProviderDB9, want: 1024}, + } + for _, tt := range tests { + t.Run(tt.provider, func(t *testing.T) { + if got := tenantBackendCacheMaxTenants(tt.provider); got != tt.want { + t.Fatalf("tenantBackendCacheMaxTenants(%q) = %d, want %d", tt.provider, got, tt.want) + } + }) + } +} + +func TestTenantBackendCacheMaxTenantsExplicitOverrideWins(t *testing.T) { + const key = "DRIVE9_POOL_MAX_TENANTS" + restore := snapshotEnv(t, []string{key}) + t.Cleanup(func() { restoreEnv(t, restore) }) + setEnv(t, key, "4096") + + for _, provider := range []string{ + tenant.ProviderTiDBCloudNativeShared, + tenant.ProviderTiDBCloudNative, + tenant.ProviderTiDBZero, + tenant.ProviderDB9, + } { + if got := tenantBackendCacheMaxTenants(provider); got != 4096 { + t.Fatalf("tenantBackendCacheMaxTenants(%q) = %d, want explicit 4096", provider, got) + } + } +} + func TestSlockOAuthFromEnvDisabledByDefault(t *testing.T) { keys := []string{ "DRIVE9_SLOCK_ORIGIN", diff --git a/pkg/backend/quota.go b/pkg/backend/quota.go index 1056e3e1..1a1e1a21 100644 --- a/pkg/backend/quota.go +++ b/pkg/backend/quota.go @@ -228,10 +228,7 @@ func (b *Dat9Backend) checkStorageQuotaServerTx(ctx context.Context, tx *sql.Tx, // falling back to a synchronous DB query when the cache is unavailable. func (b *Dat9Backend) cachedQuotaConfig(ctx context.Context) *QuotaConfigView { if b.quotaConfigCache != nil { - if cfg := b.quotaConfigCache.get(); cfg != nil { - return cfg - } - return b.quotaConfigCache.load(ctx) + return b.quotaConfigCache.get(ctx) } if b.metaStore == nil { return nil diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index 3bd961b3..fdb3a291 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -12,10 +12,12 @@ import ( ) const ( - // defaultQuotaConfigCacheRefreshInterval is the default interval for the - // background goroutine that polls the tenant quota config version from - // the central DB. Override with DRIVE9_QUOTA_CACHE_REFRESH_SECONDS. + // defaultQuotaConfigCacheRefreshInterval is the default TTL for lazily + // loaded tenant quota config. Override with DRIVE9_QUOTA_CACHE_REFRESH_SECONDS. defaultQuotaConfigCacheRefreshInterval = 30 * time.Second + // quotaConfigCacheFailureRetryInterval prevents a central DB outage from + // turning every tenant request into another config query. + quotaConfigCacheFailureRetryInterval = 5 * time.Second // quotaUsageCacheTTL bounds how long soft small-write quota checks may // reuse central usage counters. Strict upload reservations still read // central usage directly. @@ -60,11 +62,6 @@ func InitQuotaAdmissionCacheTTLs(usageTTL, pendingDeltasTTL time.Duration) { } } -type quotaConfigSnapshot struct { - config *QuotaConfigView - version string -} - func cloneQuotaConfigView(cfg *QuotaConfigView) *QuotaConfigView { if cfg == nil { return nil @@ -73,136 +70,97 @@ func cloneQuotaConfigView(cfg *QuotaConfigView) *QuotaConfigView { return &cp } -// quotaConfigCache is a per-tenant cache for low-frequency quota config. It -// only removes repeated config reads and uses version polling so config changes -// converge without a cross-server invalidation channel. +// quotaConfigCache is a passive per-tenant cache for low-frequency quota +// config. Requests refresh an expired snapshot; idle tenants create no +// goroutines and issue no quota queries. type quotaConfigCache struct { tenantID string tidbCloudOrgID string store MetaQuotaStore - mu sync.RWMutex - snapshot *quotaConfigSnapshot - loadMu sync.Mutex - - cancel context.CancelFunc - done chan struct{} + mu sync.RWMutex + snapshot *QuotaConfigView + nextRefresh time.Time + loadMu sync.Mutex } -// newQuotaConfigCache creates and starts a background-refreshing config cache. -// Backend construction must stay cheap for read-only operations such as ls, so -// the initial load is lazy: the first quota check loads config on demand and -// the background refresher keeps it current after that. +// newQuotaConfigCache creates an empty request-driven config cache. Backend +// construction stays cheap: the first quota check loads config on demand. func newQuotaConfigCache(tenantID, tidbCloudOrgID string, store MetaQuotaStore) *quotaConfigCache { - ctx, cancel := context.WithCancel(context.Background()) - c := "aConfigCache{ + return "aConfigCache{ tenantID: tenantID, tidbCloudOrgID: normalizeTenantMetricTiDBCloudOrgID(tidbCloudOrgID), store: store, - cancel: cancel, - done: make(chan struct{}), } - go c.run(ctx) - return c } -// get returns a copy of the cached quota config. Returns nil when the cache has -// not been populated yet, allowing callers to fail open or fall back. -func (c *quotaConfigCache) get() *QuotaConfigView { +// cached returns the current snapshot and whether another store read is +// suppressed until nextRefresh. A nil snapshot can still be current during +// the short retry cooldown after an initial load failure. +func (c *quotaConfigCache) cached(now time.Time) (*QuotaConfigView, bool) { c.mu.RLock() defer c.mu.RUnlock() - if c.snapshot == nil || c.snapshot.config == nil { - return nil + if !now.Before(c.nextRefresh) { + return nil, false } - return cloneQuotaConfigView(c.snapshot.config) + if c.snapshot == nil { + return nil, true + } + return cloneQuotaConfigView(c.snapshot), true +} + +// get returns a defensive copy of the cached config, refreshing it once when +// its TTL has expired. Concurrent expired requests coalesce on loadMu. +func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { + return c.load(ctx) } func (c *quotaConfigCache) load(ctx context.Context) *QuotaConfigView { - if cfg := c.get(); cfg != nil { + now := time.Now() + if cfg, current := c.cached(now); current { return cfg } c.loadMu.Lock() defer c.loadMu.Unlock() - if cfg := c.get(); cfg != nil { + now = time.Now() + if cfg, current := c.cached(now); current { return cfg } - start := time.Now() + start := now cfg, err := c.store.GetQuotaConfig(ctx, c.tenantID) if err != nil { logger.Warn(ctx, "quota_config_cache_config_failed", zap.String("tenant_id", c.tenantID), zap.Error(err)) - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "config_error", time.Since(start)) - return nil + return c.finishFailedLoad(start, "config_error") } if cfg == nil { - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "config_empty", time.Since(start)) - return nil + return c.finishFailedLoad(start, "config_empty") } c.mu.Lock() - if c.snapshot != nil && c.snapshot.config != nil { - existing := cloneQuotaConfigView(c.snapshot.config) - c.mu.Unlock() - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "raced_refresh", time.Since(start)) - return existing - } - c.snapshot = "aConfigSnapshot{config: cloneQuotaConfigView(cfg), version: ""} + c.snapshot = cloneQuotaConfigView(cfg) + c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshInterval) c.mu.Unlock() metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "ok", time.Since(start)) return cloneQuotaConfigView(cfg) } -func (c *quotaConfigCache) stop() { - c.cancel() - <-c.done -} - -func (c *quotaConfigCache) run(ctx context.Context) { - defer close(c.done) - ticker := time.NewTicker(quotaConfigCacheRefreshInterval) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - c.refresh(ctx) - } - } -} - -func (c *quotaConfigCache) refresh(ctx context.Context) { - start := time.Now() - version, err := c.store.GetQuotaConfigVersion(ctx, c.tenantID) - if err != nil { - logger.Warn(ctx, "quota_config_cache_version_failed", - zap.String("tenant_id", c.tenantID), zap.Error(err)) - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "refresh", "version_error", time.Since(start)) - return - } - - c.mu.RLock() - snapshot := c.snapshot - if snapshot != nil && snapshot.version == version { - c.mu.RUnlock() - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "refresh", "unchanged", time.Since(start)) - return - } - c.mu.RUnlock() - - cfg, err := c.store.GetQuotaConfig(ctx, c.tenantID) - if err != nil { - logger.Warn(ctx, "quota_config_cache_config_failed", - zap.String("tenant_id", c.tenantID), zap.Error(err)) - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "refresh", "config_error", time.Since(start)) - return - } +func (c *quotaConfigCache) finishFailedLoad(start time.Time, result string) *QuotaConfigView { c.mu.Lock() - c.snapshot = "aConfigSnapshot{config: cloneQuotaConfigView(cfg), version: version} + c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryInterval) + var stale *QuotaConfigView + if c.snapshot != nil { + stale = cloneQuotaConfigView(c.snapshot) + } c.mu.Unlock() - metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "refresh", "ok", time.Since(start)) + metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", result, time.Since(start)) + return stale } +// stop remains for backend-close call-site compatibility. The cache owns no +// goroutine or other lifecycle resource. +func (c *quotaConfigCache) stop() {} + type quotaUsageSnapshot struct { usage *QuotaUsageView expiresAt time.Time diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 5406f2ce..776c740e 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -2,6 +2,7 @@ package backend import ( "context" + "errors" "sync" "sync/atomic" "testing" @@ -51,16 +52,18 @@ func (m *cacheTestStore) GetQuotaConfigVersion(ctx context.Context, tenantID str return m.fakeMetaQuotaStore.GetQuotaConfigVersion(ctx, tenantID) } -func TestQuotaConfigCacheLazyLoad(t *testing.T) { +func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { + previousRefreshInterval := quotaConfigCacheRefreshInterval + quotaConfigCacheRefreshInterval = 5 * time.Millisecond + t.Cleanup(func() { quotaConfigCacheRefreshInterval = previousRefreshInterval }) + store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - defer c.stop() + t.Cleanup(c.stop) - cfg := c.get() - if cfg != nil { - t.Fatalf("config = %+v, want nil before lazy load", cfg) - } + // Construction must not start a polling loop or touch the store. + time.Sleep(20 * time.Millisecond) if got := store.versionCalls.Load(); got != 0 { t.Fatalf("versionCalls = %d, want 0", got) } @@ -68,9 +71,9 @@ func TestQuotaConfigCacheLazyLoad(t *testing.T) { t.Fatalf("configCalls = %d, want 0", got) } - cfg = c.load(context.Background()) + cfg := c.get(context.Background()) if cfg == nil { - t.Fatal("config is nil after lazy load") + t.Fatal("config is nil after first access") } if cfg.MaxStorageBytes != 1000 { t.Fatalf("MaxStorageBytes = %d, want 1000", cfg.MaxStorageBytes) @@ -86,118 +89,150 @@ func TestQuotaConfigCacheLazyLoad(t *testing.T) { } } -func TestQuotaConfigCacheRefreshFailOpenOnVersionError(t *testing.T) { +func TestQuotaConfigCacheReturnsDefensiveCopy(t *testing.T) { store := newCacheTestStore() - store.versionErr = context.DeadlineExceeded + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - defer c.stop() + t.Cleanup(c.stop) - c.refresh(context.Background()) - if cfg := c.get(); cfg != nil { - t.Fatalf("config = %+v, want nil", cfg) - } - if got := store.versionCalls.Load(); got != 1 { - t.Fatalf("versionCalls = %d, want 1", got) - } - if got := store.configCalls.Load(); got != 0 { - t.Fatalf("configCalls = %d, want 0", got) + cfg := c.get(context.Background()) + if cfg == nil { + t.Fatal("config is nil") } - if got := store.usageCalls.Load(); got != 0 { - t.Fatalf("usageCalls = %d, want 0", got) + cfg.MaxStorageBytes = 2000 + + cached := c.get(context.Background()) + if cached == nil || cached.MaxStorageBytes != 1000 { + t.Fatalf("cached config = %+v, want storage 1000", cached) } } -func TestQuotaConfigCacheRefreshOnlyLoadsConfigWhenVersionChanges(t *testing.T) { +func TestQuotaConfigCacheReusesSnapshotUntilTTLExpires(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - defer c.stop() + t.Cleanup(c.stop) - c.refresh(context.Background()) - if got := store.versionCalls.Load(); got != 1 { - t.Fatalf("versionCalls = %d, want 1", got) + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Fatalf("first config = %+v, want storage 1000", cfg) + } + store.mu.Lock() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 2000} + store.mu.Unlock() + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Fatalf("cached config = %+v, want storage 1000", cfg) } if got := store.configCalls.Load(); got != 1 { - t.Fatalf("configCalls = %d, want 1", got) + t.Fatalf("configCalls before expiry = %d, want 1", got) } - c.refresh(context.Background()) - if got := store.versionCalls.Load(); got != 2 { - t.Fatalf("versionCalls = %d, want 2", got) + c.mu.Lock() + c.nextRefresh = time.Time{} + c.mu.Unlock() + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 2000 { + t.Fatalf("refreshed config = %+v, want storage 2000", cfg) } - if got := store.configCalls.Load(); got != 1 { - t.Fatalf("configCalls = %d, want 1", got) + if got := store.configCalls.Load(); got != 2 { + t.Fatalf("configCalls after expiry = %d, want 2", got) } + if got := store.versionCalls.Load(); got != 0 { + t.Fatalf("versionCalls = %d, want 0", got) + } +} - store.mu.Lock() - store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 2000} - store.mu.Unlock() - c.refresh(context.Background()) +func TestQuotaConfigCacheCoalescesConcurrentExpiredAccess(t *testing.T) { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + c := newQuotaConfigCache("t1", "", store) + t.Cleanup(c.stop) + if cfg := c.get(context.Background()); cfg == nil { + t.Fatal("initial config is nil") + } + c.mu.Lock() + c.nextRefresh = time.Time{} + c.mu.Unlock() - cfg := c.get() - if cfg == nil { - t.Fatal("config is nil") + started := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + store.configHook = func() { + once.Do(func() { close(started) }) + <-release } - if cfg.MaxStorageBytes != 2000 { - t.Fatalf("MaxStorageBytes = %d, want 2000", cfg.MaxStorageBytes) + const callers = 32 + results := make(chan *QuotaConfigView, callers) + for range callers { + go func() { results <- c.get(context.Background()) }() } - if got := store.versionCalls.Load(); got != 3 { - t.Fatalf("versionCalls = %d, want 3", got) + <-started + close(release) + for range callers { + if cfg := <-results; cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Fatalf("concurrent config = %+v, want storage 1000", cfg) + } } if got := store.configCalls.Load(); got != 2 { - t.Fatalf("configCalls = %d, want 2", got) + t.Fatalf("configCalls = %d, want initial load plus one coalesced refresh", got) } - if got := store.usageCalls.Load(); got != 0 { - t.Fatalf("usageCalls = %d, want 0", got) + if got := store.versionCalls.Load(); got != 0 { + t.Fatalf("versionCalls = %d, want 0", got) } } -func TestQuotaConfigCacheLazyLoadDoesNotOverwriteRefreshedSnapshot(t *testing.T) { +func TestQuotaConfigCacheRefreshFailureReturnsStaleAndUsesRetryCooldown(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - defer c.stop() + t.Cleanup(c.stop) - store.configHook = func() { - c.mu.Lock() - c.snapshot = "aConfigSnapshot{ - config: &QuotaConfigView{MaxStorageBytes: 2000}, - version: "new-version", - } - c.mu.Unlock() + first := c.get(context.Background()) + if first == nil { + t.Fatal("initial config is nil") } + c.mu.Lock() + c.nextRefresh = time.Time{} + c.mu.Unlock() + store.configErr = errors.New("temporary metadb failure") - cfg := c.load(context.Background()) - if cfg == nil { - t.Fatal("config is nil") + stale := c.get(context.Background()) + if stale == nil || stale.MaxStorageBytes != 1000 { + t.Fatalf("stale config = %+v, want storage 1000", stale) } - if cfg.MaxStorageBytes != 2000 { - t.Fatalf("lazy load config = %d, want refreshed 2000", cfg.MaxStorageBytes) + if got := store.configCalls.Load(); got != 2 { + t.Fatalf("configCalls after failed refresh = %d, want 2", got) } - cached := c.get() - if cached == nil || cached.MaxStorageBytes != 2000 { - t.Fatalf("cached config = %+v, want refreshed 2000", cached) + stale = c.get(context.Background()) + if stale == nil || stale.MaxStorageBytes != 1000 { + t.Fatalf("cooldown config = %+v, want storage 1000", stale) + } + if got := store.configCalls.Load(); got != 2 { + t.Fatalf("configCalls during failure cooldown = %d, want 2", got) + } + if got := store.versionCalls.Load(); got != 0 { + t.Fatalf("versionCalls = %d, want 0", got) } } -func TestQuotaConfigCacheLazyLoadReturnsDefensiveCopy(t *testing.T) { +func TestQuotaConfigCacheInitialFailureUsesRetryCooldown(t *testing.T) { store := newCacheTestStore() - store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + store.configErr = errors.New("temporary metadb failure") c := newQuotaConfigCache("t1", "", store) - defer c.stop() + t.Cleanup(c.stop) - cfg := c.load(context.Background()) - if cfg == nil { - t.Fatal("config is nil") + if cfg := c.get(context.Background()); cfg != nil { + t.Fatalf("config = %+v, want nil on initial failure", cfg) } - cfg.MaxStorageBytes = 2000 - - cached := c.get() - if cached == nil { - t.Fatal("cached config is nil") + if cfg := c.get(context.Background()); cfg != nil { + t.Fatalf("config during cooldown = %+v, want nil", cfg) + } + if got := store.configCalls.Load(); got != 1 { + t.Fatalf("configCalls during failure cooldown = %d, want 1", got) + } + if got := store.versionCalls.Load(); got != 0 { + t.Fatalf("versionCalls = %d, want 0", got) } - if cached.MaxStorageBytes != 1000 { - t.Fatalf("cached MaxStorageBytes = %d, want 1000", cached.MaxStorageBytes) + if got := store.usageCalls.Load(); got != 0 { + t.Fatalf("usageCalls = %d, want 0", got) } } diff --git a/pkg/backend/quota_integration_test.go b/pkg/backend/quota_integration_test.go index 562091d0..7fd263ec 100644 --- a/pkg/backend/quota_integration_test.go +++ b/pkg/backend/quota_integration_test.go @@ -29,6 +29,19 @@ func newServerQuotaBackend(t *testing.T, opts Options) (*Dat9Backend, *fakeMetaQ return b, fake } +func reloadQuotaConfigForTest(t *testing.T, b *Dat9Backend, ctx context.Context) { + t.Helper() + if b.quotaConfigCache == nil { + t.Fatal("quota config cache is nil") + } + b.quotaConfigCache.mu.Lock() + b.quotaConfigCache.nextRefresh = time.Time{} + b.quotaConfigCache.mu.Unlock() + if cfg := b.quotaConfigCache.get(ctx); cfg == nil { + t.Fatal("reload quota config returned nil") + } +} + func waitForFakeCentralLLMUsage(t *testing.T, fake *fakeMetaQuotaStore, tenantID string, wantMonthly int64, wantUsageLen int) { t.Helper() deadline := time.Now().Add(2 * time.Second) @@ -117,7 +130,7 @@ func TestServerQuotaRejectsOverFileSizeLimit(t *testing.T) { fake.mu.Lock() fake.config["tenant-a"].MaxFileSizeBytes = 4 fake.mu.Unlock() - b.quotaConfigCache.refresh(ctx) + reloadQuotaConfigForTest(t, b, ctx) if _, err := b.WriteCtx(ctx, "/too-large.txt", []byte("12345"), 0, filesystem.WriteFlagCreate); !errors.Is(err, ErrFileSizeQuotaExceeded) { t.Fatalf("write error = %v, want ErrFileSizeQuotaExceeded", err) @@ -141,7 +154,7 @@ func TestCreateIfAbsentExistingPathReturnsConflictBeforeFileSizeQuota(t *testing fake.mu.Lock() fake.config["tenant-a"].MaxFileSizeBytes = 4 fake.mu.Unlock() - b.quotaConfigCache.refresh(ctx) + reloadQuotaConfigForTest(t, b, ctx) if _, _, err := b.WriteCtxIfRevisionWithTagsResult(ctx, "/size-exists.txt", []byte("12345"), 0, filesystem.WriteFlagCreate|filesystem.WriteFlagTruncate, 0, nil, ""); !errors.Is(err, datastore.ErrRevisionConflict) { t.Fatalf("duplicate create-if-absent error = %v, want ErrRevisionConflict", err) @@ -161,7 +174,7 @@ func TestServerModeBudgetGateWritesCentralOnly(t *testing.T) { fake.mu.Lock() fake.config["tenant-a"].MaxMonthlyCostMC = 100 fake.mu.Unlock() - b.quotaConfigCache.refresh(context.Background()) + reloadQuotaConfigForTest(t, b, context.Background()) b.recordImageExtractUsage("task-server-budget", ImageExtractUsage{ PromptTokens: 120, diff --git a/pkg/tenant/pool.go b/pkg/tenant/pool.go index cbcf9c20..612ee7c3 100644 --- a/pkg/tenant/pool.go +++ b/pkg/tenant/pool.go @@ -735,6 +735,13 @@ func (p *Pool) reapOnce(ctx context.Context) { if e.retired || e.refs > 0 { continue } + // Shared tenant backends are lightweight fs_id-scoped views over a + // separately managed physical DB handle. Keep them warm until the + // capacity LRU evicts them; applying the standalone idle TTL here turns + // stable high-cardinality shared traffic into continuous cold opens. + if e.sharedDBID > 0 { + continue + } if now.Sub(e.lastUsed) > p.idleTimeout { if removed := p.removeLocked(e.elem, "idle"); removed != nil { toClose = append(toClose, removed) diff --git a/pkg/tenant/pool_test.go b/pkg/tenant/pool_test.go index 39a2ffdd..06e83226 100644 --- a/pkg/tenant/pool_test.go +++ b/pkg/tenant/pool_test.go @@ -1187,6 +1187,62 @@ func TestIdleEviction(t *testing.T) { } } +func TestIdleEvictionSkipsSharedTenantEntry(t *testing.T) { + pool, tenant := newTestPoolAndTenantWithConfig(t, PoolConfig{ + MaxTenants: 2, + IdleTimeout: time.Minute, + }, "tenant-shared-idle") + ctx := context.Background() + + b, release, err := pool.Acquire(ctx, tenant) + if err != nil { + t.Fatal(err) + } + store := b.Store() + release() + + pool.mu.Lock() + e := pool.items[tenant.ID] + e.sharedDBID = 123 + e.lastUsed = time.Now().Add(-2 * time.Minute) + pool.mu.Unlock() + + pool.reapOnce(ctx) + + assertStoreOpen(t, store) + if _, ok := pool.items[tenant.ID]; !ok { + t.Fatal("shared tenant entry was removed by idle reaper") + } +} + +func TestCapacityEvictionStillRemovesSharedTenantEntry(t *testing.T) { + pool, first := newTestPoolAndTenant(t, 1, "tenant-shared-capacity-first") + ctx := context.Background() + + b, release, err := pool.Acquire(ctx, first) + if err != nil { + t.Fatal(err) + } + firstStore := b.Store() + release() + + pool.mu.Lock() + pool.items[first.ID].sharedDBID = 123 + pool.mu.Unlock() + + second := cloneTenantForID(t, pool, first, "tenant-shared-capacity-second") + _, releaseSecond, err := pool.Acquire(ctx, second) + if err != nil { + t.Fatal(err) + } + releaseSecond() + + assertStoreClosed(t, firstStore) + if _, ok := pool.items[first.ID]; ok { + t.Fatal("shared tenant entry remained after capacity eviction") + } +} + func TestIdleEvictionSkippedByRecentAcquire(t *testing.T) { pool, tenant := newTestPoolAndTenantWithConfig(t, PoolConfig{ MaxTenants: 2, From 5a998767b59e985b501bf5cd4ab7dfd97826210a Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 02:23:15 +0800 Subject: [PATCH 02/11] fix: harden lazy quota config refresh --- cmd/drive9-server/main_test.go | 25 +++++- pkg/backend/quota_cache.go | 39 +++++++--- pkg/backend/quota_cache_test.go | 133 ++++++++++++++++++++++++-------- 3 files changed, 151 insertions(+), 46 deletions(-) diff --git a/cmd/drive9-server/main_test.go b/cmd/drive9-server/main_test.go index 99c6c0ae..7fba3188 100644 --- a/cmd/drive9-server/main_test.go +++ b/cmd/drive9-server/main_test.go @@ -37,7 +37,7 @@ func TestTenantBackendCacheMaxTenantsUsesProviderDefault(t *testing.T) { for _, tt := range tests { t.Run(tt.provider, func(t *testing.T) { if got := tenantBackendCacheMaxTenants(tt.provider); got != tt.want { - t.Fatalf("tenantBackendCacheMaxTenants(%q) = %d, want %d", tt.provider, got, tt.want) + t.Errorf("tenantBackendCacheMaxTenants(%q) = %d, want %d", tt.provider, got, tt.want) } }) } @@ -56,11 +56,32 @@ func TestTenantBackendCacheMaxTenantsExplicitOverrideWins(t *testing.T) { tenant.ProviderDB9, } { if got := tenantBackendCacheMaxTenants(provider); got != 4096 { - t.Fatalf("tenantBackendCacheMaxTenants(%q) = %d, want explicit 4096", provider, got) + t.Errorf("tenantBackendCacheMaxTenants(%q) = %d, want explicit 4096", provider, got) } } } +func TestTenantBackendCacheMaxTenantsFallsBackForInvalidOverrides(t *testing.T) { + const key = "DRIVE9_POOL_MAX_TENANTS" + restore := snapshotEnv(t, []string{key}) + t.Cleanup(func() { restoreEnv(t, restore) }) + + for _, raw := range []string{"0", "-1", "bad"} { + setEnv(t, key, raw) + if got := tenantBackendCacheMaxTenants(tenant.ProviderTiDBCloudNativeShared); got != 20480 { + t.Errorf("tenantBackendCacheMaxTenants(shared) with %q = %d, want 20480", raw, got) + } + if got := tenantBackendCacheMaxTenants(tenant.ProviderTiDBCloudNative); got != 1024 { + t.Errorf("tenantBackendCacheMaxTenants(native) with %q = %d, want 1024", raw, got) + } + } + + unsetEnv(t, []string{key}) + if got := tenantBackendCacheMaxTenants("unknown-provider"); got != 1024 { + t.Errorf("tenantBackendCacheMaxTenants(unknown-provider) = %d, want 1024", got) + } +} + func TestSlockOAuthFromEnvDisabledByDefault(t *testing.T) { keys := []string{ "DRIVE9_SLOCK_ORIGIN", diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index fdb3a291..3f95a18d 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -2,6 +2,7 @@ package backend import ( "context" + "math/rand/v2" "sync" "time" @@ -15,9 +16,9 @@ const ( // defaultQuotaConfigCacheRefreshInterval is the default TTL for lazily // loaded tenant quota config. Override with DRIVE9_QUOTA_CACHE_REFRESH_SECONDS. defaultQuotaConfigCacheRefreshInterval = 30 * time.Second - // quotaConfigCacheFailureRetryInterval prevents a central DB outage from - // turning every tenant request into another config query. - quotaConfigCacheFailureRetryInterval = 5 * time.Second + // quotaConfigCacheLoadTimeout bounds a coalesced refresh independently of + // the request that happened to win the load lock. + quotaConfigCacheLoadTimeout = 5 * time.Second // quotaUsageCacheTTL bounds how long soft small-write quota checks may // reuse central usage counters. Strict upload reservations still read // central usage directly. @@ -70,6 +71,20 @@ func cloneQuotaConfigView(cfg *QuotaConfigView) *QuotaConfigView { return &cp } +// quotaConfigCacheRefreshDelay spreads refreshes over the last 10% of the +// configured TTL. The delay never exceeds the configured value, preserving +// the maximum quota-config convergence window. +func quotaConfigCacheRefreshDelay(ttl time.Duration) time.Duration { + if ttl <= 0 { + return 0 + } + maxJitter := ttl / 10 + if maxJitter <= 0 { + return ttl + } + return ttl - time.Duration(rand.Int64N(int64(maxJitter)+1)) +} + // quotaConfigCache is a passive per-tenant cache for low-frequency quota // config. Requests refresh an expired snapshot; idle tenants create no // goroutines and issue no quota queries. @@ -112,10 +127,6 @@ func (c *quotaConfigCache) cached(now time.Time) (*QuotaConfigView, bool) { // get returns a defensive copy of the cached config, refreshing it once when // its TTL has expired. Concurrent expired requests coalesce on loadMu. func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { - return c.load(ctx) -} - -func (c *quotaConfigCache) load(ctx context.Context) *QuotaConfigView { now := time.Now() if cfg, current := c.cached(now); current { return cfg @@ -128,9 +139,11 @@ func (c *quotaConfigCache) load(ctx context.Context) *QuotaConfigView { } start := now - cfg, err := c.store.GetQuotaConfig(ctx, c.tenantID) + loadCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), quotaConfigCacheLoadTimeout) + defer cancel() + cfg, err := c.store.GetQuotaConfig(loadCtx, c.tenantID) if err != nil { - logger.Warn(ctx, "quota_config_cache_config_failed", + logger.Warn(loadCtx, "quota_config_cache_config_failed", zap.String("tenant_id", c.tenantID), zap.Error(err)) return c.finishFailedLoad(start, "config_error") } @@ -139,15 +152,19 @@ func (c *quotaConfigCache) load(ctx context.Context) *QuotaConfigView { } c.mu.Lock() c.snapshot = cloneQuotaConfigView(cfg) - c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshInterval) + c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshDelay(quotaConfigCacheRefreshInterval)) c.mu.Unlock() metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "ok", time.Since(start)) return cloneQuotaConfigView(cfg) } func (c *quotaConfigCache) finishFailedLoad(start time.Time, result string) *QuotaConfigView { + retryTTL := quotaConfigCacheRefreshInterval + if retryTTL < defaultQuotaConfigCacheRefreshInterval { + retryTTL = defaultQuotaConfigCacheRefreshInterval + } c.mu.Lock() - c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryInterval) + c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshDelay(retryTTL)) var stale *QuotaConfigView if c.snapshot != nil { stale = cloneQuotaConfigView(c.snapshot) diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 776c740e..8158b989 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -12,13 +12,14 @@ import ( // cacheTestStore wraps fakeMetaQuotaStore with error injection for cache tests. type cacheTestStore struct { *fakeMetaQuotaStore - configCalls atomic.Int64 - versionCalls atomic.Int64 - usageCalls atomic.Int64 - versionErr error - configErr error - configHook func() - usageHook func() + configCalls atomic.Int64 + versionCalls atomic.Int64 + usageCalls atomic.Int64 + versionErr error + configErr error + configHook func() + configCtxHook func(context.Context) + usageHook func() } func newCacheTestStore() *cacheTestStore { @@ -35,6 +36,12 @@ func (m *cacheTestStore) GetQuotaUsage(ctx context.Context, tenantID string) (*Q func (m *cacheTestStore) GetQuotaConfig(ctx context.Context, tenantID string) (*QuotaConfigView, error) { m.configCalls.Add(1) + if m.configCtxHook != nil { + m.configCtxHook(ctx) + } + if err := ctx.Err(); err != nil { + return nil, err + } if m.configErr != nil { return nil, m.configErr } @@ -44,6 +51,66 @@ func (m *cacheTestStore) GetQuotaConfig(ctx context.Context, tenantID string) (* return m.fakeMetaQuotaStore.GetQuotaConfig(ctx, tenantID) } +func TestQuotaConfigCacheCanceledCallerDoesNotPoisonSharedRefresh(t *testing.T) { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + loadHasDeadline := false + store.configCtxHook = func(ctx context.Context) { + _, loadHasDeadline = ctx.Deadline() + } + c := newQuotaConfigCache("t1", "", store) + t.Cleanup(c.stop) + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + first := c.get(canceledCtx) + if first == nil || first.MaxStorageBytes != 1000 { + t.Errorf("first config = %+v, want storage 1000", first) + } + if !loadHasDeadline { + t.Error("shared load context has no deadline") + } + second := c.get(context.Background()) + if second == nil || second.MaxStorageBytes != 1000 { + t.Errorf("second config = %+v, want storage 1000", second) + } + if got := store.configCalls.Load(); got != 1 { + t.Errorf("configCalls = %d, want 1 shared load", got) + } +} + +func TestQuotaConfigCacheRefreshDelayStaysWithinTTL(t *testing.T) { + const ttl = 30 * time.Second + for range 1000 { + delay := quotaConfigCacheRefreshDelay(ttl) + if delay < 27*time.Second || delay > ttl { + t.Errorf("refresh delay = %s, want [27s, 30s]", delay) + } + } +} + +func TestQuotaConfigCacheFailureCooldownUsesNormalTTLMinimum(t *testing.T) { + previousRefreshInterval := quotaConfigCacheRefreshInterval + quotaConfigCacheRefreshInterval = 5 * time.Second + t.Cleanup(func() { quotaConfigCacheRefreshInterval = previousRefreshInterval }) + + store := newCacheTestStore() + store.configErr = errors.New("temporary metadb failure") + c := newQuotaConfigCache("t1", "", store) + t.Cleanup(c.stop) + + before := time.Now() + if cfg := c.get(context.Background()); cfg != nil { + t.Errorf("config = %+v, want nil", cfg) + } + c.mu.RLock() + nextRefresh := c.nextRefresh + c.mu.RUnlock() + if delay := nextRefresh.Sub(before); delay < 27*time.Second || delay > 31*time.Second { + t.Errorf("failure cooldown = %s, want approximately [27s, 30s]", delay) + } +} + func (m *cacheTestStore) GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) { m.versionCalls.Add(1) if m.versionErr != nil { @@ -65,10 +132,10 @@ func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { // Construction must not start a polling loop or touch the store. time.Sleep(20 * time.Millisecond) if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } if got := store.configCalls.Load(); got != 0 { - t.Fatalf("configCalls = %d, want 0", got) + t.Errorf("configCalls = %d, want 0", got) } cfg := c.get(context.Background()) @@ -76,16 +143,16 @@ func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { t.Fatal("config is nil after first access") } if cfg.MaxStorageBytes != 1000 { - t.Fatalf("MaxStorageBytes = %d, want 1000", cfg.MaxStorageBytes) + t.Errorf("MaxStorageBytes = %d, want 1000", cfg.MaxStorageBytes) } if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } if got := store.configCalls.Load(); got != 1 { - t.Fatalf("configCalls = %d, want 1", got) + t.Errorf("configCalls = %d, want 1", got) } if got := store.usageCalls.Load(); got != 0 { - t.Fatalf("usageCalls = %d, want 0", got) + t.Errorf("usageCalls = %d, want 0", got) } } @@ -103,7 +170,7 @@ func TestQuotaConfigCacheReturnsDefensiveCopy(t *testing.T) { cached := c.get(context.Background()) if cached == nil || cached.MaxStorageBytes != 1000 { - t.Fatalf("cached config = %+v, want storage 1000", cached) + t.Errorf("cached config = %+v, want storage 1000", cached) } } @@ -114,29 +181,29 @@ func TestQuotaConfigCacheReusesSnapshotUntilTTLExpires(t *testing.T) { t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { - t.Fatalf("first config = %+v, want storage 1000", cfg) + t.Errorf("first config = %+v, want storage 1000", cfg) } store.mu.Lock() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 2000} store.mu.Unlock() if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { - t.Fatalf("cached config = %+v, want storage 1000", cfg) + t.Errorf("cached config = %+v, want storage 1000", cfg) } if got := store.configCalls.Load(); got != 1 { - t.Fatalf("configCalls before expiry = %d, want 1", got) + t.Errorf("configCalls before expiry = %d, want 1", got) } c.mu.Lock() c.nextRefresh = time.Time{} c.mu.Unlock() if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 2000 { - t.Fatalf("refreshed config = %+v, want storage 2000", cfg) + t.Errorf("refreshed config = %+v, want storage 2000", cfg) } if got := store.configCalls.Load(); got != 2 { - t.Fatalf("configCalls after expiry = %d, want 2", got) + t.Errorf("configCalls after expiry = %d, want 2", got) } if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } } @@ -168,14 +235,14 @@ func TestQuotaConfigCacheCoalescesConcurrentExpiredAccess(t *testing.T) { close(release) for range callers { if cfg := <-results; cfg == nil || cfg.MaxStorageBytes != 1000 { - t.Fatalf("concurrent config = %+v, want storage 1000", cfg) + t.Errorf("concurrent config = %+v, want storage 1000", cfg) } } if got := store.configCalls.Load(); got != 2 { - t.Fatalf("configCalls = %d, want initial load plus one coalesced refresh", got) + t.Errorf("configCalls = %d, want initial load plus one coalesced refresh", got) } if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } } @@ -196,20 +263,20 @@ func TestQuotaConfigCacheRefreshFailureReturnsStaleAndUsesRetryCooldown(t *testi stale := c.get(context.Background()) if stale == nil || stale.MaxStorageBytes != 1000 { - t.Fatalf("stale config = %+v, want storage 1000", stale) + t.Errorf("stale config = %+v, want storage 1000", stale) } if got := store.configCalls.Load(); got != 2 { - t.Fatalf("configCalls after failed refresh = %d, want 2", got) + t.Errorf("configCalls after failed refresh = %d, want 2", got) } stale = c.get(context.Background()) if stale == nil || stale.MaxStorageBytes != 1000 { - t.Fatalf("cooldown config = %+v, want storage 1000", stale) + t.Errorf("cooldown config = %+v, want storage 1000", stale) } if got := store.configCalls.Load(); got != 2 { - t.Fatalf("configCalls during failure cooldown = %d, want 2", got) + t.Errorf("configCalls during failure cooldown = %d, want 2", got) } if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } } @@ -220,19 +287,19 @@ func TestQuotaConfigCacheInitialFailureUsesRetryCooldown(t *testing.T) { t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg != nil { - t.Fatalf("config = %+v, want nil on initial failure", cfg) + t.Errorf("config = %+v, want nil on initial failure", cfg) } if cfg := c.get(context.Background()); cfg != nil { - t.Fatalf("config during cooldown = %+v, want nil", cfg) + t.Errorf("config during cooldown = %+v, want nil", cfg) } if got := store.configCalls.Load(); got != 1 { - t.Fatalf("configCalls during failure cooldown = %d, want 1", got) + t.Errorf("configCalls during failure cooldown = %d, want 1", got) } if got := store.versionCalls.Load(); got != 0 { - t.Fatalf("versionCalls = %d, want 0", got) + t.Errorf("versionCalls = %d, want 0", got) } if got := store.usageCalls.Load(); got != 0 { - t.Fatalf("usageCalls = %d, want 0", got) + t.Errorf("usageCalls = %d, want 0", got) } } From 053a7011160eafe65f770bf6ee9ad86d68917e1b Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 02:30:22 +0800 Subject: [PATCH 03/11] test: preserve empty environment snapshots --- cmd/drive9-server/main_test.go | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/cmd/drive9-server/main_test.go b/cmd/drive9-server/main_test.go index 7fba3188..62432fd4 100644 --- a/cmd/drive9-server/main_test.go +++ b/cmd/drive9-server/main_test.go @@ -82,6 +82,19 @@ func TestTenantBackendCacheMaxTenantsFallsBackForInvalidOverrides(t *testing.T) } } +func TestRestoreEnvPreservesPresentEmptyValue(t *testing.T) { + const key = "DRIVE9_TEST_PRESENT_EMPTY_ENV" + t.Setenv(key, "") + snapshot := snapshotEnv(t, []string{key}) + unsetEnv(t, []string{key}) + + restoreEnv(t, snapshot) + value, present := os.LookupEnv(key) + if !present || value != "" { + t.Errorf("restored env = (%q, %t), want present empty value", value, present) + } +} + func TestSlockOAuthFromEnvDisabledByDefault(t *testing.T) { keys := []string{ "DRIVE9_SLOCK_ORIGIN", @@ -746,25 +759,31 @@ func TestS3ConfigValidateRejectsInvalidStaticCredentialCombinations(t *testing.T } } -func snapshotEnv(t *testing.T, keys []string) map[string]string { +type envSnapshotValue struct { + value string + present bool +} + +func snapshotEnv(t *testing.T, keys []string) map[string]envSnapshotValue { t.Helper() - out := make(map[string]string, len(keys)) + out := make(map[string]envSnapshotValue, len(keys)) for _, key := range keys { - out[key] = os.Getenv(key) + value, present := os.LookupEnv(key) + out[key] = envSnapshotValue{value: value, present: present} } return out } -func restoreEnv(t *testing.T, snapshot map[string]string) { +func restoreEnv(t *testing.T, snapshot map[string]envSnapshotValue) { t.Helper() - for key, value := range snapshot { - if value == "" { + for key, entry := range snapshot { + if !entry.present { if err := os.Unsetenv(key); err != nil { t.Fatalf("unset %s: %v", key, err) } continue } - if err := os.Setenv(key, value); err != nil { + if err := os.Setenv(key, entry.value); err != nil { t.Fatalf("restore %s: %v", key, err) } } From 5fe0d3672033b512458ea61c4b10e16109acf2ae Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 03:23:50 +0800 Subject: [PATCH 04/11] fix: keep quota refresh waiters responsive --- pkg/backend/quota_cache.go | 53 ++++++++++++++---- pkg/backend/quota_cache_test.go | 97 ++++++++++++++++++++++++++++++--- 2 files changed, 132 insertions(+), 18 deletions(-) diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index 3f95a18d..d30fc4ab 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -17,8 +17,11 @@ const ( // loaded tenant quota config. Override with DRIVE9_QUOTA_CACHE_REFRESH_SECONDS. defaultQuotaConfigCacheRefreshInterval = 30 * time.Second // quotaConfigCacheLoadTimeout bounds a coalesced refresh independently of - // the request that happened to win the load lock. + // the request that happened to claim load ownership. quotaConfigCacheLoadTimeout = 5 * time.Second + // quotaConfigCacheFailureRetryInterval keeps quota changes responsive after + // a transient MetaDB failure. Successful refreshes retain the normal TTL. + quotaConfigCacheFailureRetryInterval = 5 * time.Second // quotaUsageCacheTTL bounds how long soft small-write quota checks may // reuse central usage counters. Strict upload reservations still read // central usage directly. @@ -96,7 +99,7 @@ type quotaConfigCache struct { mu sync.RWMutex snapshot *QuotaConfigView nextRefresh time.Time - loadMu sync.Mutex + loadDone chan struct{} } // newQuotaConfigCache creates an empty request-driven config cache. Backend @@ -125,18 +128,40 @@ func (c *quotaConfigCache) cached(now time.Time) (*QuotaConfigView, bool) { } // get returns a defensive copy of the cached config, refreshing it once when -// its TTL has expired. Concurrent expired requests coalesce on loadMu. +// its TTL has expired. Warm callers serve stale config while a refresh is in +// flight; cold callers wait only until their own context is done. func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { now := time.Now() if cfg, current := c.cached(now); current { return cfg } - c.loadMu.Lock() - defer c.loadMu.Unlock() + + c.mu.Lock() now = time.Now() - if cfg, current := c.cached(now); current { + if now.Before(c.nextRefresh) { + cfg := cloneQuotaConfigView(c.snapshot) + c.mu.Unlock() return cfg } + if c.loadDone != nil { + done := c.loadDone + stale := cloneQuotaConfigView(c.snapshot) + c.mu.Unlock() + if stale != nil { + return stale + } + select { + case <-done: + c.mu.RLock() + cfg := cloneQuotaConfigView(c.snapshot) + c.mu.RUnlock() + return cfg + case <-ctx.Done(): + return nil + } + } + c.loadDone = make(chan struct{}) + c.mu.Unlock() start := now loadCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), quotaConfigCacheLoadTimeout) @@ -153,27 +178,33 @@ func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { c.mu.Lock() c.snapshot = cloneQuotaConfigView(cfg) c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshDelay(quotaConfigCacheRefreshInterval)) + c.finishConfigLoadLocked() c.mu.Unlock() metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "ok", time.Since(start)) return cloneQuotaConfigView(cfg) } func (c *quotaConfigCache) finishFailedLoad(start time.Time, result string) *QuotaConfigView { - retryTTL := quotaConfigCacheRefreshInterval - if retryTTL < defaultQuotaConfigCacheRefreshInterval { - retryTTL = defaultQuotaConfigCacheRefreshInterval - } c.mu.Lock() - c.nextRefresh = time.Now().Add(quotaConfigCacheRefreshDelay(retryTTL)) + c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryInterval) var stale *QuotaConfigView if c.snapshot != nil { stale = cloneQuotaConfigView(c.snapshot) } + c.finishConfigLoadLocked() c.mu.Unlock() metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", result, time.Since(start)) return stale } +// finishConfigLoadLocked publishes the completed load before waking cold +// waiters. c.mu must be held. +func (c *quotaConfigCache) finishConfigLoadLocked() { + done := c.loadDone + c.loadDone = nil + close(done) +} + // stop remains for backend-close call-site compatibility. The cache owns no // goroutine or other lifecycle resource. func (c *quotaConfigCache) stop() {} diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 8158b989..a9ed1532 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -89,11 +89,7 @@ func TestQuotaConfigCacheRefreshDelayStaysWithinTTL(t *testing.T) { } } -func TestQuotaConfigCacheFailureCooldownUsesNormalTTLMinimum(t *testing.T) { - previousRefreshInterval := quotaConfigCacheRefreshInterval - quotaConfigCacheRefreshInterval = 5 * time.Second - t.Cleanup(func() { quotaConfigCacheRefreshInterval = previousRefreshInterval }) - +func TestQuotaConfigCacheFailureCooldownIsFiveSeconds(t *testing.T) { store := newCacheTestStore() store.configErr = errors.New("temporary metadb failure") c := newQuotaConfigCache("t1", "", store) @@ -106,8 +102,95 @@ func TestQuotaConfigCacheFailureCooldownUsesNormalTTLMinimum(t *testing.T) { c.mu.RLock() nextRefresh := c.nextRefresh c.mu.RUnlock() - if delay := nextRefresh.Sub(before); delay < 27*time.Second || delay > 31*time.Second { - t.Errorf("failure cooldown = %s, want approximately [27s, 30s]", delay) + if delay := nextRefresh.Sub(before); delay < 5*time.Second || delay > 6*time.Second { + t.Errorf("failure cooldown = %s, want approximately 5s", delay) + } +} + +func TestQuotaConfigCacheColdWaiterHonorsOwnDeadline(t *testing.T) { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + started := make(chan struct{}) + release := make(chan struct{}) + var startedOnce sync.Once + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(unblock) + store.configHook = func() { + startedOnce.Do(func() { close(started) }) + <-release + } + c := newQuotaConfigCache("t1", "", store) + t.Cleanup(c.stop) + + leaderDone := make(chan *QuotaConfigView, 1) + go func() { leaderDone <- c.get(context.Background()) }() + <-started + + waiterCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + waiterDone := make(chan *QuotaConfigView, 1) + go func() { waiterDone <- c.get(waiterCtx) }() + + select { + case cfg := <-waiterDone: + if cfg != nil { + t.Errorf("waiter config = %+v, want nil before cold load completes", cfg) + } + case <-time.After(300 * time.Millisecond): + t.Error("cold waiter remained blocked after its context deadline") + } + + unblock() + if cfg := <-leaderDone; cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("leader config = %+v, want storage 1000", cfg) + } +} + +func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + c := newQuotaConfigCache("t1", "", store) + t.Cleanup(c.stop) + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Fatalf("initial config = %+v, want storage 1000", cfg) + } + + c.mu.Lock() + c.nextRefresh = time.Time{} + c.mu.Unlock() + store.mu.Lock() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 2000} + store.mu.Unlock() + started := make(chan struct{}) + release := make(chan struct{}) + var startedOnce sync.Once + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(unblock) + store.configHook = func() { + startedOnce.Do(func() { close(started) }) + <-release + } + + leaderDone := make(chan *QuotaConfigView, 1) + go func() { leaderDone <- c.get(context.Background()) }() + <-started + waiterDone := make(chan *QuotaConfigView, 1) + go func() { waiterDone <- c.get(context.Background()) }() + + select { + case cfg := <-waiterDone: + if cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("waiter config = %+v, want stale storage 1000", cfg) + } + case <-time.After(300 * time.Millisecond): + t.Error("warm waiter blocked behind the in-flight refresh") + } + + unblock() + if cfg := <-leaderDone; cfg == nil || cfg.MaxStorageBytes != 2000 { + t.Errorf("leader config = %+v, want refreshed storage 2000", cfg) } } From fd2a813aefabe674d39c264789ae58d82c79359c Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 05:36:54 +0800 Subject: [PATCH 05/11] fix: preserve foreground tenant LRU order --- pkg/tenant/pool.go | 1 - pkg/tenant/pool_test.go | 56 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/pkg/tenant/pool.go b/pkg/tenant/pool.go index 612ee7c3..b061403b 100644 --- a/pkg/tenant/pool.go +++ b/pkg/tenant/pool.go @@ -568,7 +568,6 @@ func (p *Pool) AcquireCached(t *meta.Tenant) (b *backend.Dat9Backend, release fu return nil, nil, false } e.refs++ - p.order.MoveToFront(e.elem) p.mu.Unlock() metrics.RecordOperation("tenant_pool", "cache_lookup", "hit", 0) metrics.RecordTenantOperationWithOrg(t.ID, e.backend.TiDBCloudOrgID(), "user_db_access", "acquire_cached", "hit", 0) diff --git a/pkg/tenant/pool_test.go b/pkg/tenant/pool_test.go index 06e83226..18a41d2a 100644 --- a/pkg/tenant/pool_test.go +++ b/pkg/tenant/pool_test.go @@ -1243,6 +1243,62 @@ func TestCapacityEvictionStillRemovesSharedTenantEntry(t *testing.T) { } } +func TestAcquireCachedDoesNotRefreshSharedCapacityLRU(t *testing.T) { + pool, busy := newTestPoolAndTenant(t, 2, "tenant-shared-lru-busy") + ctx := context.Background() + + busyBackend, releaseBusy, err := pool.Acquire(ctx, busy) + if err != nil { + t.Fatal(err) + } + busyStore := busyBackend.Store() + releaseBusy() + + idle := cloneTenantForID(t, pool, busy, "tenant-shared-lru-idle") + idleBackend, releaseIdle, err := pool.Acquire(ctx, idle) + if err != nil { + t.Fatal(err) + } + idleStore := idleBackend.Store() + releaseIdle() + + pool.mu.Lock() + pool.items[busy.ID].sharedDBID = 123 + pool.items[idle.ID].sharedDBID = 123 + pool.mu.Unlock() + + // Foreground traffic makes busy the most recently used entry. + _, releaseBusy, err = pool.Acquire(ctx, busy) + if err != nil { + t.Fatal(err) + } + releaseBusy() + + // The safety-net scan may pin idle, but must not rewrite capacity LRU order. + _, releaseCached, ok := pool.AcquireCached(idle) + if !ok { + t.Fatal("expected AcquireCached to hit idle shared tenant") + } + releaseCached() + + third := cloneTenantForID(t, pool, busy, "tenant-shared-lru-third") + _, releaseThird, err := pool.Acquire(ctx, third) + if err != nil { + t.Fatal(err) + } + releaseThird() + + pool.mu.Lock() + _, busyCached := pool.items[busy.ID] + _, idleCached := pool.items[idle.ID] + pool.mu.Unlock() + if !busyCached || idleCached { + t.Errorf("capacity LRU after safety-net scan: busy cached=%t idle cached=%t, want true false", busyCached, idleCached) + } + assertStoreOpen(t, busyStore) + assertStoreClosed(t, idleStore) +} + func TestIdleEvictionSkippedByRecentAcquire(t *testing.T) { pool, tenant := newTestPoolAndTenantWithConfig(t, PoolConfig{ MaxTenants: 2, From c0b8625a6b177c77c3e140b773de0019694dcbbc Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 05:48:32 +0800 Subject: [PATCH 06/11] refactor: remove obsolete quota cache lifecycle --- pkg/backend/llm_usage_test.go | 4 --- pkg/backend/options.go | 5 ++- pkg/backend/quota_cache.go | 4 --- pkg/backend/quota_cache_test.go | 45 ------------------------- pkg/backend/quota_migration_test.go | 9 ----- pkg/backend/quota_store.go | 1 - pkg/meta/quota.go | 29 ---------------- pkg/meta/quota_setting_test.go | 52 +---------------------------- pkg/server/quota_test.go | 35 ++++++++++--------- pkg/tenant/quota_adapter.go | 4 --- 10 files changed, 20 insertions(+), 168 deletions(-) diff --git a/pkg/backend/llm_usage_test.go b/pkg/backend/llm_usage_test.go index 60ba1eda..0a9a49e8 100644 --- a/pkg/backend/llm_usage_test.go +++ b/pkg/backend/llm_usage_test.go @@ -28,10 +28,6 @@ func (m *mockQuotaStore) GetQuotaConfig(_ context.Context, _ string) (*QuotaConf return &QuotaConfigView{}, nil } -func (m *mockQuotaStore) GetQuotaConfigVersion(_ context.Context, _ string) (string, error) { - return "", nil -} - func (m *mockQuotaStore) GetQuotaUsage(_ context.Context, _ string) (*QuotaUsageView, error) { return &QuotaUsageView{}, nil } diff --git a/pkg/backend/options.go b/pkg/backend/options.go index 67e0a8af..73344bf9 100644 --- a/pkg/backend/options.go +++ b/pkg/backend/options.go @@ -21,10 +21,10 @@ const ( defaultMaxExtractedTextBytes = DefaultImageExtractMaxTextBytes defaultAudioExtractMaxSize = int64(32 << 20) // 32 MiB defaultAudioExtractTimeout = 2 * time.Minute - defaultMaxAudioExtractedTextBytes = 8 << 10 // 8 KiB + defaultMaxAudioExtractedTextBytes = 8 << 10 // 8 KiB defaultVideoExtractMaxSize = int64(200 << 20) // 200 MiB defaultVideoExtractTimeout = 5 * time.Minute - defaultMaxVideoExtractedTextBytes = 32 << 10 // 32 KiB + defaultMaxVideoExtractedTextBytes = 32 << 10 // 32 KiB defaultMaxUploadBytes = int64(10 * (1 << 30)) // 10 GiB defaultMaxTenantStorageBytes = int64(50 * (1 << 30)) // 50 GiB defaultMaxMediaLLMFiles = int64(500) // 500 media files per tenant @@ -393,7 +393,6 @@ func (b *Dat9Backend) Close() { } b.stopMutationWorker() if b.quotaConfigCache != nil { - b.quotaConfigCache.stop() b.quotaConfigCache = nil } } diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index d30fc4ab..0b54eadd 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -205,10 +205,6 @@ func (c *quotaConfigCache) finishConfigLoadLocked() { close(done) } -// stop remains for backend-close call-site compatibility. The cache owns no -// goroutine or other lifecycle resource. -func (c *quotaConfigCache) stop() {} - type quotaUsageSnapshot struct { usage *QuotaUsageView expiresAt time.Time diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index a9ed1532..3168b276 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -13,9 +13,7 @@ import ( type cacheTestStore struct { *fakeMetaQuotaStore configCalls atomic.Int64 - versionCalls atomic.Int64 usageCalls atomic.Int64 - versionErr error configErr error configHook func() configCtxHook func(context.Context) @@ -59,7 +57,6 @@ func TestQuotaConfigCacheCanceledCallerDoesNotPoisonSharedRefresh(t *testing.T) _, loadHasDeadline = ctx.Deadline() } c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) canceledCtx, cancel := context.WithCancel(context.Background()) cancel() @@ -93,7 +90,6 @@ func TestQuotaConfigCacheFailureCooldownIsFiveSeconds(t *testing.T) { store := newCacheTestStore() store.configErr = errors.New("temporary metadb failure") c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) before := time.Now() if cfg := c.get(context.Background()); cfg != nil { @@ -121,7 +117,6 @@ func TestQuotaConfigCacheColdWaiterHonorsOwnDeadline(t *testing.T) { <-release } c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) leaderDone := make(chan *QuotaConfigView, 1) go func() { leaderDone <- c.get(context.Background()) }() @@ -151,7 +146,6 @@ func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { t.Fatalf("initial config = %+v, want storage 1000", cfg) } @@ -194,14 +188,6 @@ func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { } } -func (m *cacheTestStore) GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) { - m.versionCalls.Add(1) - if m.versionErr != nil { - return "", m.versionErr - } - return m.fakeMetaQuotaStore.GetQuotaConfigVersion(ctx, tenantID) -} - func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { previousRefreshInterval := quotaConfigCacheRefreshInterval quotaConfigCacheRefreshInterval = 5 * time.Millisecond @@ -210,13 +196,9 @@ func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) // Construction must not start a polling loop or touch the store. time.Sleep(20 * time.Millisecond) - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } if got := store.configCalls.Load(); got != 0 { t.Errorf("configCalls = %d, want 0", got) } @@ -228,9 +210,6 @@ func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { if cfg.MaxStorageBytes != 1000 { t.Errorf("MaxStorageBytes = %d, want 1000", cfg.MaxStorageBytes) } - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } if got := store.configCalls.Load(); got != 1 { t.Errorf("configCalls = %d, want 1", got) } @@ -243,7 +222,6 @@ func TestQuotaConfigCacheReturnsDefensiveCopy(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) cfg := c.get(context.Background()) if cfg == nil { @@ -261,7 +239,6 @@ func TestQuotaConfigCacheReusesSnapshotUntilTTLExpires(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { t.Errorf("first config = %+v, want storage 1000", cfg) @@ -285,16 +262,12 @@ func TestQuotaConfigCacheReusesSnapshotUntilTTLExpires(t *testing.T) { if got := store.configCalls.Load(); got != 2 { t.Errorf("configCalls after expiry = %d, want 2", got) } - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } } func TestQuotaConfigCacheCoalescesConcurrentExpiredAccess(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg == nil { t.Fatal("initial config is nil") } @@ -324,16 +297,12 @@ func TestQuotaConfigCacheCoalescesConcurrentExpiredAccess(t *testing.T) { if got := store.configCalls.Load(); got != 2 { t.Errorf("configCalls = %d, want initial load plus one coalesced refresh", got) } - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } } func TestQuotaConfigCacheRefreshFailureReturnsStaleAndUsesRetryCooldown(t *testing.T) { store := newCacheTestStore() store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) first := c.get(context.Background()) if first == nil { @@ -358,16 +327,12 @@ func TestQuotaConfigCacheRefreshFailureReturnsStaleAndUsesRetryCooldown(t *testi if got := store.configCalls.Load(); got != 2 { t.Errorf("configCalls during failure cooldown = %d, want 2", got) } - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } } func TestQuotaConfigCacheInitialFailureUsesRetryCooldown(t *testing.T) { store := newCacheTestStore() store.configErr = errors.New("temporary metadb failure") c := newQuotaConfigCache("t1", "", store) - t.Cleanup(c.stop) if cfg := c.get(context.Background()); cfg != nil { t.Errorf("config = %+v, want nil on initial failure", cfg) @@ -378,9 +343,6 @@ func TestQuotaConfigCacheInitialFailureUsesRetryCooldown(t *testing.T) { if got := store.configCalls.Load(); got != 1 { t.Errorf("configCalls during failure cooldown = %d, want 1", got) } - if got := store.versionCalls.Load(); got != 0 { - t.Errorf("versionCalls = %d, want 0", got) - } if got := store.usageCalls.Load(); got != 0 { t.Errorf("usageCalls = %d, want 0", got) } @@ -679,10 +641,3 @@ func TestQuotaPendingDeltasCacheRemovesPositiveRaceDeltasOnClearAndExpire(t *tes t.Fatalf("positive deltas after expire = %+v, want zero", got) } } - -func TestQuotaConfigCacheStop(t *testing.T) { - store := newCacheTestStore() - c := newQuotaConfigCache("t1", "", store) - c.stop() - // Should not panic or block. -} diff --git a/pkg/backend/quota_migration_test.go b/pkg/backend/quota_migration_test.go index c27a662b..92cda863 100644 --- a/pkg/backend/quota_migration_test.go +++ b/pkg/backend/quota_migration_test.go @@ -121,15 +121,6 @@ func (f *fakeMetaQuotaStore) GetQuotaConfig(ctx context.Context, tenantID string }, nil } -func (f *fakeMetaQuotaStore) GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) { - f.mu.Lock() - defer f.mu.Unlock() - if cfg, ok := f.config[tenantID]; ok { - return fmt.Sprintf("v3:%d:%d:%d:%d:%d:%d", cfg.MaxStorageBytes, cfg.MaxFileSizeBytes, cfg.MaxFileCount, cfg.MaxMediaLLMFiles, cfg.MaxVideoLLMFiles, cfg.MaxMonthlyCostMC), nil - } - return "", nil -} - func (f *fakeMetaQuotaStore) GetQuotaUsage(ctx context.Context, tenantID string) (*QuotaUsageView, error) { f.mu.Lock() defer f.mu.Unlock() diff --git a/pkg/backend/quota_store.go b/pkg/backend/quota_store.go index 1d24df5c..90199ad3 100644 --- a/pkg/backend/quota_store.go +++ b/pkg/backend/quota_store.go @@ -21,7 +21,6 @@ const defaultTenantMetricTiDBCloudOrgID = "guest" type MetaQuotaStore interface { // Config GetQuotaConfig(ctx context.Context, tenantID string) (*QuotaConfigView, error) - GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) // Counters GetQuotaUsage(ctx context.Context, tenantID string) (*QuotaUsageView, error) diff --git a/pkg/meta/quota.go b/pkg/meta/quota.go index cba4310b..29e3e48a 100644 --- a/pkg/meta/quota.go +++ b/pkg/meta/quota.go @@ -169,35 +169,6 @@ func (s *Store) GetQuotaConfig(ctx context.Context, tenantID string) (*QuotaConf return cfg, nil } -// GetQuotaConfigVersion returns a lightweight content token for a tenant's -// explicit storage quota config. An empty token means no storage-relevant row -// exists and callers should use GetQuotaConfig's default storage config. The -// token is derived from the effective config values instead of updated_at so -// updates inside the same timestamp tick cannot hide a real config change from -// cache invalidation. -func (s *Store) GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) { - start := time.Now() - var err error - defer observeMeta(ctx, "get_quota_config_version", start, &err) - - var maxStorageBytes, maxFileSizeBytes, maxFileCount, maxMediaLLMFiles, maxVideoLLMFiles, maxMonthlyCostMC int64 - err = s.db.QueryRowContext(ctx, - `SELECT max_storage_bytes, max_file_size_bytes, max_file_count, - max_media_llm_files, max_video_llm_files, max_monthly_cost_mc - FROM tenant_quota_config WHERE tenant_id = ?`, tenantID, - ).Scan(&maxStorageBytes, &maxFileSizeBytes, &maxFileCount, &maxMediaLLMFiles, &maxVideoLLMFiles, &maxMonthlyCostMC) - if err == sql.ErrNoRows { - logger.Info(ctx, "quota_config_not_found_using_defaults", - zap.String("tenant_id", tenantID)) - err = nil - return "", nil - } - if err != nil { - return "", fmt.Errorf("get quota config version for tenant %q: %w", tenantID, err) - } - return fmt.Sprintf("v3:%d:%d:%d:%d:%d:%d", maxStorageBytes, maxFileSizeBytes, maxFileCount, maxMediaLLMFiles, maxVideoLLMFiles, maxMonthlyCostMC), nil -} - // SetQuotaConfig upserts per-tenant quota configuration. func (s *Store) SetQuotaConfig(ctx context.Context, cfg *QuotaConfig) error { start := time.Now() diff --git a/pkg/meta/quota_setting_test.go b/pkg/meta/quota_setting_test.go index a80ab36e..9865646b 100644 --- a/pkg/meta/quota_setting_test.go +++ b/pkg/meta/quota_setting_test.go @@ -87,7 +87,7 @@ func TestGetQuotaConfigUsesConfiguredDefaultStorageBytes(t *testing.T) { } } -func TestQuotaConfigStoresTiDBCloudSpendingLimitWithoutStorageVersion(t *testing.T) { +func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) { originalFileSizeDefault := DefaultMaxFileSizeBytes() defer SetDefaultMaxFileSizeBytes(originalFileSizeDefault) @@ -124,14 +124,6 @@ func TestQuotaConfigStoresTiDBCloudSpendingLimitWithoutStorageVersion(t *testing if cfg.MaxFileSizeBytes != originalFileSizeDefault { t.Fatalf("materialized MaxFileSizeBytes = %d after default changed, want %d", cfg.MaxFileSizeBytes, originalFileSizeDefault) } - version, err := s.GetQuotaConfigVersion(ctx, "tenant-spending-only") - if err != nil { - t.Fatal(err) - } - if version == "" { - t.Fatalf("storage quota version should be non-empty when config row exists") - } - updated := int64(123) if err := s.SetQuotaConfigPatch(ctx, "tenant-spending-only", QuotaConfigPatch{TiDBCloudSpendingLimit: &updated}); err != nil { t.Fatal(err) @@ -176,48 +168,6 @@ func TestGetQuotaConfigUsesDefaultFileSizeForExistingZeroRow(t *testing.T) { } } -func TestGetQuotaConfigVersion(t *testing.T) { - s := newControlStore(t) - ctx := context.Background() - - version, err := s.GetQuotaConfigVersion(ctx, "tenant-without-config") - if err != nil { - t.Fatal(err) - } - if version != "" { - t.Fatalf("version for missing config = %q, want empty", version) - } - - if err := s.SetQuotaConfig(ctx, &QuotaConfig{ - TenantID: "tenant-with-config", - MaxStorageBytes: 123, - MaxFileSizeBytes: 234, - MaxFileCount: 345, - MaxMediaLLMFiles: 456, - MaxVideoLLMFiles: 567, - MaxMonthlyCostMC: 789, - }); err != nil { - t.Fatal(err) - } - version, err = s.GetQuotaConfigVersion(ctx, "tenant-with-config") - if err != nil { - t.Fatal(err) - } - if version == "" { - t.Fatal("version for explicit config is empty") - } - if err := s.SetQuotaStorageBytes(ctx, "tenant-with-config", 321); err != nil { - t.Fatal(err) - } - nextVersion, err := s.GetQuotaConfigVersion(ctx, "tenant-with-config") - if err != nil { - t.Fatal(err) - } - if nextVersion == version { - t.Fatalf("version after config value change = %q, want different from %q", nextVersion, version) - } -} - func TestSetQuotaStorageBytesUpdatesStorageOnly(t *testing.T) { s := newControlStore(t) ctx := context.Background() diff --git a/pkg/server/quota_test.go b/pkg/server/quota_test.go index edacd140..961d27d4 100644 --- a/pkg/server/quota_test.go +++ b/pkg/server/quota_test.go @@ -408,6 +408,17 @@ func newQuotaRuntime(t *testing.T, provider string) *quotaRuntime { return newQuotaRuntimeWithOptions(t, provider, quotaRuntimeOptions{}) } +func quotaConfigRowExists(t *testing.T, rt *quotaRuntime) bool { + t.Helper() + var exists bool + if err := rt.meta.DB().QueryRowContext(context.Background(), + `SELECT EXISTS(SELECT 1 FROM tenant_quota_config WHERE tenant_id = ?)`, rt.tenantID, + ).Scan(&exists); err != nil { + t.Fatal(err) + } + return exists +} + func newQuotaRuntimeWithOptions(t *testing.T, provider string, opts quotaRuntimeOptions) *quotaRuntime { t.Helper() db := newTenantDeleteDBInfo(t) @@ -1105,12 +1116,8 @@ func TestQuotaSetSpendingLimitOnlyPersistsSpendingLimitConfig(t *testing.T) { if cfg.MaxStorageBytes != meta.DefaultMaxStorageBytes() || cfg.MaxFileSizeBytes != meta.DefaultMaxFileSizeBytes() || cfg.MaxFileCount != 0 { t.Fatalf("storage quota fields = %#v, want defaults", cfg) } - version, err := rt.meta.GetQuotaConfigVersion(context.Background(), rt.tenantID) - if err != nil { - t.Fatal(err) - } - if version == "" { - t.Fatalf("storage quota config version should be non-empty when config row exists") + if !quotaConfigRowExists(t, rt) { + t.Fatal("quota config row was not persisted") } var out quotaResponse if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { @@ -1140,12 +1147,8 @@ func TestQuotaSetRejectsDrive9KeyWithoutTiDBCloudCredentials(t *testing.T) { if got := rt.prov.markCalls.Load(); got != 0 { t.Fatalf("mark calls = %d, want 0", got) } - version, err := rt.meta.GetQuotaConfigVersion(context.Background(), rt.tenantID) - if err != nil { - t.Fatal(err) - } - if version != "" { - t.Fatalf("quota config version = %q, want empty", version) + if quotaConfigRowExists(t, rt) { + t.Fatal("quota config row was written for a rejected request") } } @@ -1366,12 +1369,8 @@ func TestQuotaSetMapsTiDBCloudCredentialErrorsWithoutWritingConfig(t *testing.T) if resp.StatusCode != tc.wantStatus { t.Fatalf("status = %d, want %d", resp.StatusCode, tc.wantStatus) } - version, err := rt.meta.GetQuotaConfigVersion(context.Background(), rt.tenantID) - if err != nil { - t.Fatal(err) - } - if version != "" { - t.Fatalf("quota config version = %q, want empty", version) + if quotaConfigRowExists(t, rt) { + t.Fatal("quota config row was written after TiDB Cloud credential failure") } if got := rt.prov.markCalls.Load(); got != 1 { t.Fatalf("mark calls = %d, want 1", got) diff --git a/pkg/tenant/quota_adapter.go b/pkg/tenant/quota_adapter.go index 139d53cf..8ce18f05 100644 --- a/pkg/tenant/quota_adapter.go +++ b/pkg/tenant/quota_adapter.go @@ -41,10 +41,6 @@ func (a *metaQuotaAdapter) GetQuotaConfig(ctx context.Context, tenantID string) }, nil } -func (a *metaQuotaAdapter) GetQuotaConfigVersion(ctx context.Context, tenantID string) (string, error) { - return a.s.GetQuotaConfigVersion(ctx, tenantID) -} - func (a *metaQuotaAdapter) GetQuotaUsage(ctx context.Context, tenantID string) (*backend.QuotaUsageView, error) { u, err := a.s.GetQuotaUsage(ctx, tenantID) if err != nil { From 9d69fea9f17dc9278e112f8278c20e76e084c863 Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 11:05:43 +0800 Subject: [PATCH 07/11] fix: harden quota cache refresh ownership --- pkg/backend/quota_cache.go | 58 +++++++++++++++++++++++++- pkg/backend/quota_cache_test.go | 73 +++++++++++++++++++++++++++++---- 2 files changed, 123 insertions(+), 8 deletions(-) diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index 0b54eadd..58b725ad 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -88,6 +88,21 @@ func quotaConfigCacheRefreshDelay(ttl time.Duration) time.Duration { return ttl - time.Duration(rand.Int64N(int64(maxJitter)+1)) } +// quotaConfigCacheFailureRetryDelay spreads failed loads symmetrically around +// the five-second retry target. Unlike successful refreshes, failures may be +// retried slightly after the target so independently warmed tenants do not +// converge into a synchronized retry wave. +func quotaConfigCacheFailureRetryDelay(ttl time.Duration) time.Duration { + if ttl <= 0 { + return 0 + } + jitter := ttl / 10 + if jitter <= 0 { + return ttl + } + return ttl - jitter + time.Duration(rand.Int64N(int64(2*jitter)+1)) +} + // quotaConfigCache is a passive per-tenant cache for low-frequency quota // config. Requests refresh an expired snapshot; idle tenants create no // goroutines and issue no quota queries. @@ -164,8 +179,37 @@ func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { c.mu.Unlock() start := now + if stale := c.snapshotCopy(); stale != nil { + go func() { + defer func() { + if recovered := recover(); recovered != nil { + logger.Error(context.Background(), "quota_config_cache_async_load_panicked", + zap.String("tenant_id", c.tenantID), zap.Any("panic", recovered)) + } + }() + c.loadConfig(context.WithoutCancel(ctx), start) + }() + return stale + } + return c.loadConfig(ctx, start) +} + +func (c *quotaConfigCache) snapshotCopy() *QuotaConfigView { + c.mu.RLock() + defer c.mu.RUnlock() + return cloneQuotaConfigView(c.snapshot) +} + +func (c *quotaConfigCache) loadConfig(ctx context.Context, start time.Time) (result *QuotaConfigView) { loadCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), quotaConfigCacheLoadTimeout) defer cancel() + defer func() { + if recovered := recover(); recovered != nil { + c.finishPanickedLoad(start) + panic(recovered) + } + }() + cfg, err := c.store.GetQuotaConfig(loadCtx, c.tenantID) if err != nil { logger.Warn(loadCtx, "quota_config_cache_config_failed", @@ -186,7 +230,7 @@ func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { func (c *quotaConfigCache) finishFailedLoad(start time.Time, result string) *QuotaConfigView { c.mu.Lock() - c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryInterval) + c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryDelay(quotaConfigCacheFailureRetryInterval)) var stale *QuotaConfigView if c.snapshot != nil { stale = cloneQuotaConfigView(c.snapshot) @@ -197,6 +241,18 @@ func (c *quotaConfigCache) finishFailedLoad(start time.Time, result string) *Quo return stale } +func (c *quotaConfigCache) finishPanickedLoad(start time.Time) { + c.mu.Lock() + if c.loadDone == nil { + c.mu.Unlock() + return + } + c.nextRefresh = time.Now().Add(quotaConfigCacheFailureRetryDelay(quotaConfigCacheFailureRetryInterval)) + c.finishConfigLoadLocked() + c.mu.Unlock() + metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "panic_error", time.Since(start)) +} + // finishConfigLoadLocked publishes the completed load before waking cold // waiters. c.mu must be held. func (c *quotaConfigCache) finishConfigLoadLocked() { diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 3168b276..8124c4fa 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -98,8 +98,50 @@ func TestQuotaConfigCacheFailureCooldownIsFiveSeconds(t *testing.T) { c.mu.RLock() nextRefresh := c.nextRefresh c.mu.RUnlock() - if delay := nextRefresh.Sub(before); delay < 5*time.Second || delay > 6*time.Second { - t.Errorf("failure cooldown = %s, want approximately 5s", delay) + if delay := nextRefresh.Sub(before); delay < 4500*time.Millisecond || delay > 5500*time.Millisecond { + t.Errorf("failure cooldown = %s, want [4.5s, 5.5s]", delay) + } +} + +func TestQuotaConfigCacheFailureRetryDelayUsesSymmetricJitter(t *testing.T) { + const base = 5 * time.Second + for range 2000 { + delay := quotaConfigCacheFailureRetryDelay(base) + if delay < 4500*time.Millisecond || delay > 5500*time.Millisecond { + t.Fatalf("failure retry delay = %s, want [4.5s, 5.5s]", delay) + } + } +} + +func TestQuotaConfigCachePanicDoesNotWedgeNextLoad(t *testing.T) { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + shouldPanic := true + store.configHook = func() { + if shouldPanic { + shouldPanic = false + panic("quota config store panic") + } + } + c := newQuotaConfigCache("t1", "", store) + + func() { + defer func() { + if recover() == nil { + t.Fatal("expected quota config load panic") + } + }() + _ = c.get(context.Background()) + }() + + c.mu.Lock() + c.nextRefresh = time.Time{} + c.mu.Unlock() + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Fatalf("healthy config after panic = %+v, want storage 1000", cfg) + } + if got := store.configCalls.Load(); got != 2 { + t.Fatalf("configCalls = %d, want panic load plus healthy reload", got) } } @@ -158,6 +200,7 @@ func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { store.mu.Unlock() started := make(chan struct{}) release := make(chan struct{}) + finished := make(chan struct{}) var startedOnce sync.Once var releaseOnce sync.Once unblock := func() { releaseOnce.Do(func() { close(release) }) } @@ -165,26 +208,42 @@ func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { store.configHook = func() { startedOnce.Do(func() { close(started) }) <-release + close(finished) } + leaderCtx, cancelLeader := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancelLeader() leaderDone := make(chan *QuotaConfigView, 1) - go func() { leaderDone <- c.get(context.Background()) }() + go func() { leaderDone <- c.get(leaderCtx) }() + select { + case cfg := <-leaderDone: + if cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("expired leader config = %+v, want stale storage 1000", cfg) + } + case <-time.After(100 * time.Millisecond): + t.Error("expired warm leader blocked behind the refresh") + } + <-started waiterDone := make(chan *QuotaConfigView, 1) go func() { waiterDone <- c.get(context.Background()) }() - select { case cfg := <-waiterDone: if cfg == nil || cfg.MaxStorageBytes != 1000 { - t.Errorf("waiter config = %+v, want stale storage 1000", cfg) + t.Errorf("warm waiter config = %+v, want stale storage 1000", cfg) } case <-time.After(300 * time.Millisecond): t.Error("warm waiter blocked behind the in-flight refresh") } unblock() - if cfg := <-leaderDone; cfg == nil || cfg.MaxStorageBytes != 2000 { - t.Errorf("leader config = %+v, want refreshed storage 2000", cfg) + select { + case <-finished: + case <-time.After(300 * time.Millisecond): + t.Fatal("async refresh did not finish") + } + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 2000 { + t.Errorf("refreshed config = %+v, want storage 2000", cfg) } } From 06591fd1121b5241cc745fca83f513b09074240c Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 11:10:20 +0800 Subject: [PATCH 08/11] fix: keep shared handle reaper independent --- pkg/tenant/pool.go | 62 +++++++++++++++++++++-------------------- pkg/tenant/pool_test.go | 23 +++++++++++++++ 2 files changed, 55 insertions(+), 30 deletions(-) diff --git a/pkg/tenant/pool.go b/pkg/tenant/pool.go index b061403b..78708e76 100644 --- a/pkg/tenant/pool.go +++ b/pkg/tenant/pool.go @@ -66,9 +66,11 @@ type PoolConfig struct { // IsLeader() snapshot. LeaderChecker LeaderChecker - // IdleTimeout controls how long a cached backend can stay in the warm - // cache without any activity before the idle reaper evicts it. 0 - // disables idle eviction (LRU capacity eviction still applies). + // IdleTimeout controls how long a standalone cached backend can stay in the + // warm cache without activity before the idle reaper evicts it. Shared + // tenant entries intentionally bypass this TTL and are reclaimed by the + // capacity LRU instead. 0 disables standalone idle eviction; the shared + // physical-DB handle reaper remains active. // // "Activity" means any access through Get, Acquire, or S3Backend — // including foreground user requests (HTTP/FUSE), tenant-specific @@ -83,8 +85,9 @@ type PoolConfig struct { // per due attempt. IdleTimeout time.Duration - // IdleReapInterval is how often the idle reaper scans for idle backends. - // Defaults to defaultTenantPoolIdleReapInterval when IdleTimeout > 0. + // IdleReapInterval is how often the reaper scans standalone entries and + // unreferenced shared physical-DB handles. It defaults to + // defaultTenantPoolIdleReapInterval when not set. IdleReapInterval time.Duration // SharedDBForcePlaintext, when true, opens shared-schema DB handles without @@ -248,7 +251,7 @@ func NewPool(cfg PoolConfig, enc encrypt.Encryptor) *Pool { } idleTimeout := cfg.IdleTimeout reapInterval := cfg.IdleReapInterval - if reapInterval <= 0 && idleTimeout > 0 { + if reapInterval <= 0 { reapInterval = defaultTenantPoolIdleReapInterval } metrics.RecordGauge("tenant_pool", "cached_backends", 0) @@ -691,11 +694,11 @@ func withTenantPoolDrainTimeout(ctx context.Context) (context.Context, context.C return context.WithTimeout(ctx, defaultTenantPoolDrainTimeout) } -// Start launches the idle reaper goroutine if IdleTimeout is configured. -// Safe to call on a pool with IdleTimeout=0 (no-op). The reaper is stopped -// by Close. +// Start launches the reaper goroutine. IdleTimeout=0 disables standalone +// tenant eviction, but shared physical-DB handles still need this reaper. +// The reaper is stopped by Close. func (p *Pool) Start(ctx context.Context) { - if p == nil || p.idleTimeout <= 0 { + if p == nil || p.reapInterval <= 0 { return } workerCtx, cancel := context.WithCancel(ctx) @@ -724,31 +727,30 @@ func (p *Pool) reapLoop(ctx context.Context) { } func (p *Pool) reapOnce(ctx context.Context) { - if p.idleTimeout <= 0 { - return - } now := time.Now() var toClose []*entry - p.mu.Lock() - for _, e := range p.items { - if e.retired || e.refs > 0 { - continue - } - // Shared tenant backends are lightweight fs_id-scoped views over a - // separately managed physical DB handle. Keep them warm until the - // capacity LRU evicts them; applying the standalone idle TTL here turns - // stable high-cardinality shared traffic into continuous cold opens. - if e.sharedDBID > 0 { - continue - } - if now.Sub(e.lastUsed) > p.idleTimeout { - if removed := p.removeLocked(e.elem, "idle"); removed != nil { - toClose = append(toClose, removed) + if p.idleTimeout > 0 { + p.mu.Lock() + for _, e := range p.items { + if e.retired || e.refs > 0 { + continue + } + // Shared tenant backends are lightweight fs_id-scoped views over a + // separately managed physical DB handle. Keep them warm until the + // capacity LRU evicts them; applying the standalone idle TTL here turns + // stable high-cardinality shared traffic into continuous cold opens. + if e.sharedDBID > 0 { + continue + } + if now.Sub(e.lastUsed) > p.idleTimeout { + if removed := p.removeLocked(e.elem, "idle"); removed != nil { + toClose = append(toClose, removed) + } } } + p.recordCachedBackendCountLocked() + p.mu.Unlock() } - p.recordCachedBackendCountLocked() - p.mu.Unlock() for _, retired := range toClose { p.closeEntry(retired) } diff --git a/pkg/tenant/pool_test.go b/pkg/tenant/pool_test.go index 18a41d2a..a1756aa4 100644 --- a/pkg/tenant/pool_test.go +++ b/pkg/tenant/pool_test.go @@ -182,6 +182,29 @@ func TestSharedDBHandleIdleReapUsesLongerTTLAndSkipsReferencedHandles(t *testing } } +func TestSharedDBHandleIdleReapRunsWhenTenantIdleEvictionDisabled(t *testing.T) { + p := NewPool(PoolConfig{IdleTimeout: 0}, nil) + t.Cleanup(p.Close) + db, err := sql.Open("mysql", testDSN) + if err != nil { + t.Fatal(err) + } + if err := db.Ping(); err != nil { + t.Fatal(err) + } + p.sharedDBs[1] = db + p.sharedDBLastUsed[1] = time.Now().Add(-defaultSharedDBHandleIdleTTL - time.Second) + + p.reapOnce(context.Background()) + + if _, ok := p.sharedDBs[1]; ok { + t.Fatal("expired unreferenced shared handle remains cached when tenant idle eviction is disabled") + } + if err := db.Ping(); err == nil { + t.Fatal("expired unreferenced shared handle remains open when tenant idle eviction is disabled") + } +} + func registerSharedDBForCacheMetrics(t *testing.T, orgID string) (*meta.Store, *Pool, int64, string) { t.Helper() metaStore, err := meta.OpenContext(context.Background(), testDSN) From 6f1570695550adc59680217f6412fce38e4a7265 Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 11:16:06 +0800 Subject: [PATCH 09/11] test: use nonfatal assertion failures --- pkg/meta/quota_setting_test.go | 12 ++++++------ pkg/server/quota_test.go | 16 ++++++++-------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pkg/meta/quota_setting_test.go b/pkg/meta/quota_setting_test.go index 9865646b..2b7cadaf 100644 --- a/pkg/meta/quota_setting_test.go +++ b/pkg/meta/quota_setting_test.go @@ -99,7 +99,7 @@ func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) t.Fatal(err) } if cfg.TiDBCloudSpendingLimit != nil { - t.Fatalf("default spending limit = %#v, want nil", cfg.TiDBCloudSpendingLimit) + t.Errorf("default spending limit = %#v, want nil", cfg.TiDBCloudSpendingLimit) } zero := int64(0) @@ -111,10 +111,10 @@ func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) t.Fatal(err) } if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != 0 { - t.Fatalf("spending limit = %#v, want 0", cfg.TiDBCloudSpendingLimit) + t.Errorf("spending limit = %#v, want 0", cfg.TiDBCloudSpendingLimit) } if cfg.MaxStorageBytes != DefaultMaxStorageBytes() || cfg.MaxFileSizeBytes != DefaultMaxFileSizeBytes() || cfg.MaxFileCount != 0 { - t.Fatalf("storage quota fields = %#v, want defaults", cfg) + t.Errorf("storage quota fields = %#v, want defaults", cfg) } SetDefaultMaxFileSizeBytes(originalFileSizeDefault + 1) cfg, err = s.GetQuotaConfig(ctx, "tenant-spending-only") @@ -122,7 +122,7 @@ func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) t.Fatal(err) } if cfg.MaxFileSizeBytes != originalFileSizeDefault { - t.Fatalf("materialized MaxFileSizeBytes = %d after default changed, want %d", cfg.MaxFileSizeBytes, originalFileSizeDefault) + t.Errorf("materialized MaxFileSizeBytes = %d after default changed, want %d", cfg.MaxFileSizeBytes, originalFileSizeDefault) } updated := int64(123) if err := s.SetQuotaConfigPatch(ctx, "tenant-spending-only", QuotaConfigPatch{TiDBCloudSpendingLimit: &updated}); err != nil { @@ -133,7 +133,7 @@ func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) t.Fatal(err) } if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != updated { - t.Fatalf("updated spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, updated) + t.Errorf("updated spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, updated) } checkedAt := time.Now().UTC() if err := s.SetQuotaConfigPatch(ctx, "tenant-spending-only", QuotaConfigPatch{TiDBCloudSpendingLimitCheckedAt: &checkedAt}); err != nil { @@ -144,7 +144,7 @@ func TestQuotaConfigSpendingLimitPatchMaterializesStorageDefaults(t *testing.T) t.Fatal(err) } if cfg.TiDBCloudSpendingLimitCheckedAt == nil { - t.Fatal("spending limit checked_at = nil, want timestamp") + t.Error("spending limit checked_at = nil, want timestamp") } } diff --git a/pkg/server/quota_test.go b/pkg/server/quota_test.go index 961d27d4..2f4255c2 100644 --- a/pkg/server/quota_test.go +++ b/pkg/server/quota_test.go @@ -1100,31 +1100,31 @@ func TestQuotaSetSpendingLimitOnlyPersistsSpendingLimitConfig(t *testing.T) { } calls := rt.prov.callsSnapshot() if len(calls) != 2 || calls[0] != "mark" || calls[1] != "update" { - t.Fatalf("calls = %#v, want mark before update", calls) + t.Errorf("calls = %#v, want mark before update", calls) } lastOptions := rt.prov.lastOptionsSnapshot() if lastOptions.TiDBCloudSpendingLimitMonthly == nil || *lastOptions.TiDBCloudSpendingLimitMonthly != spendingLimit { - t.Fatalf("spending limit option = %#v, want %d", lastOptions.TiDBCloudSpendingLimitMonthly, spendingLimit) + t.Errorf("spending limit option = %#v, want %d", lastOptions.TiDBCloudSpendingLimitMonthly, spendingLimit) } cfg, err := rt.meta.GetQuotaConfig(context.Background(), rt.tenantID) if err != nil { t.Fatal(err) } if cfg.TiDBCloudSpendingLimit == nil || *cfg.TiDBCloudSpendingLimit != spendingLimit { - t.Fatalf("persisted spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, spendingLimit) + t.Errorf("persisted spending limit = %#v, want %d", cfg.TiDBCloudSpendingLimit, spendingLimit) } if cfg.MaxStorageBytes != meta.DefaultMaxStorageBytes() || cfg.MaxFileSizeBytes != meta.DefaultMaxFileSizeBytes() || cfg.MaxFileCount != 0 { - t.Fatalf("storage quota fields = %#v, want defaults", cfg) + t.Errorf("storage quota fields = %#v, want defaults", cfg) } if !quotaConfigRowExists(t, rt) { - t.Fatal("quota config row was not persisted") + t.Error("quota config row was not persisted") } var out quotaResponse if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { t.Fatal(err) } if out.Config.TiDBCloudSpendingLimit == nil || *out.Config.TiDBCloudSpendingLimit != spendingLimit { - t.Fatalf("config = %#v", out.Config) + t.Errorf("config = %#v", out.Config) } } @@ -1148,7 +1148,7 @@ func TestQuotaSetRejectsDrive9KeyWithoutTiDBCloudCredentials(t *testing.T) { t.Fatalf("mark calls = %d, want 0", got) } if quotaConfigRowExists(t, rt) { - t.Fatal("quota config row was written for a rejected request") + t.Error("quota config row was written for a rejected request") } } @@ -1370,7 +1370,7 @@ func TestQuotaSetMapsTiDBCloudCredentialErrorsWithoutWritingConfig(t *testing.T) t.Fatalf("status = %d, want %d", resp.StatusCode, tc.wantStatus) } if quotaConfigRowExists(t, rt) { - t.Fatal("quota config row was written after TiDB Cloud credential failure") + t.Error("quota config row was written after TiDB Cloud credential failure") } if got := rt.prov.markCalls.Load(); got != 1 { t.Fatalf("mark calls = %d, want 1", got) From 60dfea1999ab825aebbabe403d00eb5e3113d699 Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 11:28:46 +0800 Subject: [PATCH 10/11] test: align cache tests with async refresh semantics --- pkg/backend/quota_cache_test.go | 20 ++++++++++++++++++++ pkg/tenant/pool_test.go | 7 ++++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 8124c4fa..5b9238e1 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -24,6 +24,21 @@ func newCacheTestStore() *cacheTestStore { return &cacheTestStore{fakeMetaQuotaStore: newFakeMetaQuotaStore()} } +func waitForQuotaConfigLoad(t *testing.T, c *quotaConfigCache) { + t.Helper() + c.mu.RLock() + done := c.loadDone + c.mu.RUnlock() + if done == nil { + return + } + select { + case <-done: + case <-time.After(300 * time.Millisecond): + t.Fatal("quota config load did not finish") + } +} + func (m *cacheTestStore) GetQuotaUsage(ctx context.Context, tenantID string) (*QuotaUsageView, error) { m.usageCalls.Add(1) if m.usageHook != nil { @@ -315,6 +330,10 @@ func TestQuotaConfigCacheReusesSnapshotUntilTTLExpires(t *testing.T) { c.mu.Lock() c.nextRefresh = time.Time{} c.mu.Unlock() + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("expired config = %+v, want stale storage 1000", cfg) + } + waitForQuotaConfigLoad(t, c) if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 2000 { t.Errorf("refreshed config = %+v, want storage 2000", cfg) } @@ -376,6 +395,7 @@ func TestQuotaConfigCacheRefreshFailureReturnsStaleAndUsesRetryCooldown(t *testi if stale == nil || stale.MaxStorageBytes != 1000 { t.Errorf("stale config = %+v, want storage 1000", stale) } + waitForQuotaConfigLoad(t, c) if got := store.configCalls.Load(); got != 2 { t.Errorf("configCalls after failed refresh = %d, want 2", got) } diff --git a/pkg/tenant/pool_test.go b/pkg/tenant/pool_test.go index a1756aa4..f97c6d61 100644 --- a/pkg/tenant/pool_test.go +++ b/pkg/tenant/pool_test.go @@ -1453,14 +1453,15 @@ func TestIdleEvictionDisabled(t *testing.T) { } } -func TestIdleEvictionStartNoOpWhenDisabled(t *testing.T) { +func TestIdleEvictionStartKeepsSharedHandleReaperWhenDisabled(t *testing.T) { pool := NewPool(PoolConfig{ MaxTenants: 2, IdleTimeout: 0, }, nil) pool.Start(context.Background()) - if pool.reapStop != nil { - t.Fatal("expected reapStop to be nil when IdleTimeout=0") + t.Cleanup(pool.Close) + if pool.reapStop == nil { + t.Fatal("expected shared handle reaper to run when IdleTimeout=0") } } From cac4d5abf90624be66a32752edcc9d5764856a93 Mon Sep 17 00:00:00 2001 From: srstack Date: Mon, 27 Jul 2026 12:23:09 +0800 Subject: [PATCH 11/11] fix: bound async quota refresh concurrency --- pkg/backend/quota_cache.go | 44 ++++++++++++++- pkg/backend/quota_cache_test.go | 99 +++++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+), 3 deletions(-) diff --git a/pkg/backend/quota_cache.go b/pkg/backend/quota_cache.go index 58b725ad..d6d25d26 100644 --- a/pkg/backend/quota_cache.go +++ b/pkg/backend/quota_cache.go @@ -19,6 +19,15 @@ const ( // quotaConfigCacheLoadTimeout bounds a coalesced refresh independently of // the request that happened to claim load ownership. quotaConfigCacheLoadTimeout = 5 * time.Second + // defaultQuotaConfigCacheAsyncRefreshSlots bounds detached warm-cache + // refreshes so a MetaDB slowdown cannot create one in-flight query per + // resident tenant. Keep headroom in the shared metadata connection pool for + // foreground requests and other control-plane work. + defaultQuotaConfigCacheAsyncRefreshSlots = 32 + // quotaConfigCacheSlotRetryInterval prevents a tenant that could not claim + // an async refresh slot from retrying on every request while the budget is + // exhausted. + quotaConfigCacheSlotRetryInterval = time.Second // quotaConfigCacheFailureRetryInterval keeps quota changes responsive after // a transient MetaDB failure. Successful refreshes retain the normal TTL. quotaConfigCacheFailureRetryInterval = 5 * time.Second @@ -36,9 +45,10 @@ const ( // quotaConfigCacheRefreshInterval is the resolved refresh interval (package-level // var so it can be set from env at startup). var ( - quotaConfigCacheRefreshInterval = defaultQuotaConfigCacheRefreshInterval - quotaUsageCacheTTL = defaultQuotaUsageCacheTTL - quotaPendingDeltasCacheTTL = defaultQuotaPendingDeltasCacheTTL + quotaConfigCacheRefreshInterval = defaultQuotaConfigCacheRefreshInterval + quotaUsageCacheTTL = defaultQuotaUsageCacheTTL + quotaPendingDeltasCacheTTL = defaultQuotaPendingDeltasCacheTTL + quotaConfigCacheAsyncRefreshSlots = make(chan struct{}, defaultQuotaConfigCacheAsyncRefreshSlots) ) // InitQuotaConfigCacheRefreshInterval overrides the default refresh interval. @@ -180,7 +190,12 @@ func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { start := now if stale := c.snapshotCopy(); stale != nil { + if !tryAcquireQuotaConfigAsyncRefreshSlot() { + c.deferAsyncRefresh(start) + return stale + } go func() { + defer releaseQuotaConfigAsyncRefreshSlot() defer func() { if recovered := recover(); recovered != nil { logger.Error(context.Background(), "quota_config_cache_async_load_panicked", @@ -194,6 +209,29 @@ func (c *quotaConfigCache) get(ctx context.Context) *QuotaConfigView { return c.loadConfig(ctx, start) } +func tryAcquireQuotaConfigAsyncRefreshSlot() bool { + select { + case quotaConfigCacheAsyncRefreshSlots <- struct{}{}: + return true + default: + return false + } +} + +func releaseQuotaConfigAsyncRefreshSlot() { + <-quotaConfigCacheAsyncRefreshSlots +} + +func (c *quotaConfigCache) deferAsyncRefresh(start time.Time) { + c.mu.Lock() + if c.loadDone != nil { + c.nextRefresh = time.Now().Add(quotaConfigCacheSlotRetryInterval) + c.finishConfigLoadLocked() + } + c.mu.Unlock() + metrics.RecordTenantOperationWithOrg(c.tenantID, c.tidbCloudOrgID, "quota_config_cache", "load", "deferred", time.Since(start)) +} + func (c *quotaConfigCache) snapshotCopy() *QuotaConfigView { c.mu.RLock() defer c.mu.RUnlock() diff --git a/pkg/backend/quota_cache_test.go b/pkg/backend/quota_cache_test.go index 5b9238e1..a2e12598 100644 --- a/pkg/backend/quota_cache_test.go +++ b/pkg/backend/quota_cache_test.go @@ -262,6 +262,105 @@ func TestQuotaConfigCacheWarmWaiterReturnsStaleWithoutWaiting(t *testing.T) { } } +func TestQuotaConfigCacheAsyncRefreshHasGlobalSlotBudget(t *testing.T) { + const tenants = defaultQuotaConfigCacheAsyncRefreshSlots + 8 + stores := make([]*cacheTestStore, tenants) + caches := make([]*quotaConfigCache, tenants) + refreshStarted := make(chan struct{}, tenants) + release := make(chan struct{}) + for i := range tenants { + store := newCacheTestStore() + store.config["t1"] = &QuotaConfigView{MaxStorageBytes: 1000} + store.configHook = func() { + refreshStarted <- struct{}{} + <-release + } + stores[i] = store + cache := newQuotaConfigCache("t1", "", store) + // The first load is cold and must not use the async refresh budget. + store.configHook = nil + if cfg := cache.get(context.Background()); cfg == nil { + t.Fatalf("initial config for tenant %d is nil", i) + } + store.configHook = func() { + refreshStarted <- struct{}{} + <-release + } + cache.mu.Lock() + cache.nextRefresh = time.Time{} + cache.mu.Unlock() + caches[i] = cache + } + + var callers sync.WaitGroup + for _, cache := range caches { + callers.Add(1) + go func(c *quotaConfigCache) { + defer callers.Done() + if cfg := c.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("expired config = %+v, want stale storage 1000", cfg) + } + }(cache) + } + callers.Wait() + + for range defaultQuotaConfigCacheAsyncRefreshSlots { + select { + case <-refreshStarted: + case <-time.After(300 * time.Millisecond): + t.Fatal("async refresh did not consume its available slots") + } + } + select { + case <-refreshStarted: + t.Fatal("async refresh exceeded the global slot budget") + default: + } + + deferred := 0 + for i, cache := range caches { + cache.mu.RLock() + loadInFlight := cache.loadDone != nil + cache.mu.RUnlock() + if stores[i].configCalls.Load() == 1 { + deferred++ + if loadInFlight { + t.Errorf("deferred tenant %d retained load ownership", i) + } + } + } + if deferred == 0 { + t.Error("expected at least one tenant refresh to be deferred") + } + + close(release) + for _, cache := range caches { + waitForQuotaConfigLoad(t, cache) + } + for i, store := range stores { + if store.configCalls.Load() > 2 { + t.Errorf("tenant %d configCalls = %d, want at most 2", i, store.configCalls.Load()) + } + } + + for i, cache := range caches { + if stores[i].configCalls.Load() != 1 { + continue + } + cache.mu.Lock() + cache.nextRefresh = time.Time{} + cache.mu.Unlock() + if cfg := cache.get(context.Background()); cfg == nil || cfg.MaxStorageBytes != 1000 { + t.Errorf("deferred tenant %d config = %+v, want stale storage 1000", i, cfg) + } + waitForQuotaConfigLoad(t, cache) + if got := stores[i].configCalls.Load(); got != 2 { + t.Errorf("deferred tenant %d configCalls = %d, want 2 after retry", i, got) + } + break + } +} + func TestQuotaConfigCacheIsPassiveUntilFirstAccess(t *testing.T) { previousRefreshInterval := quotaConfigCacheRefreshInterval quotaConfigCacheRefreshInterval = 5 * time.Millisecond