perf: scale shared tenant backend cache

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 == "" {

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,82 @@ 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.Errorf("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.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 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",
@@ -682,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)
 		}
 	}

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
 }

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
 	}
 }

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