diff --git a/docker-compose.yml b/docker-compose.yml index 870963f8a..18898d4a4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -187,6 +187,16 @@ services: SPINDLE_NIX_CACHE_READ_URLS: http://ncps:8501 SPINDLE_NIX_CACHE_TRUSTED_PUBLIC_KEYS: cache.local:F7YqpMzuBdILYd/v+wMZN2YKxCzliXQyFmeezOxw7rU= SPINDLE_NIX_CACHE_UPLOAD_URL: http://ncps:8501/upload + SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB: 0 + SPINDLE_QUOTA_USER_WORKFLOWS: 0 + SPINDLE_QUOTA_USER_VCPUS: 0 + SPINDLE_QUOTA_USER_MEMORY_MIB: 0 + SPINDLE_QUOTA_USER_DISK_MIB: 0 + SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB: 0 + SPINDLE_QUOTA_REPO_WORKFLOWS: 0 + SPINDLE_QUOTA_REPO_VCPUS: 0 + SPINDLE_QUOTA_REPO_MEMORY_MIB: 0 + SPINDLE_QUOTA_REPO_DISK_MIB: 0 SPINDLE_MICROVM_PIPELINES_DEBUG_SSH_ENABLED: true SPINDLE_MICROVM_PIPELINES_DEBUG_SSH_LISTEN_ADDR: 0.0.0.0:2223 SPINDLE_MICROVM_PIPELINES_DEBUG_SSH_GRACE_PERIOD: 10m diff --git a/localinfra/observability/grafana/provisioning/dashboards/spindle.json b/localinfra/observability/grafana/provisioning/dashboards/spindle.json index 2d773f043..af6e3f1d0 100644 --- a/localinfra/observability/grafana/provisioning/dashboards/spindle.json +++ b/localinfra/observability/grafana/provisioning/dashboards/spindle.json @@ -1309,7 +1309,7 @@ }, "id": 1001, "title": "Cache Storage Quota Usage & Subjects", - "description": "Committed cache storage quota usage in bytes and count of quota subjects by scope and status.", + "description": "Committed cache storage quota usage, configured defaults (zero means unlimited), and count of quota subjects by scope and status.", "type": "timeseries", "targets": [ { @@ -1333,6 +1333,17 @@ "legendFormat": "Subjects: {{scope}} ({{status}})", "range": true, "refId": "B" + }, + { + "datasource": { + "type": "prometheus", + "uid": "$datasource" + }, + "editorMode": "code", + "expr": "spindle_quota_default_limit{job=\"spindle\", instance=~\"$instance\", resource=\"cache_storage_bytes\"}", + "legendFormat": "Default limit: {{scope}}", + "range": true, + "refId": "C" } ], "options": { @@ -1376,7 +1387,7 @@ }, "id": 1002, "title": "Workflow, vCPU, Memory & Disk Quota Usage", - "description": "Committed workflow, vCPU, memory (MiB) and disk (MiB) quota usage and subject counts by scope and resource.", + "description": "Committed workflow, vCPU, memory (MiB) and disk (MiB) quota usage, configured defaults (zero means unlimited), and subject counts by scope and resource.", "type": "timeseries", "targets": [ { @@ -1400,6 +1411,17 @@ "legendFormat": "Subjects: {{scope}} {{resource}} ({{status}})", "range": true, "refId": "B" + }, + { + "datasource": { + "type": "prometheus", + "uid": "$datasource" + }, + "editorMode": "code", + "expr": "spindle_quota_default_limit{job=\"spindle\", instance=~\"$instance\", resource=~\"workflows|vcpus|memory_mib|disk_mib\"}", + "legendFormat": "Default limit: {{scope}} {{resource}}", + "range": true, + "refId": "C" } ], "options": { @@ -2522,4 +2544,4 @@ "version": 1, "weekStart": "", "refresh": "10s" -} +} \ No newline at end of file diff --git a/nix/modules/spindle.nix b/nix/modules/spindle.nix index 7011c155f..ad878257a 100644 --- a/nix/modules/spindle.nix +++ b/nix/modules/spindle.nix @@ -403,6 +403,63 @@ in }; }; + quota = { + user = { + cacheStorageMiB = mkOption { + type = types.int; + default = 0; + description = "Default cache storage limit for users in MiB. Zero is unlimited."; + }; + workflows = mkOption { + type = types.int; + default = 0; + description = "Default concurrent workflows limit for users. Zero is unlimited."; + }; + vCPUs = mkOption { + type = types.int; + default = 0; + description = "Default vCPU limit for users. Zero is unlimited."; + }; + memoryMiB = mkOption { + type = types.int; + default = 0; + description = "Default memory limit for users in MiB. Zero is unlimited."; + }; + diskMiB = mkOption { + type = types.int; + default = 0; + description = "Default disk limit for users in MiB. Zero is unlimited."; + }; + }; + repo = { + cacheStorageMiB = mkOption { + type = types.int; + default = 0; + description = "Default cache storage limit for repositories in MiB. Zero is unlimited."; + }; + workflows = mkOption { + type = types.int; + default = 0; + description = "Default concurrent workflows limit for repositories. Zero is unlimited."; + }; + vCPUs = mkOption { + type = types.int; + default = 0; + description = "Default vCPU limit for repositories. Zero is unlimited."; + }; + memoryMiB = mkOption { + type = types.int; + default = 0; + description = "Default memory limit for repositories in MiB. Zero is unlimited."; + }; + diskMiB = mkOption { + type = types.int; + default = 0; + description = "Default disk limit for repositories in MiB. Zero is unlimited."; + }; + }; + }; + environmentFile = mkOption { type = with types; nullOr path; default = null; @@ -673,6 +730,16 @@ in "SPINDLE_NIX_CACHE_READ_URLS=${concatStringsSep "," instance.pipelines.nixCache.readUrls}" "SPINDLE_NIX_CACHE_TRUSTED_PUBLIC_KEYS=${concatStringsSep "," instance.pipelines.nixCache.trustedPublicKeys}" "SPINDLE_NIX_CACHE_UPLOAD_URL=${instance.pipelines.nixCache.uploadUrl}" + "SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB=${toString instance.quota.user.cacheStorageMiB}" + "SPINDLE_QUOTA_USER_WORKFLOWS=${toString instance.quota.user.workflows}" + "SPINDLE_QUOTA_USER_VCPUS=${toString instance.quota.user.vCPUs}" + "SPINDLE_QUOTA_USER_MEMORY_MIB=${toString instance.quota.user.memoryMiB}" + "SPINDLE_QUOTA_USER_DISK_MIB=${toString instance.quota.user.diskMiB}" + "SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB=${toString instance.quota.repo.cacheStorageMiB}" + "SPINDLE_QUOTA_REPO_WORKFLOWS=${toString instance.quota.repo.workflows}" + "SPINDLE_QUOTA_REPO_VCPUS=${toString instance.quota.repo.vCPUs}" + "SPINDLE_QUOTA_REPO_MEMORY_MIB=${toString instance.quota.repo.memoryMiB}" + "SPINDLE_QUOTA_REPO_DISK_MIB=${toString instance.quota.repo.diskMiB}" "SPINDLE_ARTIFACT_STORES_DISK_DIR=${instance.artifactStores.disk.dir}" "SPINDLE_ARTIFACT_STORES_S3_BUCKET=${instance.artifactStores.s3.bucket}" "SPINDLE_ARTIFACT_STORES_S3_REGION=${instance.artifactStores.s3.region}" diff --git a/spindle/config/config.go b/spindle/config/config.go index 0b329656a..d79a86a01 100644 --- a/spindle/config/config.go +++ b/spindle/config/config.go @@ -8,6 +8,7 @@ import ( "github.com/bluesky-social/indigo/atproto/syntax" "github.com/sethvargo/go-envconfig" + "tangled.org/core/spindle/quota" "tangled.org/core/xrpc/serviceauth" ) @@ -129,6 +130,27 @@ type NixCache struct { UploadURL string `env:"UPLOAD_URL"` } +type Quota struct { + User struct { + CacheStorageMiB int64 `env:"CACHE_STORAGE_MIB, default=0"` + Workflows int64 `env:"WORKFLOWS, default=0"` + VCPUs int64 `env:"VCPUS, default=0"` + MemoryMiB int64 `env:"MEMORY_MIB, default=0"` + DiskMiB int64 `env:"DISK_MIB, default=0"` + } `env:",prefix=USER_"` + Repo struct { + CacheStorageMiB int64 `env:"CACHE_STORAGE_MIB, default=0"` + Workflows int64 `env:"WORKFLOWS, default=0"` + VCPUs int64 `env:"VCPUS, default=0"` + MemoryMiB int64 `env:"MEMORY_MIB, default=0"` + DiskMiB int64 `env:"DISK_MIB, default=0"` + } `env:",prefix=REPO_"` +} + +type quotaEnvironment struct { + Quota Quota `env:",prefix=SPINDLE_QUOTA_"` +} + // governs how spindle places and runs jobs type Role string @@ -172,6 +194,7 @@ type Config struct { NixeryPipelines NixeryPipelines `env:",prefix=SPINDLE_NIXERY_PIPELINES_"` MicroVMPipelines MicroVMPipelines `env:",prefix=SPINDLE_MICROVM_PIPELINES_"` NixCache NixCache `env:",prefix=SPINDLE_NIX_CACHE_"` + Quota Quota `env:",prefix=SPINDLE_QUOTA_"` ArtifactStores ArtifactStores `env:",prefix=SPINDLE_ARTIFACT_STORES_"` LegacyS3 LegacyS3 `env:",prefix=SPINDLE_S3_"` Mill Mill `env:",prefix=SPINDLE_MILL_"` @@ -221,9 +244,117 @@ func (c *Config) validate() error { if c.Logging.Endpoint != "" && c.Logging.ServiceName == "" { return fmt.Errorf("SPINDLE_LOGGING_SERVICE_NAME must not be empty when logging is enabled") } + if c.Quota.User.CacheStorageMiB < 0 { + return fmt.Errorf("user cache storage limit cannot be negative, got %d", c.Quota.User.CacheStorageMiB) + } + if c.Quota.User.Workflows < 0 { + return fmt.Errorf("user workflows limit cannot be negative, got %d", c.Quota.User.Workflows) + } + if c.Quota.User.VCPUs < 0 { + return fmt.Errorf("user vcpus limit cannot be negative, got %d", c.Quota.User.VCPUs) + } + if c.Quota.User.MemoryMiB < 0 { + return fmt.Errorf("user memory limit cannot be negative, got %d", c.Quota.User.MemoryMiB) + } + if c.Quota.User.DiskMiB < 0 { + return fmt.Errorf("user disk limit cannot be negative, got %d", c.Quota.User.DiskMiB) + } + + if c.Quota.Repo.CacheStorageMiB < 0 { + return fmt.Errorf("repo cache storage limit cannot be negative, got %d", c.Quota.Repo.CacheStorageMiB) + } + if c.Quota.Repo.Workflows < 0 { + return fmt.Errorf("repo workflows limit cannot be negative, got %d", c.Quota.Repo.Workflows) + } + if c.Quota.Repo.VCPUs < 0 { + return fmt.Errorf("repo vcpus limit cannot be negative, got %d", c.Quota.Repo.VCPUs) + } + if c.Quota.Repo.MemoryMiB < 0 { + return fmt.Errorf("repo memory limit cannot be negative, got %d", c.Quota.Repo.MemoryMiB) + } + if c.Quota.Repo.DiskMiB < 0 { + return fmt.Errorf("repo disk limit cannot be negative, got %d", c.Quota.Repo.DiskMiB) + } + + if c.Quota.User.CacheStorageMiB > 8796093022207 { + return fmt.Errorf("user cache storage limit %d MiB would overflow bytes", c.Quota.User.CacheStorageMiB) + } + if c.Quota.Repo.CacheStorageMiB > 8796093022207 { + return fmt.Errorf("repo cache storage limit %d MiB would overflow bytes", c.Quota.Repo.CacheStorageMiB) + } return nil } +func (c *Config) ToDefaults() (quota.Defaults, error) { + mibToBytes := func(mib int64) (int64, error) { + if mib < 0 { + return 0, fmt.Errorf("negative limit: %d", mib) + } + if mib == 0 { + return 0, nil + } + if mib > 8796093022207 { + return 0, fmt.Errorf("limit %d MiB would overflow bytes", mib) + } + return mib * 1024 * 1024, nil + } + + limits := []struct { + scope string + resource string + limit int64 + }{ + {scope: "user", resource: quota.ResourceWorkflows, limit: c.Quota.User.Workflows}, + {scope: "user", resource: quota.ResourceVCPUs, limit: c.Quota.User.VCPUs}, + {scope: "user", resource: quota.ResourceMemoryMiB, limit: c.Quota.User.MemoryMiB}, + {scope: "user", resource: quota.ResourceDiskMiB, limit: c.Quota.User.DiskMiB}, + {scope: "repo", resource: quota.ResourceWorkflows, limit: c.Quota.Repo.Workflows}, + {scope: "repo", resource: quota.ResourceVCPUs, limit: c.Quota.Repo.VCPUs}, + {scope: "repo", resource: quota.ResourceMemoryMiB, limit: c.Quota.Repo.MemoryMiB}, + {scope: "repo", resource: quota.ResourceDiskMiB, limit: c.Quota.Repo.DiskMiB}, + } + for _, item := range limits { + if item.limit < 0 { + return nil, fmt.Errorf("%s %s limit cannot be negative, got %d", item.scope, item.resource, item.limit) + } + } + + userCacheBytes, err := mibToBytes(c.Quota.User.CacheStorageMiB) + if err != nil { + return nil, err + } + repoCacheBytes, err := mibToBytes(c.Quota.Repo.CacheStorageMiB) + if err != nil { + return nil, err + } + + return quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: userCacheBytes, + quota.ResourceWorkflows: c.Quota.User.Workflows, + quota.ResourceVCPUs: c.Quota.User.VCPUs, + quota.ResourceMemoryMiB: c.Quota.User.MemoryMiB, + quota.ResourceDiskMiB: c.Quota.User.DiskMiB, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: repoCacheBytes, + quota.ResourceWorkflows: c.Quota.Repo.Workflows, + quota.ResourceVCPUs: c.Quota.Repo.VCPUs, + quota.ResourceMemoryMiB: c.Quota.Repo.MemoryMiB, + quota.ResourceDiskMiB: c.Quota.Repo.DiskMiB, + }, + }, nil +} + +func LoadQuotaDefaults(ctx context.Context) (quota.Defaults, error) { + var env quotaEnvironment + if err := envconfig.Process(ctx, &env); err != nil { + return nil, err + } + cfg := Config{Quota: env.Quota} + return cfg.ToDefaults() +} + func Load(ctx context.Context) (*Config, error) { var cfg Config err := envconfig.Process(ctx, &cfg) diff --git a/spindle/config/config_test.go b/spindle/config/config_test.go index 976126d42..e269e61e9 100644 --- a/spindle/config/config_test.go +++ b/spindle/config/config_test.go @@ -4,6 +4,8 @@ import ( "context" "testing" "time" + + "tangled.org/core/spindle/quota" ) func TestLoadAllowsUnconfiguredMicroVMEngine(t *testing.T) { @@ -180,6 +182,154 @@ func TestLoadTracingConfigValidation(t *testing.T) { } } +func TestLoadQuotaLimitsValidation(t *testing.T) { + tests := []struct { + name string + envs map[string]string + wantErr bool + }{ + { + name: "valid defaults", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB": "0", + "SPINDLE_QUOTA_USER_WORKFLOWS": "0", + "SPINDLE_QUOTA_USER_VCPUS": "0", + "SPINDLE_QUOTA_USER_MEMORY_MIB": "0", + "SPINDLE_QUOTA_USER_DISK_MIB": "0", + "SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB": "0", + "SPINDLE_QUOTA_REPO_WORKFLOWS": "0", + "SPINDLE_QUOTA_REPO_VCPUS": "0", + "SPINDLE_QUOTA_REPO_MEMORY_MIB": "0", + "SPINDLE_QUOTA_REPO_DISK_MIB": "0", + }, + wantErr: false, + }, + { + name: "valid positive values", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB": "100", + "SPINDLE_QUOTA_USER_WORKFLOWS": "5", + "SPINDLE_QUOTA_USER_VCPUS": "4", + "SPINDLE_QUOTA_USER_MEMORY_MIB": "4096", + "SPINDLE_QUOTA_USER_DISK_MIB": "10240", + }, + wantErr: false, + }, + { + name: "negative user cache storage", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB": "-1", + }, + wantErr: true, + }, + { + name: "negative user workflows", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_WORKFLOWS": "-1", + }, + wantErr: true, + }, + { + name: "negative user vcpus", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_VCPUS": "-1", + }, + wantErr: true, + }, + { + name: "negative user memory", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_MEMORY_MIB": "-1", + }, + wantErr: true, + }, + { + name: "negative user disk", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_DISK_MIB": "-1", + }, + wantErr: true, + }, + { + name: "negative repo cache storage", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB": "-1", + }, + wantErr: true, + }, + { + name: "negative repo workflows", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_WORKFLOWS": "-1", + }, + wantErr: true, + }, + { + name: "negative repo vcpus", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_VCPUS": "-1", + }, + wantErr: true, + }, + { + name: "negative repo memory", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_MEMORY_MIB": "-1", + }, + wantErr: true, + }, + { + name: "negative repo disk", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_DISK_MIB": "-1", + }, + wantErr: true, + }, + { + name: "user cache storage overflow", + envs: map[string]string{ + "SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB": "9223372036854775807", + }, + wantErr: true, + }, + { + name: "repo cache storage overflow", + envs: map[string]string{ + "SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB": "9223372036854775807", + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("SPINDLE_QUOTA_USER_CACHE_STORAGE_MIB", "") + t.Setenv("SPINDLE_QUOTA_USER_WORKFLOWS", "") + t.Setenv("SPINDLE_QUOTA_USER_VCPUS", "") + t.Setenv("SPINDLE_QUOTA_USER_MEMORY_MIB", "") + t.Setenv("SPINDLE_QUOTA_USER_DISK_MIB", "") + t.Setenv("SPINDLE_QUOTA_REPO_CACHE_STORAGE_MIB", "") + t.Setenv("SPINDLE_QUOTA_REPO_WORKFLOWS", "") + t.Setenv("SPINDLE_QUOTA_REPO_VCPUS", "") + t.Setenv("SPINDLE_QUOTA_REPO_MEMORY_MIB", "") + t.Setenv("SPINDLE_QUOTA_REPO_DISK_MIB", "") + + t.Setenv("SPINDLE_SERVER_HOSTNAME", "spindle.example.com") + t.Setenv("SPINDLE_SERVER_OWNER", "did:web:spindle.example.com") + t.Setenv("SPINDLE_TRACING_SAMPLE_RATIO", "1.0") + + for k, v := range tt.envs { + t.Setenv(k, v) + } + + _, err := Load(context.Background()) + if (err != nil) != tt.wantErr { + t.Errorf("Load() error = %v, wantErr = %v", err, tt.wantErr) + } + }) + } +} + func TestLoadLoggingConfig(t *testing.T) { t.Setenv("SPINDLE_SERVER_HOSTNAME", "spindle.example.com") t.Setenv("SPINDLE_SERVER_OWNER", "did:web:spindle.example.com") @@ -213,3 +363,21 @@ func TestLoadLoggingConfigValidation(t *testing.T) { t.Fatal("expected error for empty ServiceName when logging endpoint is configured") } } + +func TestLoadQuotaDefaultsDoesNotRequireServerConfiguration(t *testing.T) { + t.Setenv("SPINDLE_QUOTA_USER_WORKFLOWS", "3") + defaults, err := LoadQuotaDefaults(context.Background()) + if err != nil { + t.Fatal(err) + } + if got := defaults[quota.ScopeUser][quota.ResourceWorkflows]; got != 3 { + t.Fatalf("user workflow limit = %d, want 3", got) + } +} + +func TestLoadQuotaDefaultsRejectsNegativeLimits(t *testing.T) { + t.Setenv("SPINDLE_QUOTA_REPO_VCPUS", "-1") + if _, err := LoadQuotaDefaults(context.Background()); err == nil { + t.Fatal("negative quota limit accepted") + } +} diff --git a/spindle/db/db.go b/spindle/db/db.go index 276f835ed..34ca9d5b0 100644 --- a/spindle/db/db.go +++ b/spindle/db/db.go @@ -388,6 +388,178 @@ func runMigrations(_ context.Context, conn *sql.Conn, logger *slog.Logger) error return err } + if err := orm.RunMigration(conn, logger, "cache-quotas-schema", func(tx *sql.Tx) error { + _, err := tx.Exec(` + create table if not exists cache_objects ( + repo_did text not null, + object_key text not null, + byte_size integer not null check (byte_size >= 0), + primary key (repo_did, object_key) + ); + create index if not exists idx_cache_objects_key on cache_objects(object_key); + + create table if not exists cache_repo_owners ( + repo_did text primary key, + owner_did text not null + ); + create index if not exists idx_cache_repo_owners_owner on cache_repo_owners(owner_did); + + create table if not exists cache_reservations ( + id text primary key, + object_key text not null, + byte_size integer not null check (byte_size >= 0), + repo_did text not null, + owner_did text not null, + phase text not null check (phase in ('reserved', 'publishing')), + created_at integer not null + ); + create index if not exists idx_cache_reservations_repo on cache_reservations(repo_did); + create index if not exists idx_cache_reservations_owner on cache_reservations(owner_did); + + create table if not exists cache_overrides ( + scope text not null check (scope in ('user', 'repo')), + did text not null, + max_bytes integer check (max_bytes >= 0 or max_bytes is null), + primary key (scope, did) + ); + create index if not exists idx_cache_overrides_scope on cache_overrides(scope); + `) + return err + }); err != nil { + return err + } + + if err := orm.RunMigration(conn, logger, "quota-schema", func(tx *sql.Tx) error { + _, err := tx.Exec(` + create table if not exists quota_limits ( + did text not null, + resource text not null, + max_amount integer check (max_amount >= 0 or max_amount is null), + primary key (did, resource) + ); + + create table if not exists quota_reservations ( + id text not null, + resource text not null, + kind text not null, + key text not null, + amount integer not null check (amount >= 0), + repo_did text not null, + owner_did text not null, + phase text not null check (phase in ('reserved', 'publishing', 'active')), + created_at integer not null, + primary key (id, resource) + ); + create index if not exists idx_quota_reservations_repo on quota_reservations(repo_did); + create index if not exists idx_quota_reservations_owner on quota_reservations(owner_did); + + create table if not exists quota_allocations ( + repo_did text not null, + resource text not null, + kind text not null, + key text not null, + amount integer not null check (amount >= 0), + primary key (repo_did, resource, kind, key) + ); + create index if not exists idx_quota_allocations_key on quota_allocations(key); + + create table if not exists quota_repo_owners ( + repo_did text primary key, + owner_did text not null + ); + create index if not exists idx_quota_repo_owners_owner on quota_repo_owners(owner_did); + `) + if err != nil { + return err + } + + var hasOverrides int + err = tx.QueryRow(`select count(*) from sqlite_master where type='table' and name='cache_overrides'`).Scan(&hasOverrides) + if err != nil { + return err + } + if hasOverrides > 0 { + _, err = tx.Exec(` + insert or ignore into quota_limits (did, resource, max_amount) + select did, 'cache_storage_bytes', max_bytes from cache_overrides + `) + if err != nil { + return err + } + } + + var hasReservations int + err = tx.QueryRow(`select count(*) from sqlite_master where type='table' and name='cache_reservations'`).Scan(&hasReservations) + if err != nil { + return err + } + if hasReservations > 0 { + _, err = tx.Exec(` + insert or ignore into quota_reservations (id, resource, kind, key, amount, repo_did, owner_did, phase, created_at) + select id, 'cache_storage_bytes', 'nix_cache', object_key, byte_size, repo_did, owner_did, phase, created_at from cache_reservations + `) + if err != nil { + return err + } + } + + var hasObjects int + err = tx.QueryRow(`select count(*) from sqlite_master where type='table' and name='cache_objects'`).Scan(&hasObjects) + if err != nil { + return err + } + if hasObjects > 0 { + _, err = tx.Exec(` + insert or ignore into quota_allocations (repo_did, resource, kind, key, amount) + select repo_did, 'cache_storage_bytes', 'nix_cache', object_key, byte_size from cache_objects + `) + if err != nil { + return err + } + } + + var hasRepoOwners int + err = tx.QueryRow(`select count(*) from sqlite_master where type='table' and name='cache_repo_owners'`).Scan(&hasRepoOwners) + if err != nil { + return err + } + if hasRepoOwners > 0 { + _, err = tx.Exec(` + insert or ignore into quota_repo_owners (repo_did, owner_did) + select repo_did, owner_did from cache_repo_owners + `) + if err != nil { + return err + } + } + + _, err = tx.Exec(` + drop index if exists idx_cache_objects_key; + drop index if exists idx_cache_repo_owners_owner; + drop index if exists idx_cache_reservations_repo; + drop index if exists idx_cache_reservations_owner; + drop index if exists idx_cache_overrides_scope; + drop table if exists cache_objects; + drop table if exists cache_repo_owners; + drop table if exists cache_reservations; + drop table if exists cache_overrides; + `) + if err != nil { + return err + } + + return nil + }); err != nil { + return err + } + + if err := orm.RunMigration(conn, logger, "quota-zero-means-unlimited", func(tx *sql.Tx) error { + _, err := tx.Exec(`update quota_limits set max_amount = null where max_amount = 0`) + return err + }); err != nil { + return err + } + return nil } diff --git a/spindle/db/quota.go b/spindle/db/quota.go new file mode 100644 index 000000000..f0b80b32c --- /dev/null +++ b/spindle/db/quota.go @@ -0,0 +1,843 @@ +package db + +import ( + "context" + "crypto/rand" + "database/sql" + "encoding/hex" + "errors" + "fmt" + "math" + "sort" + "strings" + "time" + + "tangled.org/core/spindle/quota" +) + +type QuotaStore struct { + db *DB + defaults quota.Defaults +} + +var _ quota.Store = (*QuotaStore)(nil) + +func NewQuotaStore(db *DB, defaults quota.Defaults) *QuotaStore { + return &QuotaStore{ + db: db, + defaults: defaults, + } +} + +func addQuotaUsage(total, amount int64) (int64, error) { + if amount < 0 || amount > math.MaxInt64-total { + return 0, fmt.Errorf("quota usage overflow: %d + %d", total, amount) + } + return total + amount, nil +} + +func sumQuotaUsage(ctx context.Context, conn *sql.Conn, kind, key string, amount int64, query string, args ...any) (total int64, exists bool, err error) { + rows, err := conn.QueryContext(ctx, query, args...) + if err != nil { + return 0, false, err + } + defer rows.Close() + + for rows.Next() { + var rowKind, rowKey string + var rowAmount int64 + if err := rows.Scan(&rowKind, &rowKey, &rowAmount); err != nil { + return 0, false, err + } + total, err = addQuotaUsage(total, rowAmount) + if err != nil { + return 0, false, err + } + exists = exists || rowKind == kind && rowKey == key && rowAmount == amount + } + return total, exists, rows.Err() +} + +func (d *QuotaStore) Close() error { + return d.db.Close() +} + +func (d *QuotaStore) withTx(ctx context.Context, fn func(conn *sql.Conn) error) (retErr error) { + conn, err := d.db.Conn(ctx) + if err != nil { + return err + } + defer conn.Close() + + _, err = conn.ExecContext(ctx, "BEGIN IMMEDIATE") + if err != nil { + return err + } + var committed bool + defer func() { + if !committed { + rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if _, rollbackErr := conn.ExecContext(rollbackCtx, "ROLLBACK"); rollbackErr != nil { + retErr = errors.Join(retErr, fmt.Errorf("rollback quota transaction: %w", rollbackErr)) + } + } + }() + + if err := fn(conn); err != nil { + return err + } + + _, err = conn.ExecContext(ctx, "COMMIT") + if err != nil { + return err + } + committed = true + return nil +} + +func (d *QuotaStore) withReadTx(ctx context.Context, fn func(tx *sql.Tx) error) error { + tx, err := d.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return err + } + defer tx.Rollback() + + if err := fn(tx); err != nil { + return err + } + + return tx.Commit() +} + +func (d *QuotaStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + if err := quota.Validate(req); err != nil { + return quota.Reservation{}, err + } + + var res quota.Reservation + err := d.withTx(ctx, func(conn *sql.Conn) error { + _, err := conn.ExecContext(ctx, ` + INSERT INTO quota_repo_owners (repo_did, owner_did) + VALUES (?, ?) + ON CONFLICT(repo_did) DO UPDATE SET owner_did = excluded.owner_did + `, req.Identity.RepoDID, req.Identity.OwnerDID) + if err != nil { + return err + } + + rows, err := conn.QueryContext(ctx, ` + SELECT id, resource, amount FROM quota_reservations + WHERE repo_did = ? AND kind = ? AND key = ? AND phase IN ('reserved', 'publishing', 'active') + `, req.Identity.RepoDID, string(req.Kind), req.Key) + if err != nil { + return err + } + defer rows.Close() + + resMapByID := make(map[string]map[string]int64) + for rows.Next() { + var id, resourceStr string + var amount int64 + if err := rows.Scan(&id, &resourceStr, &amount); err != nil { + return err + } + if _, ok := resMapByID[id]; !ok { + resMapByID[id] = make(map[string]int64) + } + resMapByID[id][resourceStr] = amount + } + for id, resources := range resMapByID { + if len(resources) == len(req.Resources) { + match := true + for r, amt := range req.Resources { + if resources[r] != amt { + match = false + break + } + } + if match { + res = quota.Reservation{ + ID: id, + Allowed: true, + Temporary: false, + Reason: quota.ReasonUnlimited, + } + return nil + } + } + } + + allocRows, err := conn.QueryContext(ctx, ` + SELECT resource, amount FROM quota_allocations + WHERE repo_did = ? AND kind = ? AND key = ? + `, req.Identity.RepoDID, string(req.Kind), req.Key) + if err != nil { + return err + } + defer allocRows.Close() + + allocResources := make(map[string]int64) + for allocRows.Next() { + var resourceStr string + var amount int64 + if err := allocRows.Scan(&resourceStr, &amount); err != nil { + return err + } + allocResources[resourceStr] = amount + } + if len(allocResources) == len(req.Resources) { + match := true + for r, amt := range req.Resources { + if allocResources[r] != amt { + match = false + break + } + } + if match { + res = quota.Reservation{ + ID: "", + Allowed: true, + Temporary: false, + Reason: quota.ReasonUnlimited, + } + return nil + } + } + + deny := func(reason, resource string, temporary bool) { + res = quota.Reservation{ + Allowed: false, + Temporary: temporary, + Reason: reason, + Resource: resource, + } + } + + var sortedResources []string + for r := range req.Resources { + sortedResources = append(sortedResources, r) + } + sort.Strings(sortedResources) + + for _, resource := range sortedResources { + amount := req.Resources[resource] + + var repoLimit *int64 + var maxAmount *int64 + err = conn.QueryRowContext(ctx, ` + SELECT max_amount FROM quota_limits + WHERE did = ? AND resource = ? + `, req.Identity.RepoDID, resource).Scan(&maxAmount) + if err == nil { + if maxAmount != nil { + val := *maxAmount + repoLimit = &val + } + } else if errors.Is(err, sql.ErrNoRows) { + if d.defaults != nil { + if resMap, ok := d.defaults[quota.ScopeRepo]; ok { + if val, ok := resMap[resource]; ok && val > 0 { + repoLimit = &val + } + } + } + } else { + return err + } + + currentRepoUsage, _, err := sumQuotaUsage(ctx, conn, string(req.Kind), req.Key, amount, ` + SELECT kind, key, amount FROM ( + SELECT kind, key, amount FROM quota_allocations WHERE repo_did = ? AND resource = ? + UNION + SELECT kind, key, amount FROM quota_reservations WHERE repo_did = ? AND resource = ? AND phase IN ('reserved', 'publishing', 'active') + ) + `, req.Identity.RepoDID, resource, req.Identity.RepoDID, resource) + if err != nil { + return err + } + if amount > math.MaxInt64-currentRepoUsage { + deny(quota.ReasonRepoLimit, resource, true) + return nil + } + if repoLimit != nil && *repoLimit >= 0 && amount > *repoLimit-currentRepoUsage { + deny(quota.ReasonRepoLimit, resource, amount <= *repoLimit) + return nil + } + + var userLimit *int64 + var maxUserAmount *int64 + err = conn.QueryRowContext(ctx, ` + SELECT max_amount FROM quota_limits + WHERE did = ? AND resource = ? + `, req.Identity.OwnerDID, resource).Scan(&maxUserAmount) + if err == nil { + if maxUserAmount != nil { + val := *maxUserAmount + userLimit = &val + } + } else if errors.Is(err, sql.ErrNoRows) { + if d.defaults != nil { + if resMap, ok := d.defaults[quota.ScopeUser]; ok { + if val, ok := resMap[resource]; ok && val > 0 { + userLimit = &val + } + } + } + } else { + return err + } + + currentUserUsage, existsOwner, err := sumQuotaUsage(ctx, conn, string(req.Kind), req.Key, amount, ` + SELECT kind, key, amount FROM ( + SELECT a.kind, a.key, a.amount + FROM quota_allocations a + JOIN quota_repo_owners r ON a.repo_did = r.repo_did + WHERE r.owner_did = ? AND a.resource = ? + UNION + SELECT res.kind, res.key, res.amount + FROM quota_reservations res + WHERE res.owner_did = ? AND res.resource = ? AND res.phase IN ('reserved', 'publishing', 'active') + ) + `, req.Identity.OwnerDID, resource, req.Identity.OwnerDID, resource) + if err != nil { + return err + } + + var additional int64 + if !existsOwner { + additional = amount + } + if additional > math.MaxInt64-currentUserUsage { + deny(quota.ReasonUserLimit, resource, true) + return nil + } + + // existing reservations survive lower limits + if userLimit != nil && *userLimit >= 0 && additional > 0 && additional > *userLimit-currentUserUsage { + deny(quota.ReasonUserLimit, resource, additional <= *userLimit) + return nil + } + } + + resID := req.ID + if resID == "" { + var b [16]byte + if _, err = rand.Read(b[:]); err != nil { + return err + } + resID = hex.EncodeToString(b[:]) + } + + phase := "reserved" + if req.Kind == quota.KindWorkflow { + phase = "active" + } + + for _, resource := range sortedResources { + amount := req.Resources[resource] + _, err = conn.ExecContext(ctx, ` + INSERT INTO quota_reservations (id, resource, kind, key, amount, repo_did, owner_did, phase, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + `, resID, resource, string(req.Kind), req.Key, amount, req.Identity.RepoDID, req.Identity.OwnerDID, phase, time.Now().Unix()) + if err != nil { + return err + } + } + + hasFiniteLimit := func(scope quota.Scope, did string, resource string) (bool, error) { + // the did is unique across repo and user axes + var override sql.NullInt64 + err := conn.QueryRowContext(ctx, ` + SELECT max_amount FROM quota_limits + WHERE did = ? AND resource = ? + `, did, resource).Scan(&override) + switch { + case err == nil: + return override.Valid, nil + case !errors.Is(err, sql.ErrNoRows): + return false, err + case d.defaults == nil || d.defaults[scope] == nil: + return false, nil + default: + return d.defaults[scope][resource] > 0, nil + } + } + + reason := quota.ReasonUnlimited + for _, resource := range sortedResources { + repoFinite, err := hasFiniteLimit(quota.ScopeRepo, req.Identity.RepoDID, resource) + if err != nil { + return err + } + userFinite, err := hasFiniteLimit(quota.ScopeUser, req.Identity.OwnerDID, resource) + if err != nil { + return err + } + if repoFinite || userFinite { + reason = quota.ReasonWithinLimit + break + } + } + + res = quota.Reservation{ + ID: resID, + Allowed: true, + Temporary: false, + Reason: reason, + } + return nil + }) + + if err != nil { + return quota.Reservation{}, err + } + return res, nil +} + +func (d *QuotaStore) BeginPublish(ctx context.Context, reservationID string) error { + if reservationID == "" { + return nil + } + return d.withTx(ctx, func(conn *sql.Conn) error { + var phase string + err := conn.QueryRowContext(ctx, ` + SELECT phase FROM quota_reservations WHERE id = ? LIMIT 1 + `, reservationID).Scan(&phase) + if errors.Is(err, sql.ErrNoRows) { + return nil + } else if err != nil { + return err + } + + if phase == "reserved" { + _, err = conn.ExecContext(ctx, ` + UPDATE quota_reservations SET phase = 'publishing' WHERE id = ? + `, reservationID) + return err + } + return nil + }) +} + +func (d *QuotaStore) Commit(ctx context.Context, reservationID string) error { + if reservationID == "" { + return nil + } + return d.withTx(ctx, func(conn *sql.Conn) error { + rows, err := conn.QueryContext(ctx, ` + SELECT resource, kind, key, amount, repo_did FROM quota_reservations WHERE id = ? + `, reservationID) + if err != nil { + return err + } + defer rows.Close() + + type resItem struct { + resource string + kind string + key string + amount int64 + repo string + } + var items []resItem + for rows.Next() { + var item resItem + if err := rows.Scan(&item.resource, &item.kind, &item.key, &item.amount, &item.repo); err != nil { + return err + } + items = append(items, item) + } + + if len(items) == 0 { + return nil + } + + for _, item := range items { + _, err = conn.ExecContext(ctx, ` + INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(repo_did, resource, kind, key) + DO UPDATE SET amount = excluded.amount + `, item.repo, item.resource, item.kind, item.key, item.amount) + if err != nil { + return err + } + } + + _, err = conn.ExecContext(ctx, ` + DELETE FROM quota_reservations WHERE id = ? + `, reservationID) + return err + }) +} + +func (d *QuotaStore) Release(ctx context.Context, reservationID string) error { + if reservationID == "" { + return nil + } + return d.withTx(ctx, func(conn *sql.Conn) error { + _, err := conn.ExecContext(ctx, ` + DELETE FROM quota_reservations WHERE id = ? + `, reservationID) + return err + }) +} + +func (d *QuotaStore) SetLimit(ctx context.Context, did string, resource string, limit int64) error { + if err := quota.ValidateOverride(did, resource); err != nil { + return err + } + return d.withTx(ctx, func(conn *sql.Conn) error { + var val any + if limit <= 0 { + val = nil + } else { + val = limit + } + _, err := conn.ExecContext(ctx, ` + INSERT OR REPLACE INTO quota_limits (did, resource, max_amount) + VALUES (?, ?, ?) + `, did, resource, val) + return err + }) +} + +func (d *QuotaStore) UnsetLimit(ctx context.Context, did string, resource string) error { + if err := quota.ValidateOverride(did, resource); err != nil { + return err + } + return d.withTx(ctx, func(conn *sql.Conn) error { + _, err := conn.ExecContext(ctx, ` + DELETE FROM quota_limits WHERE did = ? AND resource = ? + `, did, resource) + return err + }) +} + +func (d *QuotaStore) GetLimit(ctx context.Context, did string, resource string) (*quota.Limit, error) { + if err := quota.ValidateOverride(did, resource); err != nil { + return nil, err + } + var limit *quota.Limit + err := d.withReadTx(ctx, func(tx *sql.Tx) error { + var maxAmt *int64 + err := tx.QueryRowContext(ctx, ` + SELECT max_amount FROM quota_limits WHERE did = ? AND resource = ? + `, did, resource).Scan(&maxAmt) + if errors.Is(err, sql.ErrNoRows) { + return nil + } + if err != nil { + return err + } + l := quota.Limit{DID: did, Resource: resource, Limit: -1} + if maxAmt != nil { + l.Limit = *maxAmt + } + limit = &l + return nil + }) + if err != nil { + return nil, err + } + return limit, nil +} + +func (d *QuotaStore) ListLimits(ctx context.Context) ([]quota.Limit, error) { + var limits []quota.Limit + err := d.withReadTx(ctx, func(tx *sql.Tx) error { + rows, err := tx.QueryContext(ctx, ` + SELECT did, resource, max_amount FROM quota_limits ORDER BY did, resource + `) + if err != nil { + return err + } + defer rows.Close() + + for rows.Next() { + var l quota.Limit + var resStr string + var maxAmt *int64 + if err := rows.Scan(&l.DID, &resStr, &maxAmt); err != nil { + return err + } + l.Resource = resStr + if maxAmt != nil { + l.Limit = *maxAmt + } else { + l.Limit = -1 + } + limits = append(limits, l) + } + return nil + }) + if err != nil { + return nil, err + } + return limits, nil +} + +func (d *QuotaStore) ListUsage(ctx context.Context) ([]quota.Usage, error) { + var usages []quota.Usage + err := d.withReadTx(ctx, func(tx *sql.Tx) error { + collect := func(scope quota.Scope, query string) error { + rows, err := tx.QueryContext(ctx, query) + if err != nil { + return err + } + defer rows.Close() + + for rows.Next() { + var did, resource string + var amount int64 + if err := rows.Scan(&did, &resource, &amount); err != nil { + return err + } + if amount == 0 { + continue + } + if len(usages) == 0 || usages[len(usages)-1].Scope != scope || + usages[len(usages)-1].DID != did || usages[len(usages)-1].Resource != resource { + usages = append(usages, quota.Usage{Scope: scope, DID: did, Resource: resource}) + } + usage := &usages[len(usages)-1] + usage.Used, err = addQuotaUsage(usage.Used, amount) + if err != nil { + return fmt.Errorf("%s %s %s: %w", scope, did, resource, err) + } + } + return rows.Err() + } + + if err := collect(quota.ScopeRepo, ` + SELECT repo_did, resource, amount + FROM ( + SELECT repo_did, resource, kind, key, amount + FROM quota_allocations + UNION + SELECT repo_did, resource, kind, key, amount + FROM quota_reservations + WHERE phase IN ('reserved', 'publishing', 'active') + ) + ORDER BY repo_did, resource + `); err != nil { + return err + } + + return collect(quota.ScopeUser, ` + SELECT owner_did, resource, amount + FROM ( + SELECT owners.owner_did, allocations.resource, + allocations.kind, allocations.key, allocations.amount + FROM quota_allocations AS allocations + JOIN quota_repo_owners AS owners + ON allocations.repo_did = owners.repo_did + UNION + SELECT owner_did, resource, kind, key, amount + FROM quota_reservations + WHERE phase IN ('reserved', 'publishing', 'active') + ) + ORDER BY owner_did, resource + `) + }) + if err != nil { + return nil, err + } + return usages, nil +} + +func (d *QuotaStore) MetricsSnapshot(ctx context.Context) (quota.MetricsSnapshot, error) { + snapshot := quota.MetricsSnapshot{ + Usage: map[quota.Scope]map[string]float64{ + quota.ScopeUser: make(map[string]float64), + quota.ScopeRepo: make(map[string]float64), + }, + Subjects: map[quota.Scope]map[string]map[string]int64{ + quota.ScopeUser: make(map[string]map[string]int64), + quota.ScopeRepo: make(map[string]map[string]int64), + }, + } + + resources := make(map[string]struct{}) + for _, defaults := range d.defaults { + for resource := range defaults { + resources[resource] = struct{}{} + } + } + + scopes := []quota.Scope{quota.ScopeUser, quota.ScopeRepo} + statuses := []string{"unlimited", "under_limit", "near_limit", "at_limit", "over_limit"} + + for _, sc := range scopes { + for res := range resources { + snapshot.Subjects[sc][res] = make(map[string]int64) + for _, st := range statuses { + snapshot.Subjects[sc][res][st] = 0 + } + } + } + + limits, err := d.ListLimits(ctx) + if err != nil { + return snapshot, err + } + usages, err := d.ListUsage(ctx) + if err != nil { + return snapshot, err + } + + rows, err := d.db.QueryContext(ctx, `SELECT repo_did, owner_did FROM quota_repo_owners`) + if err != nil { + return snapshot, err + } + repoDids := map[string]struct{}{} + subjects := map[quota.Scope]map[string]struct{}{ + quota.ScopeUser: {}, + quota.ScopeRepo: {}, + } + for rows.Next() { + var repoDID, ownerDID string + if err := rows.Scan(&repoDID, &ownerDID); err != nil { + rows.Close() + return snapshot, err + } + if repoDID != "" { + repoDids[repoDID] = struct{}{} + subjects[quota.ScopeRepo][repoDID] = struct{}{} + } + if ownerDID != "" { + subjects[quota.ScopeUser][ownerDID] = struct{}{} + } + } + if err := rows.Close(); err != nil { + return snapshot, err + } + if err := rows.Err(); err != nil { + return snapshot, err + } + + type subjectResource struct { + scope quota.Scope + did string + resource string + } + limitMap := make(map[[2]string]int64) + for _, l := range limits { + limitMap[[2]string{l.DID, l.Resource}] = l.Limit + // unknown dids use the user axis + scope := quota.ScopeUser + if _, ok := repoDids[l.DID]; ok { + scope = quota.ScopeRepo + } + subjects[scope][l.DID] = struct{}{} + for _, sc := range scopes { + if _, ok := snapshot.Subjects[sc][l.Resource]; !ok { + snapshot.Subjects[sc][l.Resource] = make(map[string]int64) + for _, st := range statuses { + snapshot.Subjects[sc][l.Resource][st] = 0 + } + } + } + } + + usageMap := make(map[subjectResource]int64) + for _, u := range usages { + key := subjectResource{scope: u.Scope, did: u.DID, resource: u.Resource} + usageMap[key] = u.Used + subjects[u.Scope][u.DID] = struct{}{} + snapshot.Usage[u.Scope][u.Resource] += float64(u.Used) + for _, sc := range scopes { + if _, ok := snapshot.Subjects[sc][u.Resource]; !ok { + snapshot.Subjects[sc][u.Resource] = make(map[string]int64) + for _, st := range statuses { + snapshot.Subjects[sc][u.Resource][st] = 0 + } + } + } + } + + for _, scope := range scopes { + for did := range subjects[scope] { + for resource := range snapshot.Subjects[scope] { + key := subjectResource{scope: scope, did: did, resource: resource} + used := usageMap[key] + limit, hasLimit := limitMap[[2]string{did, resource}] + if !hasLimit && d.defaults != nil { + if defaults, ok := d.defaults[scope]; ok { + if value, ok := defaults[resource]; ok && value > 0 { + limit = value + hasLimit = true + } + } + } + + status := "unlimited" + if hasLimit && limit >= 0 { + status = classifyStatus(used, limit) + } + snapshot.Subjects[scope][resource][status]++ + } + } + } + + return snapshot, nil +} + +func classifyStatus(used, limit int64) string { + if limit < 0 { + return "unlimited" + } + if used > limit { + return "over_limit" + } + if used == limit { + return "at_limit" + } + if used >= limit-limit/5 { + return "near_limit" + } + return "under_limit" +} + +func (d *QuotaStore) Recover(ctx context.Context, liveIDs []string) error { + return d.withTx(ctx, func(conn *sql.Conn) error { + _, err := conn.ExecContext(ctx, ` + INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) + SELECT repo_did, resource, kind, key, amount FROM quota_reservations WHERE phase = 'publishing' + ON CONFLICT(repo_did, resource, kind, key) + DO UPDATE SET amount = excluded.amount + `) + if err != nil { + return err + } + + _, err = conn.ExecContext(ctx, ` + DELETE FROM quota_reservations WHERE phase IN ('reserved', 'publishing') + `) + if err != nil { + return err + } + + if len(liveIDs) == 0 { + _, err = conn.ExecContext(ctx, ` + DELETE FROM quota_reservations WHERE phase = 'active' + `) + return err + } + + placeholders := make([]string, len(liveIDs)) + args := make([]any, len(liveIDs)) + for i, id := range liveIDs { + placeholders[i] = "?" + args[i] = id + } + inClause := strings.Join(placeholders, ", ") + + query := fmt.Sprintf(` + DELETE FROM quota_reservations WHERE phase = 'active' AND id NOT IN (%s) + `, inClause) + _, err = conn.ExecContext(ctx, query, args...) + return err + }) +} diff --git a/spindle/db/quota_test.go b/spindle/db/quota_test.go new file mode 100644 index 000000000..f3cc76777 --- /dev/null +++ b/spindle/db/quota_test.go @@ -0,0 +1,1056 @@ +package db + +import ( + "context" + "database/sql" + "errors" + "path/filepath" + "sync" + "testing" + + "tangled.org/core/spindle/quota" +) + +func TestRepoDedup(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res.Allowed || res.ID == "" { + t.Fatalf("expected allowed reservation, got %v", res) + } + + err = qs.BeginPublish(ctx, res.ID) + if err != nil { + t.Fatal(err) + } + err = qs.Commit(ctx, res.ID) + if err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected reservation to be allowed due to repo dedup") + } + if res2.ID != "" { + t.Fatalf("expected empty reservation ID for already committed object, got %q", res2.ID) + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash456", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res3.Allowed { + t.Fatal("expected reservation to be rejected (exceeds limit)") + } +} + +func TestOwnerUnionDedup(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + id2 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo2"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res1.Allowed { + t.Fatal("expected res1 to be allowed") + } + + if err := qs.BeginPublish(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected res2 to be allowed due to owner union dedup") + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash456", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if res3.Allowed { + t.Fatal("expected res3 to be rejected") + } +} + +func TestDifferentOwnerLogicalCharging(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + id2 := quota.Identity{OwnerDID: "did:web:bob", RepoDID: "did:web:bob/repo1"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res1.Allowed { + t.Fatal("expected res1 to be allowed") + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected res2 to be allowed") + } +} + +func TestBothLimits(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res.Allowed || res.Reason != string(quota.ReasonRepoLimit) { + t.Fatalf("expected rejection due to repo limit, got allowed=%t reason=%s", res.Allowed, res.Reason) + } +} + +func TestDefaultOverrideUnlimitedPrecedence(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res.Allowed { + t.Fatal("expected rejection") + } + + if err := qs.SetLimit(ctx, id1.RepoDID, quota.ResourceCacheStorageBytes, 200); err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected allowed after repo override") + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash456", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 120, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res3.Allowed { + t.Fatal("expected rejection due to user limit") + } + + if err := qs.SetLimit(ctx, id1.OwnerDID, quota.ResourceCacheStorageBytes, 0); err != nil { + t.Fatal(err) + } + if err := qs.SetLimit(ctx, id1.RepoDID, quota.ResourceCacheStorageBytes, 0); err != nil { + t.Fatal(err) + } + + res4, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash456", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 120, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res4.Allowed { + t.Fatal("expected zero override to make user quota unlimited") + } + if res4.Reason != quota.ReasonUnlimited { + t.Fatalf("reason = %q, want %q", res4.Reason, quota.ReasonUnlimited) + } +} + +func TestIdempotence(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 40, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 40, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res1.ID != res2.ID { + t.Fatalf("expected identical reservation IDs, got %q and %q", res1.ID, res2.ID) + } + + if err := qs.BeginPublish(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.BeginPublish(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + if err := qs.Commit(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + if err := qs.Release(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.Release(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + if err := qs.BeginPublish(ctx, ""); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, ""); err != nil { + t.Fatal(err) + } + if err := qs.Release(ctx, ""); err != nil { + t.Fatal(err) + } +} + +func TestQuotaTransactionRollsBackAfterContextCancellation(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + qs := NewQuotaStore(database, nil) + + txCtx, cancel := context.WithCancel(ctx) + err = qs.withTx(txCtx, func(*sql.Conn) error { + cancel() + return txCtx.Err() + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("transaction error = %v, want context canceled", err) + } + + if err := qs.SetLimit(ctx, "did:web:alice", quota.ResourceWorkflows, 1); err != nil { + t.Fatalf("write after canceled transaction failed: %v", err) + } +} +func TestCommitUpdatesChangedAmountForSameKey(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, nil) + identity := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"} + reserveAndCommit := func(amount int64) { + t.Helper() + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "nar/object.nar", + Identity: identity, + Resources: quota.Resources{quota.ResourceCacheStorageBytes: amount}, + }) + if err != nil { + t.Fatal(err) + } + if !res.Allowed || res.ID == "" { + t.Fatalf("reservation = %+v, want a new allowed reservation", res) + } + if err := qs.BeginPublish(ctx, res.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res.ID); err != nil { + t.Fatal(err) + } + } + + reserveAndCommit(60) + reserveAndCommit(80) + + var amount int64 + if err := database.QueryRow(` + SELECT amount FROM quota_allocations + WHERE repo_did = ? AND resource = ? AND kind = ? AND key = ? + `, identity.RepoDID, quota.ResourceCacheStorageBytes, quota.KindNixCache, "nar/object.nar").Scan(&amount); err != nil { + t.Fatal(err) + } + if amount != 80 { + t.Fatalf("committed amount = %d, want 80", amount) + } +} + +func TestRecoverUpdatesChangedAmountForPublishingKey(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + qs := NewQuotaStore(database, nil) + identity := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"} + reserve := func(amount int64) quota.Reservation { + t.Helper() + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "nar/object.nar", + Identity: identity, + Resources: quota.Resources{quota.ResourceCacheStorageBytes: amount}, + }) + if err != nil { + t.Fatal(err) + } + if !res.Allowed || res.ID == "" { + t.Fatalf("reservation = %+v, want a new allowed reservation", res) + } + if err := qs.BeginPublish(ctx, res.ID); err != nil { + t.Fatal(err) + } + return res + } + + first := reserve(60) + if err := qs.Commit(ctx, first.ID); err != nil { + t.Fatal(err) + } + reserve(80) + if err := qs.Recover(ctx, nil); err != nil { + t.Fatal(err) + } + + var amount int64 + if err := database.QueryRow(` + SELECT amount FROM quota_allocations + WHERE repo_did = ? AND resource = ? AND kind = ? AND key = ? + `, identity.RepoDID, quota.ResourceCacheStorageBytes, quota.KindNixCache, "nar/object.nar").Scan(&amount); err != nil { + t.Fatal(err) + } + if amount != 80 { + t.Fatalf("recovered amount = %d, want 80", amount) + } +} + +func TestRecoveryPhases(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + const resource = "gpu_slices" + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + resource: 10, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + resource: 10, + }, + }) + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash_reserved", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 10, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash_publishing", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 20, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if err := qs.BeginPublish(ctx, res2.ID); err != nil { + t.Fatal(err) + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "workflow_preserved", + Resources: quota.Resources{ + resource: 1, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if err := qs.Recover(ctx, []string{res3.ID}); err != nil { + t.Fatal(err) + } + + var count int + err = database.QueryRowContext(ctx, ` + SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_reserved' + `).Scan(&count) + if err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatal("expected reserved object to NOT be committed") + } + + err = database.QueryRowContext(ctx, ` + SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_publishing' + `).Scan(&count) + if err != nil { + t.Fatal(err) + } + if count != 1 { + t.Fatal("expected publishing object to be committed") + } + + err = database.QueryRowContext(ctx, ` + SELECT COUNT(*) FROM quota_reservations WHERE id = ? + `, res3.ID).Scan(&count) + if err != nil { + t.Fatal(err) + } + if count != 1 { + t.Fatal("expected preserved reservation to remain in reservations table") + } + + err = database.QueryRowContext(ctx, ` + SELECT COUNT(*) FROM quota_reservations WHERE id IN (?, ?) + `, res1.ID, res2.ID).Scan(&count) + if err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatal("expected unpreserved reservations to be cleaned up") + } +} + +func TestConcurrency(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + var wg sync.WaitGroup + const numGoroutines = 10 + results := make([]quota.Reservation, numGoroutines) + errorsList := make([]error, numGoroutines) + + for i := range numGoroutines { + wg.Add(1) + go func(idx int) { + defer wg.Done() + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash_" + string(rune('a'+idx)), + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 15, + }, + Identity: id1, + }) + results[idx] = res + errorsList[idx] = err + }(i) + } + wg.Wait() + + allowedCount := 0 + rejectedCount := 0 + for i := range numGoroutines { + if errorsList[i] != nil { + t.Fatalf("goroutine %d failed with error: %v", i, errorsList[i]) + } + if results[i].Allowed { + allowedCount++ + } else { + rejectedCount++ + } + } + + if allowedCount > 6 { + t.Fatalf("expected at most 6 allowed reservations, got %d", allowedCount) + } + if allowedCount == 0 { + t.Fatal("expected at least one reservation to be allowed") + } +} + +func TestMetricsSnapshot(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeRepo: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + + repos := []struct { + repo string + owner string + obj string + bytes int64 + override int64 + }{ + {"did:web:alice/repo1", "did:web:alice", "obj1", 50, -2}, + {"did:web:alice/repo2", "did:web:alice", "obj2", 80, -2}, + {"did:web:alice/repo3", "did:web:alice", "obj3", 90, -2}, + {"did:web:alice/repo4", "did:web:alice", "obj4", 100, -2}, + {"did:web:alice/repo5", "did:web:alice", "obj5", 120, -2}, + {"did:web:alice/repo6", "did:web:alice", "obj6", 200, -1}, + } + + for _, r := range repos { + _, err = database.ExecContext(ctx, ` + INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?) + `, r.repo, r.owner) + if err != nil { + t.Fatal(err) + } + _, err = database.ExecContext(ctx, ` + INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) VALUES (?, 'cache_storage_bytes', 'nix_cache', ?, ?) + `, r.repo, r.obj, r.bytes) + if err != nil { + t.Fatal(err) + } + if r.override != -2 { + var val any = nil + if r.override >= 0 { + val = r.override + } + _, err = database.ExecContext(ctx, ` + INSERT INTO quota_limits (did, resource, max_amount) VALUES (?, 'cache_storage_bytes', ?) + `, r.repo, val) + if err != nil { + t.Fatal(err) + } + } + } + _, err = database.ExecContext(ctx, ` + INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?) + `, "did:web:bob/idle", "did:web:bob") + if err != nil { + t.Fatal(err) + } + + snapshot, err := qs.MetricsSnapshot(ctx) + if err != nil { + t.Fatal(err) + } + + repoCounts := snapshot.Subjects[quota.ScopeRepo][quota.ResourceCacheStorageBytes] + if repoCounts["under_limit"] != 2 { + t.Errorf("expected 2 under_limit repos including the idle default-limited repo, got %d", repoCounts["under_limit"]) + } + if repoCounts["near_limit"] != 2 { + t.Errorf("expected 2 near_limit repos, got %d", repoCounts["near_limit"]) + } + if repoCounts["at_limit"] != 1 { + t.Errorf("expected 1 at_limit repo, got %d", repoCounts["at_limit"]) + } + if repoCounts["over_limit"] != 1 { + t.Errorf("expected 1 over_limit repo, got %d", repoCounts["over_limit"]) + } + if repoCounts["unlimited"] != 1 { + t.Errorf("expected 1 unlimited repo, got %d", repoCounts["unlimited"]) + } +} + +func TestUserLimitAllowZeroAdditional(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res1.Allowed { + t.Fatal("expected reservation to be allowed") + } + + if err := qs.SetLimit(ctx, id1.OwnerDID, quota.ResourceCacheStorageBytes, 50); err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected deduplicated reservation to be allowed even when over limit") + } +} + +func TestUnequalOwnerClaimsRejected(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + id2 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo2"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res1.Allowed || res1.ID == "" { + t.Fatalf("expected allowed reservation, got %v", res1) + } + + if err := qs.BeginPublish(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if res2.Allowed { + t.Fatal("expected unequal larger claim to be rejected at user limit") + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if !res3.Allowed { + t.Fatal("expected equal-amount cross-repo dedup claim to be allowed") + } + + res4, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 30, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if !res4.Allowed { + t.Fatal("expected smaller unequal claim to be allowed since total is within limit") + } +} + +func TestUnequalOwnerClaimsCharged(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{ + quota.ScopeUser: quota.Resources{ + quota.ResourceCacheStorageBytes: 200, + }, + }) + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + id2 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo2"} + + res1, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res1.Allowed || res1.ID == "" { + t.Fatalf("expected allowed reservation, got %v", res1) + } + + if err := qs.BeginPublish(ctx, res1.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res1.ID); err != nil { + t.Fatal(err) + } + + res2, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash123", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id2, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed || res2.ID == "" { + t.Fatalf("expected allowed unequal claim, got %v", res2) + } + + if err := qs.BeginPublish(ctx, res2.ID); err != nil { + t.Fatal(err) + } + if err := qs.Commit(ctx, res2.ID); err != nil { + t.Fatal(err) + } + + res3, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash456", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 70, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res3.Allowed { + t.Fatal("expected new claim to be rejected since user limit has been exceeded by charged unequal claims") + } +} + +func TestLimitArithmeticDoesNotOverflow(t *testing.T) { + ctx := context.Background() + database, err := Make(ctx, filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + + qs := NewQuotaStore(database, quota.Defaults{}) + id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"} + const limit = int64(1 << 62) + if err := qs.SetLimit(ctx, id.RepoDID, "capacity", limit); err != nil { + t.Fatal(err) + } + reserve := func(key string, amount int64) bool { + t.Helper() + res, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: key, + Identity: id, + Resources: quota.Resources{"capacity": amount}, + }) + if err != nil { + t.Fatal(err) + } + return res.Allowed + } + if !reserve("large", limit-1) || reserve("overflow", 2) { + t.Fatal("overflow-safe limit admission produced the wrong decisions") + } + + first, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "unlimited-1", + Identity: id, + Resources: quota.Resources{"unlimited": quota.MaxResourceAmount}, + }) + if err != nil || !first.Allowed { + t.Fatalf("first unlimited reservation = %+v, %v", first, err) + } + second, err := qs.Reserve(ctx, quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "unlimited-2", + Identity: id, + Resources: quota.Resources{"unlimited": quota.MaxResourceAmount}, + }) + if err != nil || second.Allowed || !second.Temporary { + t.Fatalf("overflowing unlimited reservation = %+v, %v", second, err) + } + if _, err := qs.ListUsage(ctx); err != nil { + t.Fatalf("safe maximum usage failed to aggregate: %v", err) + } + if _, err := database.ExecContext(ctx, ` + INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) + VALUES (?, 'unlimited', 'generic_cache', 'forced-overflow', ?) + `, id.RepoDID, quota.MaxResourceAmount); err != nil { + t.Fatal(err) + } + if _, err := qs.ListUsage(ctx); err == nil { + t.Fatal("expected pre-existing overflowing usage to return an error") + } +} diff --git a/spindle/quota/manager.go b/spindle/quota/manager.go new file mode 100644 index 000000000..a6c151f17 --- /dev/null +++ b/spindle/quota/manager.go @@ -0,0 +1,841 @@ +package quota + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "strings" + "sync" + "time" +) + +type waiter struct { + ctx context.Context + req ReserveRequest + resource string + leaseChan chan Lease + errChan chan error + start time.Time +} + +type noopLease struct{} + +func (noopLease) ID() string { return "" } +func (noopLease) Release() {} +func (noopLease) ReleaseWithError() error { return nil } + +type managerLease struct { + id string + manager *Manager + mu sync.Mutex + released bool +} + +func (l *managerLease) ID() string { + return l.id +} + +func (l *managerLease) Release() { + l.mu.Lock() + if l.released { + l.mu.Unlock() + return + } + l.released = true + l.mu.Unlock() + + if err := l.manager.releaseLease(l, false); err != nil { + l.manager.releaseOrRetry(l.id) + } +} + +func (l *managerLease) ReleaseWithError() error { + l.mu.Lock() + if l.released { + l.mu.Unlock() + return nil + } + l.mu.Unlock() + + err := l.manager.releaseLease(l, true) + if err == nil { + l.mu.Lock() + l.released = true + l.mu.Unlock() + } + return err +} + +func (m *Manager) releaseLease(l *managerLease, retainOnError bool) error { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + origID, _ := parseID(l.id) + unlock := m.lockID(origID) + defer unlock() + + m.mu.Lock() + if m.activeLeases[origID] != l { + m.mu.Unlock() + return nil + } + delete(m.activeLeases, origID) + m.mu.Unlock() + + err := m.store.Release(context.Background(), origID) + if err != nil && retainOnError { + m.mu.Lock() + if _, exists := m.activeLeases[origID]; !exists { + m.activeLeases[origID] = l + } + m.mu.Unlock() + } + if err == nil { + go m.processQueue() + } + return err +} + +type refLock struct { + mu sync.Mutex + ref int +} + +type Manager struct { + store ReservationStore + observer Observer + lifecycleMu sync.Mutex + mu sync.Mutex + queue []*waiter + stop chan struct{} + pendingReleases []string + activeLeases map[string]*managerLease + idLocks map[string]*refLock +} + +func NewManager(store ReservationStore, retryInterval time.Duration, observer Observer) *Manager { + m := &Manager{ + store: store, + observer: observer, + stop: make(chan struct{}), + activeLeases: make(map[string]*managerLease), + idLocks: make(map[string]*refLock), + } + if retryInterval > 0 { + m.startTicker(retryInterval) + } + return m +} + +func (m *Manager) lockID(id string) func() { + m.mu.Lock() + if m.idLocks == nil { + m.idLocks = make(map[string]*refLock) + } + lk, ok := m.idLocks[id] + if !ok { + lk = &refLock{} + m.idLocks[id] = lk + } + lk.ref++ + m.mu.Unlock() + + lk.mu.Lock() + return func() { + lk.mu.Unlock() + + m.mu.Lock() + lk.ref-- + if lk.ref == 0 { + delete(m.idLocks, id) + } + m.mu.Unlock() + } +} + +func parseID(id string) (string, string) { + parts := strings.SplitN(id, "#", 2) + if len(parts) == 2 { + return parts[0], parts[1] + } + return id, "" +} + +func StorageReservationID(id string) string { + storageID, _ := parseID(id) + return storageID +} + +func (m *Manager) Close() { + m.mu.Lock() + select { + case <-m.stop: + m.mu.Unlock() + return + default: + } + close(m.stop) + + queueCopy := m.queue + m.queue = nil + m.mu.Unlock() + + // fail queued waiters outside the lock + kinds := make(map[string]bool) + for _, w := range queueCopy { + kinds[string(w.req.Kind)] = true + select { + case w.errChan <- errors.New("manager closed"): + default: + } + m.recordDecisionForResources(w.req, false, false, "no_wait_slot_unavailable") + m.recordWait(w, time.Since(w.start)) + } + + for k := range kinds { + m.setQueueDepth(k) + } +} + +func (m *Manager) startTicker(interval time.Duration) { + ticker := time.NewTicker(interval) + go func() { + for { + select { + case <-ticker.C: + m.processQueue() + case <-m.stop: + ticker.Stop() + return + } + } + }() +} + +func (m *Manager) recordDecision(kind, resource string, allowed, temporary bool, reason string) { + if obs := m.observer; obs != nil { + obs.RecordDecision(kind, resource, allowed, temporary, reason) + } +} + +func (m *Manager) recordDecisionForResources(req ReserveRequest, allowed, temporary bool, reason string) { + if obs := m.observer; obs != nil { + for res := range req.Resources { + obs.RecordDecision(string(req.Kind), res, allowed, temporary, reason) + } + } +} +func (m *Manager) recordWait(w *waiter, duration time.Duration) { + if obs := m.observer; obs != nil { + obs.RecordWait(string(w.req.Kind), w.resource, duration) + } +} + +func (m *Manager) setQueueDepth(kind string) { + m.mu.Lock() + var depth int64 + for _, w := range m.queue { + if string(w.req.Kind) == kind && w.ctx.Err() == nil { + depth++ + } + } + m.mu.Unlock() + + if obs := m.observer; obs != nil { + obs.SetWaitDepth(kind, depth) + } +} + +func (m *Manager) hasConflict(req ReserveRequest) bool { + for _, w := range m.queue { + if w.ctx.Err() != nil { + continue + } + if w.req.Identity.OwnerDID != req.Identity.OwnerDID && w.req.Identity.RepoDID != req.Identity.RepoDID { + continue + } + for res := range req.Resources { + if _, ok := w.req.Resources[res]; ok { + return true + } + } + } + return false +} + +func (m *Manager) TryAcquire(ctx context.Context, req ReserveRequest) (Lease, Reservation, error) { + if err := Validate(req); err != nil { + return nil, Reservation{}, err + } + + m.mu.Lock() + select { + case <-m.stop: + m.mu.Unlock() + return nil, Reservation{}, errors.New("manager closed") + default: + } + + if m.hasConflict(req) { + m.mu.Unlock() + res := Reservation{ + Allowed: false, + Temporary: true, + Reason: "queued", + } + m.recordDecisionForResources(req, false, true, "queued") + return nil, res, nil + } + m.mu.Unlock() + var unlock func() + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + + if req.ID != "" { + unlock = m.lockID(req.ID) + } + + res, err := m.store.Reserve(ctx, req) + if err != nil { + if unlock != nil { + unlock() + } + m.recordDecisionForResources(req, false, false, "store_error") + return nil, Reservation{}, err + } + + if res.Allowed { + if res.ID == "" { + if unlock != nil { + unlock() + } + m.recordDecisionForResources(req, true, false, res.Reason) + return noopLease{}, res, nil + } + + if unlock == nil { + unlock = m.lockID(res.ID) + } + token := make([]byte, 4) + _, _ = rand.Read(token) + gen := hex.EncodeToString(token) + modifiedID := fmt.Sprintf("%s#%s", res.ID, gen) + + l := &managerLease{ + id: modifiedID, + manager: m, + } + m.mu.Lock() + if _, exists := m.activeLeases[res.ID]; exists { + m.mu.Unlock() + if unlock != nil { + unlock() + } + res.ID = "" + res.Allowed = false + res.Temporary = true + res.Reason = "queued" + m.recordDecisionForResources(req, false, true, res.Reason) + return nil, res, nil + } + if m.activeLeases == nil { + m.activeLeases = make(map[string]*managerLease) + } + m.activeLeases[res.ID] = l + m.mu.Unlock() + + if unlock != nil { + unlock() + } + + res.ID = modifiedID + m.recordDecisionForResources(req, true, false, res.Reason) + return l, res, nil + } + + if unlock != nil { + unlock() + } + + m.recordDecision(string(req.Kind), res.Resource, false, res.Temporary, res.Reason) + return nil, res, nil +} + +func (m *Manager) Acquire(ctx context.Context, req ReserveRequest) (Lease, error) { + if err := Validate(req); err != nil { + return nil, err + } + + start := time.Now() + + m.mu.Lock() + select { + case <-m.stop: + m.mu.Unlock() + return nil, errors.New("manager closed") + default: + } + + if m.hasConflict(req) { + w := &waiter{ + ctx: ctx, + req: req, + leaseChan: make(chan Lease, 1), + errChan: make(chan error, 1), + start: start, + } + m.queue = append(m.queue, w) + m.mu.Unlock() + + m.recordDecisionForResources(req, false, true, "queued") + m.setQueueDepth(string(req.Kind)) + + go m.processQueue() + + return m.wait(ctx, w, start, req) + } + m.mu.Unlock() + m.lifecycleMu.Lock() + m.mu.Lock() + + res, err := m.store.Reserve(ctx, req) + if err != nil { + m.mu.Unlock() + m.lifecycleMu.Unlock() + m.recordDecisionForResources(req, false, false, "store_error") + return nil, err + } + + if res.Allowed { + if res.ID == "" { + m.mu.Unlock() + m.lifecycleMu.Unlock() + m.recordDecisionForResources(req, true, false, res.Reason) + return noopLease{}, nil + } + + if _, exists := m.activeLeases[res.ID]; exists { + w := &waiter{ + ctx: ctx, + req: req, + resource: res.Resource, + leaseChan: make(chan Lease, 1), + errChan: make(chan error, 1), + start: start, + } + m.queue = append(m.queue, w) + m.mu.Unlock() + m.lifecycleMu.Unlock() + m.recordDecisionForResources(req, false, true, "queued") + m.setQueueDepth(string(req.Kind)) + go m.processQueue() + return m.wait(ctx, w, start, req) + } + + token := make([]byte, 4) + _, _ = rand.Read(token) + gen := hex.EncodeToString(token) + modifiedID := fmt.Sprintf("%s#%s", res.ID, gen) + + l := &managerLease{ + id: modifiedID, + manager: m, + } + if m.activeLeases == nil { + m.activeLeases = make(map[string]*managerLease) + } + m.activeLeases[res.ID] = l + m.mu.Unlock() + m.lifecycleMu.Unlock() + m.recordDecisionForResources(req, true, false, res.Reason) + return l, nil + } + + if !res.Temporary { + m.mu.Unlock() + m.lifecycleMu.Unlock() + m.recordDecision(string(req.Kind), res.Resource, false, false, res.Reason) + return nil, fmt.Errorf("quota denied permanently: %s", res.Reason) + } + + w := &waiter{ + ctx: ctx, + req: req, + resource: res.Resource, + leaseChan: make(chan Lease, 1), + errChan: make(chan error, 1), + start: start, + } + m.queue = append(m.queue, w) + m.mu.Unlock() + m.lifecycleMu.Unlock() + + m.recordDecision(string(req.Kind), res.Resource, false, true, res.Reason) + m.setQueueDepth(string(req.Kind)) + + go m.processQueue() + + return m.wait(ctx, w, start, req) +} + +func (m *Manager) wait(ctx context.Context, w *waiter, start time.Time, req ReserveRequest) (Lease, error) { + select { + case <-ctx.Done(): + m.mu.Lock() + inQueue := false + for i, q := range m.queue { + if q == w { + m.queue = append(m.queue[:i], m.queue[i+1:]...) + inQueue = true + break + } + } + m.mu.Unlock() + + select { + case lease := <-w.leaseChan: + m.mu.Lock() + if m.activeLeases != nil { + origID, _ := parseID(lease.ID()) + delete(m.activeLeases, origID) + } + m.mu.Unlock() + m.releaseOrRetry(lease.ID()) + m.recordDecisionForResources(w.req, false, false, "context_done_after_ready") + default: + if inQueue { + m.recordDecisionForResources(w.req, false, false, "context_done_in_queue") + } + } + + m.recordWait(w, time.Since(start)) + m.setQueueDepth(string(req.Kind)) + return nil, ctx.Err() + case err := <-w.errChan: + return nil, err + case lease := <-w.leaseChan: + if err := ctx.Err(); err != nil { + m.mu.Lock() + if m.activeLeases != nil { + origID, _ := parseID(lease.ID()) + delete(m.activeLeases, origID) + } + m.mu.Unlock() + m.releaseOrRetry(lease.ID()) + m.recordDecisionForResources(w.req, false, false, "context_done_after_ready") + m.recordWait(w, time.Since(start)) + m.setQueueDepth(string(req.Kind)) + return nil, err + } + return lease, nil + } +} + +func (m *Manager) releaseOrRetry(id string) { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + origID, _ := parseID(id) + unlock := m.lockID(origID) + defer unlock() + + m.mu.Lock() + if m.activeLeases != nil { + if _, active := m.activeLeases[origID]; active { + m.mu.Unlock() + return + } + } + m.mu.Unlock() + + err := m.store.Release(context.Background(), origID) + if err != nil { + m.mu.Lock() + if m.activeLeases != nil { + if _, active := m.activeLeases[origID]; active { + m.mu.Unlock() + return + } + } + select { + case <-m.stop: + default: + m.pendingReleases = append(m.pendingReleases, origID) + } + m.mu.Unlock() + go m.processQueue() + } +} + +func (m *Manager) processQueue() { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + m.mu.Lock() + select { + case <-m.stop: + m.mu.Unlock() + return + default: + } + + toRetry := m.pendingReleases + m.pendingReleases = nil + m.mu.Unlock() + + var failed []string + for _, id := range toRetry { + m.mu.Lock() + var active bool + if m.activeLeases != nil { + _, active = m.activeLeases[id] + } + m.mu.Unlock() + if active { + continue + } + + if err := m.store.Release(context.Background(), id); err != nil { + m.mu.Lock() + var activeNow bool + if m.activeLeases != nil { + _, activeNow = m.activeLeases[id] + } + if !activeNow { + failed = append(failed, id) + } + m.mu.Unlock() + } + } + + m.mu.Lock() + select { + case <-m.stop: + if len(failed) > 0 { + m.pendingReleases = append(m.pendingReleases, failed...) + } + m.mu.Unlock() + return + default: + } + if len(failed) > 0 { + m.pendingReleases = append(m.pendingReleases, failed...) + } + blockedOwners := make(map[string]bool) + blockedRepos := make(map[string]bool) + grantedOwners := make(map[string]bool) + kinds := make(map[string]bool) + + var nextQueue []*waiter + var releasesToMake []string + + for i, w := range m.queue { + if w.ctx.Err() != nil { + continue + } + kinds[string(w.req.Kind)] = true + + ownerBlocked := blockedOwners[w.req.Identity.OwnerDID] + repoBlocked := blockedRepos[w.req.Identity.RepoDID] + + if ownerBlocked || repoBlocked { + blockedOwners[w.req.Identity.OwnerDID] = true + blockedRepos[w.req.Identity.RepoDID] = true + nextQueue = append(nextQueue, w) + continue + } + + if grantedOwners[w.req.Identity.OwnerDID] { + hasOtherOwner := false + for _, other := range m.queue[i+1:] { + if other.req.Identity.OwnerDID != w.req.Identity.OwnerDID && other.ctx.Err() == nil { + hasOtherOwner = true + break + } + } + if hasOtherOwner { + nextQueue = append(nextQueue, w) + continue + } + } + + res, err := m.store.Reserve(w.ctx, w.req) + if err != nil { + select { + case w.errChan <- err: + default: + } + continue + } + + if res.Allowed { + var l Lease = noopLease{} + if res.ID != "" { + if _, exists := m.activeLeases[res.ID]; exists { + blockedOwners[w.req.Identity.OwnerDID] = true + blockedRepos[w.req.Identity.RepoDID] = true + nextQueue = append(nextQueue, w) + w.resource = res.Resource + continue + } + token := make([]byte, 4) + _, _ = rand.Read(token) + modifiedID := fmt.Sprintf("%s#%s", res.ID, hex.EncodeToString(token)) + managed := &managerLease{ + id: modifiedID, + manager: m, + } + if m.activeLeases == nil { + m.activeLeases = make(map[string]*managerLease) + } + m.activeLeases[res.ID] = managed + l = managed + } + select { + case w.leaseChan <- l: + grantedOwners[w.req.Identity.OwnerDID] = true + m.recordDecisionForResources(w.req, true, false, "allowed_from_queue") + m.recordWait(w, time.Since(w.start)) + case <-w.ctx.Done(): + if res.ID != "" { + delete(m.activeLeases, res.ID) + releasesToMake = append(releasesToMake, res.ID) + } + } + } else { + if !res.Temporary { + select { + case w.errChan <- fmt.Errorf("quota denied permanently: %s", res.Reason): + default: + } + m.recordDecision(string(w.req.Kind), res.Resource, false, false, res.Reason) + m.recordWait(w, time.Since(w.start)) + } else { + blockedOwners[w.req.Identity.OwnerDID] = true + blockedRepos[w.req.Identity.RepoDID] = true + nextQueue = append(nextQueue, w) + w.resource = res.Resource + } + } + } + + m.queue = nextQueue + + // update gauges after releasing the queue lock + m.mu.Unlock() + for k := range kinds { + m.setQueueDepth(k) + } + + // release abandoned grants without blocking the queue lock + var failedReleases []string + for _, id := range releasesToMake { + if err := m.store.Release(context.Background(), id); err != nil { + failedReleases = append(failedReleases, id) + } + } + + if len(failedReleases) > 0 { + m.mu.Lock() + select { + case <-m.stop: + default: + m.pendingReleases = append(m.pendingReleases, failedReleases...) + } + m.mu.Unlock() + } +} + +func (m *Manager) Release(ctx context.Context, reservationID string) error { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + origID, _ := parseID(reservationID) + unlock := m.lockID(origID) + defer unlock() + + var owned *managerLease + m.mu.Lock() + if m.activeLeases != nil { + activeL, exists := m.activeLeases[origID] + owned = activeL + if exists && activeL.id != reservationID { + m.mu.Unlock() + return nil + } + if exists { + delete(m.activeLeases, origID) + } + } + m.mu.Unlock() + + err := m.store.Release(ctx, origID) + if err != nil { + if owned == nil { + owned = &managerLease{id: reservationID, manager: m} + } + m.mu.Lock() + if _, exists := m.activeLeases[origID]; !exists { + m.activeLeases[origID] = owned + } + m.mu.Unlock() + } + if err == nil { + go m.processQueue() + } + return err +} + +func (m *Manager) Commit(ctx context.Context, reservationID string) error { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + origID, _ := parseID(reservationID) + unlock := m.lockID(origID) + defer unlock() + + var owned *managerLease + m.mu.Lock() + if m.activeLeases != nil { + activeL, exists := m.activeLeases[origID] + owned = activeL + if exists && activeL.id != reservationID { + m.mu.Unlock() + return nil + } + if exists { + delete(m.activeLeases, origID) + } + } + m.mu.Unlock() + + err := m.store.Commit(ctx, origID) + if err != nil { + if owned == nil { + owned = &managerLease{id: reservationID, manager: m} + } + m.mu.Lock() + if _, exists := m.activeLeases[origID]; !exists { + m.activeLeases[origID] = owned + } + m.mu.Unlock() + } + if err == nil { + go m.processQueue() + } + return err +} + +func (m *Manager) BeginPublish(ctx context.Context, reservationID string) error { + m.lifecycleMu.Lock() + defer m.lifecycleMu.Unlock() + origID, _ := parseID(reservationID) + unlock := m.lockID(origID) + defer unlock() + + m.mu.Lock() + active, exists := m.activeLeases[origID] + if exists && active.id != reservationID { + m.mu.Unlock() + return errors.New("quota reservation ownership changed") + } + m.mu.Unlock() + return m.store.BeginPublish(ctx, origID) +} diff --git a/spindle/quota/manager_test.go b/spindle/quota/manager_test.go new file mode 100644 index 000000000..b3e1bc746 --- /dev/null +++ b/spindle/quota/manager_test.go @@ -0,0 +1,1204 @@ +package quota_test + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "strings" + "sync" + "testing" + "time" + + "tangled.org/core/spindle/quota" +) + +type mockStore struct { + mu sync.Mutex + limits map[string]int64 + allocations map[string]int64 + reservations map[string]quota.ReserveRequest + phases map[string]string + repoOwners map[string]string + defaultUser int64 + defaultRepo int64 +} + +func newMockStore(defaultUser, defaultRepo int64) *mockStore { + return &mockStore{ + limits: make(map[string]int64), + allocations: make(map[string]int64), + reservations: make(map[string]quota.ReserveRequest), + phases: make(map[string]string), + repoOwners: make(map[string]string), + defaultUser: defaultUser, + defaultRepo: defaultRepo, + } +} + +var _ quota.Store = (*mockStore)(nil) + +func (m *mockStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + if err := quota.Validate(req); err != nil { + return quota.Reservation{}, err + } + + m.mu.Lock() + defer m.mu.Unlock() + + m.repoOwners[req.Identity.RepoDID] = req.Identity.OwnerDID + + // committed keys remain idempotent + for res := range req.Resources { + allocKey := req.Identity.RepoDID + "/" + string(res) + "/" + string(req.Kind) + "/" + req.Key + if _, ok := m.allocations[allocKey]; ok { + return quota.Reservation{ + ID: "", + Allowed: true, + Temporary: false, + Reason: quota.ReasonUnlimited, + }, nil + } + } + + // reserved keys remain idempotent + for id, resReq := range m.reservations { + if resReq.Identity.RepoDID == req.Identity.RepoDID && + resReq.Kind == req.Kind && + resReq.Key == req.Key && + (m.phases[id] == "reserved" || m.phases[id] == "publishing") { + return quota.Reservation{ + ID: id, + Allowed: true, + Temporary: false, + Reason: quota.ReasonUnlimited, + }, nil + } + } + for res, amount := range req.Resources { + + repoLimitKey := req.Identity.RepoDID + "/" + string(res) + var repoLimit int64 = -1 + if lim, ok := m.limits[repoLimitKey]; ok { + repoLimit = lim + } else if res == quota.ResourceCacheStorageBytes && m.defaultRepo > 0 { + repoLimit = m.defaultRepo + } + + if repoLimit >= 0 { + var currentUsage int64 + seenKeys := make(map[string]bool) + for k, amt := range m.allocations { + parts := m.split(k) + if parts[0] == req.Identity.RepoDID && parts[1] == string(res) { + seenKeys[parts[3]] = true + currentUsage += amt + } + } + for id, r := range m.reservations { + if r.Identity.RepoDID == req.Identity.RepoDID && (m.phases[id] == "reserved" || m.phases[id] == "publishing") { + if amt, ok := r.Resources[res]; ok { + if !seenKeys[r.Key] { + seenKeys[r.Key] = true + currentUsage += amt + } + } + } + } + + if currentUsage+amount > repoLimit { + temporary := true + if amount > repoLimit { + temporary = false + } + return quota.Reservation{ + ID: "", + Allowed: false, + Temporary: temporary, + Reason: quota.ReasonRepoLimit, + Resource: res, + }, nil + } + } + + userLimitKey := req.Identity.OwnerDID + "/" + string(res) + var userLimit int64 = -1 + if lim, ok := m.limits[userLimitKey]; ok { + userLimit = lim + } else if res == quota.ResourceCacheStorageBytes && m.defaultUser > 0 { + userLimit = m.defaultUser + } + + if userLimit >= 0 { + var currentUsage int64 + seenKeys := make(map[string]bool) + for k, amt := range m.allocations { + parts := m.split(k) + owner := m.repoOwners[parts[0]] + if owner == req.Identity.OwnerDID && parts[1] == string(res) { + if !seenKeys[parts[3]] { + seenKeys[parts[3]] = true + currentUsage += amt + } + } + } + for id, r := range m.reservations { + if r.Identity.OwnerDID == req.Identity.OwnerDID && (m.phases[id] == "reserved" || m.phases[id] == "publishing") { + if amt, ok := r.Resources[res]; ok { + if !seenKeys[r.Key] { + seenKeys[r.Key] = true + currentUsage += amt + } + } + } + } + + // identical owner-level claims do not consume quota twice + keyExists := false + for k := range m.allocations { + parts := m.split(k) + owner := m.repoOwners[parts[0]] + if owner == req.Identity.OwnerDID && parts[1] == string(res) && parts[3] == req.Key { + keyExists = true + break + } + } + if !keyExists { + for id, r := range m.reservations { + if r.Identity.OwnerDID == req.Identity.OwnerDID && r.Key == req.Key && (m.phases[id] == "reserved" || m.phases[id] == "publishing") { + if _, ok := r.Resources[res]; ok { + keyExists = true + break + } + } + } + } + + var additional int64 + if !keyExists { + additional = amount + } + + if currentUsage+additional > userLimit { + temporary := true + if additional > userLimit { + temporary = false + } + return quota.Reservation{ + ID: "", + Allowed: false, + Temporary: temporary, + Reason: quota.ReasonUserLimit, + Resource: res, + }, nil + } + } + } + + id := req.ID + if id == "" { + var b [16]byte + _, _ = rand.Read(b[:]) + id = hex.EncodeToString(b[:]) + } + + m.reservations[id] = req + m.phases[id] = "reserved" + + repoLimitKey := req.Identity.RepoDID + "/cache_storage_bytes" + userLimitKey := req.Identity.OwnerDID + "/cache_storage_bytes" + _, hasRepoLimit := m.limits[repoLimitKey] + _, hasUserLimit := m.limits[userLimitKey] + + reason := quota.ReasonUnlimited + if hasRepoLimit || hasUserLimit || m.defaultRepo > 0 || m.defaultUser > 0 { + reason = quota.ReasonWithinLimit + } + + return quota.Reservation{ + ID: id, + Allowed: true, + Temporary: false, + Reason: reason, + }, nil +} + +func (m *mockStore) split(k string) []string { + idx := m.findResourceIndex(k) + if idx != -1 { + repo := k[:idx] + rest := k[idx+1:] + sub := m.splitRest(rest) + parts := []string{repo} + return append(parts, sub...) + } + return []string{k, "", "", ""} +} + +func (m *mockStore) findResourceIndex(k string) int { + resources := []string{ + "/cache_storage_bytes/", + "/workflows/", + "/vcpus/", + "/memory_mib/", + "/disk_mib/", + } + for _, res := range resources { + idx := findSubstr(k, res) + if idx != -1 { + return idx + } + } + return -1 +} + +func findSubstr(s, sub string) int { + n := len(s) + m := len(sub) + if n < m { + return -1 + } + for i := 0; i <= n-m; i++ { + if s[i:i+m] == sub { + return i + } + } + return -1 +} + +func (m *mockStore) splitRest(r string) []string { + idx1 := -1 + for i := 0; i < len(r); i++ { + if r[i] == '/' { + idx1 = i + break + } + } + if idx1 == -1 { + return []string{r, "", ""} + } + res := r[:idx1] + rest := r[idx1+1:] + idx2 := -1 + for i := 0; i < len(rest); i++ { + if rest[i] == '/' { + idx2 = i + break + } + } + if idx2 == -1 { + return []string{res, rest, ""} + } + kind := rest[:idx2] + key := rest[idx2+1:] + return []string{res, kind, key} +} + +func (m *mockStore) BeginPublish(ctx context.Context, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.reservations[id]; ok { + m.phases[id] = "publishing" + } + return nil +} + +func (m *mockStore) Commit(ctx context.Context, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + req, ok := m.reservations[id] + if !ok { + return nil + } + for res, amt := range req.Resources { + allocKey := req.Identity.RepoDID + "/" + string(res) + "/" + string(req.Kind) + "/" + req.Key + m.allocations[allocKey] = amt + } + delete(m.reservations, id) + delete(m.phases, id) + return nil +} + +func (m *mockStore) Release(ctx context.Context, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.reservations, id) + delete(m.phases, id) + return nil +} + +func (m *mockStore) MetricsSnapshot(ctx context.Context) (quota.MetricsSnapshot, error) { + return quota.MetricsSnapshot{}, errors.New("unimplemented") +} +func (m *mockStore) SetLimit(ctx context.Context, did string, resource string, limit int64) error { + m.mu.Lock() + defer m.mu.Unlock() + key := did + "/" + resource + m.limits[key] = limit + return nil +} +func (m *mockStore) UnsetLimit(ctx context.Context, did string, resource string) error { + m.mu.Lock() + defer m.mu.Unlock() + key := did + "/" + resource + delete(m.limits, key) + return nil +} +func (m *mockStore) GetLimit(ctx context.Context, did string, resource string) (*quota.Limit, error) { + m.mu.Lock() + defer m.mu.Unlock() + limit, ok := m.limits[did+"/"+resource] + if !ok { + return nil, nil + } + return "a.Limit{DID: did, Resource: resource, Limit: limit}, nil +} + +func (m *mockStore) ListLimits(ctx context.Context) ([]quota.Limit, error) { + return nil, errors.New("unimplemented") +} +func (m *mockStore) ListUsage(ctx context.Context) ([]quota.Usage, error) { + return nil, errors.New("unimplemented") +} +func (m *mockStore) Recover(ctx context.Context, liveIDs []string) error { + m.mu.Lock() + defer m.mu.Unlock() + + liveMap := make(map[string]bool) + for _, id := range liveIDs { + liveMap[id] = true + } + + for id, req := range m.reservations { + if liveMap[id] { + continue + } + if m.phases[id] == "publishing" { + for res, amt := range req.Resources { + allocKey := req.Identity.RepoDID + "/" + res + "/" + string(req.Kind) + "/" + req.Key + m.allocations[allocKey] = amt + } + } + delete(m.reservations, id) + delete(m.phases, id) + } + return nil +} + +func TestManagerValidation(t *testing.T) { + req := quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "job", + Identity: quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:repo"}, + Resources: quota.Resources{"gpu_slices": 1}, + } + if err := quota.Validate(req); err != nil { + t.Fatalf("generic resource rejected: %v", err) + } + for i := range quota.MaxResourcePairs { + req.Resources[strings.Repeat("x", i+1)] = 1 + } + if err := quota.Validate(req); err == nil { + t.Fatal("expected too many resources to be rejected") + } + req.Resources = quota.Resources{strings.Repeat("x", quota.MaxResourceKeyBytes+1): 1} + if err := quota.Validate(req); err == nil { + t.Fatal("expected an oversized resource name to be rejected") + } + for _, resource := range []string{strings.Repeat("é", 33), string([]byte{0xff}), "esc\x1b[2J", "new\nline", "tab\tname", "del\x7fname"} { + req.Resources = quota.Resources{resource: 1} + if err := quota.Validate(req); err == nil { + t.Fatalf("expected invalid resource name %q to be rejected", resource) + } + } + req.Resources = quota.Resources{"large": quota.MaxResourceAmount + 1} + if err := quota.Validate(req); err == nil { + t.Fatal("expected an oversized resource amount to be rejected") + } + if err := quota.ValidateOverride("did:web:alice", "gpu_slices"); err != nil { + t.Fatalf("generic override rejected: %v", err) + } +} + +func TestManagerPrecedenceAndDenials(t *testing.T) { + store := newMockStore(100, 50) + mgr := quota.NewManager(store, 50*time.Millisecond, nil) + defer mgr.Close() + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + _, res, err := mgr.TryAcquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res.Allowed || res.Temporary { + t.Fatalf("expected permanent denial, got %+v", res) + } + + _ = store.SetLimit(context.Background(), id1.RepoDID, quota.ResourceCacheStorageBytes, 200) + + lease2, res2, err := mgr.TryAcquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash2", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if !res2.Allowed { + t.Fatal("expected allowed reservation") + } + if lease2 == nil { + t.Fatal("expected non-nil lease") + } + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + _, err = mgr.Acquire(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash3", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id1, + }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected DeadlineExceeded, got %v", err) + } +} + +func TestManagerConcurrencyAndFairness(t *testing.T) { + store := newMockStore(100, 100) + mgr := quota.NewManager(store, 10*time.Millisecond, nil) + defer mgr.Close() + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + res1, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + if res1 == nil { + t.Fatal("expected non-nil lease") + } + + _ = store.SetLimit(context.Background(), id1.OwnerDID, quota.ResourceCacheStorageBytes, 90) + + type acquireResult struct { + lease quota.Lease + err error + } + resultA := make(chan acquireResult, 1) + resultB := make(chan acquireResult, 1) + go func() { + lease, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hashA", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id1, + }) + resultA <- acquireResult{lease: lease, err: err} + }() + + time.Sleep(20 * time.Millisecond) + + go func() { + lease, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hashB", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id1, + }) + resultB <- acquireResult{lease: lease, err: err} + }() + + time.Sleep(20 * time.Millisecond) + + err = mgr.Release(context.Background(), res1.ID()) + if err != nil { + t.Fatal(err) + } + + var acquiredA acquireResult + select { + case acquiredA = <-resultA: + if acquiredA.err != nil || acquiredA.lease == nil { + t.Fatalf("expected waiter A to be allowed, got lease=%v, err=%v", acquiredA.lease, acquiredA.err) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for waiter A") + } + select { + case acquiredB := <-resultB: + t.Fatalf("waiter B completed before waiter A released capacity: lease=%v, err=%v", acquiredB.lease, acquiredB.err) + default: + } + + var hashAResID string + store.mu.Lock() + for id, r := range store.reservations { + if r.Key == "hashA" { + hashAResID = id + break + } + } + store.mu.Unlock() + + if hashAResID == "" { + t.Fatal("could not find reservation ID for hashA") + } + + acquiredA.lease.Release() + + select { + case acquiredB := <-resultB: + if acquiredB.err != nil || acquiredB.lease == nil { + t.Fatalf("expected waiter B to be allowed, got lease=%v, err=%v", acquiredB.lease, acquiredB.err) + } + acquiredB.lease.Release() + case <-time.After(time.Second): + t.Fatal("timed out waiting for waiter B") + } +} + +func TestManagerCancellation(t *testing.T) { + store := newMockStore(100, 100) + mgr := quota.NewManager(store, 50*time.Millisecond, nil) + defer mgr.Close() + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + _, _ = mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id1, + }) + + ctx, cancel := context.WithCancel(context.Background()) + var errA error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + _, errA = mgr.Acquire(ctx, quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hashA", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id1, + }) + }() + + time.Sleep(20 * time.Millisecond) + cancel() + + wg.Wait() + if !errors.Is(errA, context.Canceled) { + t.Fatalf("expected context Canceled, got %v", errA) + } + + store.mu.Lock() + found := false + for _, r := range store.reservations { + if r.Key == "hashA" { + found = true + } + } + store.mu.Unlock() + if found { + t.Fatal("expected hashA reservation to be cleaned up and not leaked") + } +} + +func TestExternalLimitWake(t *testing.T) { + store := newMockStore(100, 100) + mgr := quota.NewManager(store, 20*time.Millisecond, nil) + defer mgr.Close() + + id1 := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + lease0, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash0", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 80, + }, + Identity: id1, + }) + if err != nil { + t.Fatal(err) + } + defer lease0.Release() + + var allowed bool + var eErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + lease, e := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 30, + }, + Identity: id1, + }) + allowed = (e == nil && lease != nil) + eErr = e + }() + + time.Sleep(20 * time.Millisecond) + + _ = store.SetLimit(context.Background(), id1.RepoDID, quota.ResourceCacheStorageBytes, 200) + _ = store.SetLimit(context.Background(), id1.OwnerDID, quota.ResourceCacheStorageBytes, 200) + + wg.Wait() + + if !allowed || eErr != nil { + t.Fatalf("expected waiter to be woken up and allowed, got allowed=%t, err=%v", allowed, eErr) + } +} + +func TestWorkflowReservationID(t *testing.T) { + id1 := quota.WorkflowReservationID("run-1", "owner", "repo", "knot", "rkey", "name") + id2 := quota.WorkflowReservationID("run-1", "owner", "repo", "knot", "rkey", "name") + if id1 != id2 { + t.Errorf("expected deterministic IDs, got %q and %q", id1, id2) + } + + if len(id1) != 73 { + t.Errorf("expected length 73, got %d for %q", len(id1), id1) + } + + // length-prefixing keeps shifted field boundaries distinct + idShift1 := quota.WorkflowReservationID("run-1", "ab", "c", "knot", "rkey", "name") + idShift2 := quota.WorkflowReservationID("run-1", "a", "bc", "knot", "rkey", "name") + if idShift1 == idShift2 { + t.Errorf("expected different IDs for shifted boundaries, both got %q", idShift1) + } + + idDiff := quota.WorkflowReservationID("run-1", "owner", "repo", "knot", "rkey", "name-changed") + if id1 == idDiff { + t.Errorf("expected different IDs on changed argument, both got %q", id1) + } + + idOtherRun := quota.WorkflowReservationID("run-2", "owner", "repo", "knot", "rkey", "name") + if id1 == idOtherRun { + t.Errorf("expected different IDs for distinct runs, both got %q", id1) + } +} + +type waitObserver struct { + mu sync.Mutex + resources []string +} + +func (o *waitObserver) RecordDecision(string, string, bool, bool, string) {} +func (o *waitObserver) SetWaitDepth(string, int64) {} +func (o *waitObserver) RecordWait(kind, resource string, duration time.Duration) { + o.mu.Lock() + defer o.mu.Unlock() + o.resources = append(o.resources, resource) +} + +func TestManagerWaitMetricUsesBlockingResource(t *testing.T) { + store := newMockStore(0, 0) + id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"} + if err := store.SetLimit(context.Background(), id.RepoDID, quota.ResourceVCPUs, 2); err != nil { + t.Fatal(err) + } + + observer := &waitObserver{} + mgr := quota.NewManager(store, 5*time.Millisecond, observer) + defer mgr.Close() + + lease, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "first", + Identity: id, + Resources: quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceVCPUs: 2, + }, + }) + if err != nil { + t.Fatal(err) + } + defer lease.Release() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + _, err = mgr.Acquire(ctx, quota.ReserveRequest{ + Kind: quota.KindWorkflow, + Key: "second", + Identity: id, + Resources: quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceVCPUs: 1, + }, + }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected deadline exceeded, got %v", err) + } + + observer.mu.Lock() + defer observer.mu.Unlock() + if len(observer.resources) != 1 || observer.resources[0] != quota.ResourceVCPUs { + t.Fatalf("wait resources = %v, want [%s]", observer.resources, quota.ResourceVCPUs) + } +} + +type failReleaseStore struct { + *mockStore + mu sync.Mutex + failCounts map[string]int + cancelFn context.CancelFunc +} + +func (s *failReleaseStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + res, err := s.mockStore.Reserve(ctx, req) + if err == nil && req.Key == "hash2" && res.Allowed { + s.mu.Lock() + s.failCounts[res.ID] = 2 + if s.cancelFn != nil { + s.cancelFn() + } + s.mu.Unlock() + } + return res, err +} + +func (s *failReleaseStore) Release(ctx context.Context, id string) error { + s.mu.Lock() + n := s.failCounts[id] + if n > 0 { + s.failCounts[id]-- + s.mu.Unlock() + return errors.New("transient release failure") + } + s.mu.Unlock() + return s.mockStore.Release(ctx, id) +} + +func TestManagerAbandonedReleaseRetry(t *testing.T) { + mock := newMockStore(100, 100) + ctx, cancel := context.WithCancel(context.Background()) + store := &failReleaseStore{ + mockStore: mock, + failCounts: make(map[string]int), + cancelFn: cancel, + } + mgr := quota.NewManager(store, 100*time.Millisecond, nil) + defer mgr.Close() + + id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + req1 := quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + Identity: id, + } + lease1, err := mgr.Acquire(context.Background(), req1) + if err != nil { + t.Fatal(err) + } + + req2 := quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash2", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 50, + }, + Identity: id, + } + + errChan := make(chan error, 1) + go func() { + _, err := mgr.Acquire(ctx, req2) + errChan <- err + }() + + // wait until the request has joined the queue before releasing capacity + time.Sleep(20 * time.Millisecond) + + // cancellation after a grant must release the abandoned reservation + lease1.Release() + + err = <-errChan + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + + mock.mu.Lock() + var resID string + for id, req := range mock.reservations { + if req.Key == "hash2" { + resID = id + break + } + } + mock.mu.Unlock() + + if resID == "" { + t.Fatal("expected reservation to be created for req2") + } + + time.Sleep(5 * time.Millisecond) + mock.mu.Lock() + _, exists := mock.reservations[resID] + mock.mu.Unlock() + if !exists { + t.Fatal("expected reservation to still exist due to failed release") + } + + time.Sleep(40 * time.Millisecond) + mock.mu.Lock() + _, exists = mock.reservations[resID] + mock.mu.Unlock() + if !exists { + t.Fatal("expected reservation to still exist before ticker runs") + } + + deadline := time.Now().Add(1 * time.Second) + released := false + for time.Now().Before(deadline) { + mock.mu.Lock() + _, exists = mock.reservations[resID] + mock.mu.Unlock() + if !exists { + released = true + break + } + time.Sleep(10 * time.Millisecond) + } + if !released { + t.Fatal("expected reservation to be released after successful ticker retry") + } +} + +type trackingStore struct { + *mockStore + mu sync.Mutex + releaseCalls map[string]int + failRelease bool +} + +func (s *trackingStore) Release(ctx context.Context, id string) error { + s.mu.Lock() + s.releaseCalls[id]++ + fail := s.failRelease + s.mu.Unlock() + + if fail { + return errors.New("transient release failure") + } + return s.mockStore.Release(ctx, id) +} + +func TestConcurrentReleaseAttempts(t *testing.T) { + mock := newMockStore(100, 100) + store := &trackingStore{ + mockStore: mock, + releaseCalls: make(map[string]int), + failRelease: true, + } + // keep the ticker out of the release-attempt count + mgr := quota.NewManager(store, 10*time.Second, nil) + defer mgr.Close() + + id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + lease, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 10, + }, + Identity: id, + }) + if err != nil { + t.Fatal(err) + } + + // concurrent lease release remains idempotent + var wg sync.WaitGroup + numConcurrent := 10 + wg.Add(numConcurrent) + for range numConcurrent { + go func() { + defer wg.Done() + lease.Release() + }() + } + wg.Wait() + + store.mu.Lock() + origID := strings.Split(lease.ID(), "#")[0] + calls := store.releaseCalls[origID] + store.mu.Unlock() + + if calls < 2 || calls > 3 { + t.Fatalf("expected 2 or 3 store.Release calls (1 initial + 1 immediate retry + optional background retry), got %d", calls) + } + + store.mu.Lock() + store.failRelease = false + store.mu.Unlock() + + err = mgr.Release(context.Background(), "dummy") + if err != nil { + t.Fatal(err) + } + + time.Sleep(50 * time.Millisecond) + + mock.mu.Lock() + origID = strings.Split(lease.ID(), "#")[0] + _, exists := mock.reservations[origID] + mock.mu.Unlock() + if exists { + t.Fatal("expected lease reservation to be successfully released after retry") + } +} + +type fairnessMockStore struct { + *mockStore + reserveStartChan chan struct{} + reserveHoldChan chan struct{} + reserveBChan chan struct{} + onceStart sync.Once + onceB sync.Once +} + +func (s *fairnessMockStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + switch req.Key { + case "hashA": + s.onceStart.Do(func() { + close(s.reserveStartChan) + <-s.reserveHoldChan + }) + case "hashB": + s.onceB.Do(func() { + close(s.reserveBChan) + }) + } + return s.mockStore.Reserve(ctx, req) +} + +func TestAcquireFairnessTransition(t *testing.T) { + mock := newMockStore(100, 100) + store := &fairnessMockStore{ + mockStore: mock, + reserveStartChan: make(chan struct{}), + reserveHoldChan: make(chan struct{}), + reserveBChan: make(chan struct{}), + } + mgr := quota.NewManager(store, 10*time.Millisecond, nil) + defer mgr.Close() + + id := quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} + + lease1, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hash1", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 60, + }, + Identity: id, + }) + if err != nil { + t.Fatal(err) + } + + // hold the earlier request inside the store while a later request arrives + errChanA := make(chan error, 1) + var leaseA quota.Lease + go func() { + l, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hashA", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + Identity: id, + }) + leaseA = l + errChanA <- err + }() + + <-store.reserveStartChan + + go lease1.Release() + + errChanB := make(chan error, 1) + go func() { + _, err := mgr.Acquire(context.Background(), quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: "hashB", + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: 100, + }, + Identity: id, + }) + errChanB <- err + }() + + select { + case <-store.reserveBChan: + close(store.reserveHoldChan) + t.Fatal("later request reached the store before the earlier reservation completed") + case <-time.After(20 * time.Millisecond): + } + + close(store.reserveHoldChan) + + select { + case err := <-errChanA: + if err != nil { + t.Fatalf("expected A to succeed, got error: %v", err) + } + if leaseA == nil { + t.Fatal("expected non-nil lease for A") + } + case <-time.After(500 * time.Millisecond): + t.Fatal("timeout waiting for A to finish") + } + + select { + case err := <-errChanB: + t.Fatalf("later request completed while the earlier request held capacity: %v", err) + default: + } + if leaseA != nil { + leaseA.Release() + } +} + +func TestManagerFencesDuplicateReservationOwners(t *testing.T) { + store := newMockStore(100, 100) + mgr := quota.NewManager(store, 0, nil) + defer mgr.Close() + req := quota.ReserveRequest{ + ID: "fixed-reservation", + Kind: quota.KindNixCache, + Key: "same-object", + Identity: quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"}, + Resources: quota.Resources{quota.ResourceCacheStorageBytes: 10}, + } + + lease1, res1, err := mgr.TryAcquire(context.Background(), req) + if err != nil || lease1 == nil || !res1.Allowed { + t.Fatalf("first acquisition = lease %v, reservation %+v, error %v", lease1, res1, err) + } + lease2, res2, err := mgr.TryAcquire(context.Background(), req) + if err != nil { + t.Fatal(err) + } + if lease2 != nil || res2.Allowed || !res2.Temporary { + t.Fatalf("duplicate active acquisition = lease %v, reservation %+v, want a temporary denial", lease2, res2) + } + + lease1.Release() + lease3, res3, err := mgr.TryAcquire(context.Background(), req) + if err != nil || lease3 == nil || !res3.Allowed { + t.Fatalf("replacement acquisition = lease %v, reservation %+v, error %v", lease3, res3, err) + } + if res1.ID == res3.ID { + t.Fatalf("replacement owner reused fenced ID %q", res1.ID) + } + if quota.StorageReservationID(res1.ID) != quota.StorageReservationID(res3.ID) { + t.Fatalf("fenced IDs refer to different storage reservations: %q and %q", res1.ID, res3.ID) + } + if err := mgr.BeginPublish(context.Background(), res1.ID); err == nil { + t.Fatal("superseded owner began publication") + } + if err := mgr.BeginPublish(context.Background(), res3.ID); err != nil { + t.Fatalf("current owner could not begin publication: %v", err) + } + + lease1.Release() + storageID := quota.StorageReservationID(res3.ID) + store.mu.Lock() + _, stillReserved := store.reservations[storageID] + store.mu.Unlock() + if !stillReserved { + t.Fatal("superseded owner released the current reservation") + } + lease3.Release() +} + +func TestManagerQueuesDuplicateBlockingAcquisition(t *testing.T) { + store := newMockStore(100, 100) + mgr := quota.NewManager(store, 5*time.Millisecond, nil) + defer mgr.Close() + req := quota.ReserveRequest{ + ID: "fixed-reservation", + Kind: quota.KindWorkflow, + Key: "same-workflow", + Identity: quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo"}, + Resources: quota.Resources{quota.ResourceWorkflows: 1}, + } + + first, _, err := mgr.TryAcquire(context.Background(), req) + if err != nil || first == nil { + t.Fatalf("first acquisition = lease %v, error %v", first, err) + } + result := make(chan struct { + lease quota.Lease + err error + }, 1) + go func() { + lease, err := mgr.Acquire(context.Background(), req) + result <- struct { + lease quota.Lease + err error + }{lease: lease, err: err} + }() + + select { + case got := <-result: + t.Fatalf("duplicate acquisition completed before release: lease %v, error %v", got.lease, got.err) + case <-time.After(30 * time.Millisecond): + } + + first.Release() + select { + case got := <-result: + if got.err != nil || got.lease == nil { + t.Fatalf("queued acquisition = lease %v, error %v", got.lease, got.err) + } + got.lease.Release() + case <-time.After(time.Second): + t.Fatal("queued duplicate acquisition did not complete after release") + } +} diff --git a/spindle/quota/quota.go b/spindle/quota/quota.go new file mode 100644 index 000000000..c44b00038 --- /dev/null +++ b/spindle/quota/quota.go @@ -0,0 +1,203 @@ +package quota + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "strings" + "time" + "unicode" + "unicode/utf8" +) + +type Scope string + +const ( + ScopeUser Scope = "user" + ScopeRepo Scope = "repo" +) + +const ( + ResourceCacheStorageBytes = "cache_storage_bytes" + ResourceWorkflows = "workflows" + ResourceVCPUs = "vcpus" + ResourceMemoryMiB = "memory_mib" + ResourceDiskMiB = "disk_mib" +) + +const ( + MaxResourcePairs = 16 + MaxResourceKeyBytes = 64 + MaxResourceAmount = int64(1 << 62) +) + +type Kind string + +const ( + KindWorkflow Kind = "workflow" + KindNixCache Kind = "nix_cache" + KindGenericCache Kind = "generic_cache" +) + +const ( + ReasonUnlimited = "unlimited" + ReasonWithinLimit = "within_limit" + ReasonUserLimit = "user_limit" + ReasonRepoLimit = "repo_limit" + ReasonAllowedImmediately = "allowed_immediately" +) + +type Identity struct { + OwnerDID string + RepoDID string +} + +type Resources = map[string]int64 + +type Defaults map[Scope]Resources + +type ReserveRequest struct { + ID string + Kind Kind + Key string + Identity Identity + Resources Resources +} + +type Reservation struct { + ID string + Allowed bool + Temporary bool + Reason string + Resource string +} + +type Lease interface { + ID() string + Release() + ReleaseWithError() error +} + +type ReservationStore interface { + Reserve(ctx context.Context, req ReserveRequest) (Reservation, error) + BeginPublish(ctx context.Context, reservationID string) error + Commit(ctx context.Context, reservationID string) error + Release(ctx context.Context, reservationID string) error +} + +// the did is unique across repo and user axes +type Limit struct { + DID string + Resource string + Limit int64 +} + +type Usage struct { + Scope Scope + DID string + Resource string + Used int64 +} + +type MetricsSnapshot struct { + Usage map[Scope]map[string]float64 + Subjects map[Scope]map[string]map[string]int64 +} + +type Observer interface { + RecordDecision(kind, resource string, allowed, temporary bool, reason string) + SetWaitDepth(kind string, depth int64) + RecordWait(kind, resource string, duration time.Duration) +} + +type Store interface { + ReservationStore + Recover(ctx context.Context, liveIDs []string) error + SetLimit(ctx context.Context, did string, resource string, limit int64) error + GetLimit(ctx context.Context, did string, resource string) (*Limit, error) + UnsetLimit(ctx context.Context, did string, resource string) error + ListLimits(ctx context.Context) ([]Limit, error) + ListUsage(ctx context.Context) ([]Usage, error) + MetricsSnapshot(ctx context.Context) (MetricsSnapshot, error) +} + +func Validate(req ReserveRequest) error { + if req.Kind != KindWorkflow && req.Kind != KindNixCache && req.Kind != KindGenericCache { + return fmt.Errorf("invalid kind: %s", req.Kind) + } + if req.Key == "" { + return errors.New("empty key") + } + if len(req.Resources) == 0 { + return errors.New("empty resources") + } + if err := ValidateResources(req.Resources); err != nil { + return err + } + if req.Identity.OwnerDID == "" { + return errors.New("empty owner DID") + } + if req.Identity.RepoDID == "" { + return errors.New("empty repo DID") + } + if !strings.HasPrefix(req.Identity.OwnerDID, "did:") { + return fmt.Errorf("invalid owner DID format: %s", req.Identity.OwnerDID) + } + if !strings.HasPrefix(req.Identity.RepoDID, "did:") { + return fmt.Errorf("invalid repo DID format: %s", req.Identity.RepoDID) + } + return nil +} + +func ValidateOverride(did string, resource string) error { + if did == "" || !strings.HasPrefix(did, "did:") { + return fmt.Errorf("invalid DID: %s", did) + } + if err := validateResourceName(resource); err != nil { + return err + } + return nil +} + +func ValidateResources(resources Resources) error { + if len(resources) > MaxResourcePairs { + return fmt.Errorf("too many resource pairs: %d (max %d)", len(resources), MaxResourcePairs) + } + for resource, amount := range resources { + if err := validateResourceName(resource); err != nil { + return err + } + if amount < 0 || amount > MaxResourceAmount { + return fmt.Errorf("invalid resource amount for %s: %d", resource, amount) + } + } + return nil +} + +func validateResourceName(resource string) error { + if resource == "" || len(resource) > MaxResourceKeyBytes || !utf8.ValidString(resource) { + return fmt.Errorf("invalid resource: %q", resource) + } + // control runes break terminals and log scrapers + for _, r := range resource { + if unicode.IsControl(r) { + return fmt.Errorf("invalid resource: %q", resource) + } + } + return nil +} + +func WorkflowReservationID(runID, owner, repoDid, knot, rkey, name string) string { + raw := fmt.Sprintf("%d:%s:%d:%s:%d:%s:%d:%s:%d:%s:%d:%s", + len(runID), runID, + len(owner), owner, + len(repoDid), repoDid, + len(knot), knot, + len(rkey), rkey, + len(name), name, + ) + hash := sha256.Sum256([]byte(raw)) + return "workflow:" + hex.EncodeToString(hash[:]) +}