Something went wrong. Try again.
Monorepo for Tangled tangled.org
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489package db
import ( "context" "database/sql" "errors" "sync" "testing"
"tangled.org/core/spindle/quota")
var ( alice1 = quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo1"} alice2 = quota.Identity{OwnerDID: "did:web:alice", RepoDID: "did:web:alice/repo2"})
func cacheDefaults(user, repo int64) quota.Defaults { defaults := quota.Defaults{} if user != 0 { defaults[quota.ScopeUser] = quota.Resources{quota.ResourceCacheStorageBytes: user} } if repo != 0 { defaults[quota.ScopeRepo] = quota.Resources{quota.ResourceCacheStorageBytes: repo} } return defaults}
func reserveBytes(t *testing.T, qs *QuotaStore, id quota.Identity, key string, amount int64) quota.Reservation { t.Helper() res, err := qs.Reserve(context.Background(), quota.ReserveRequest{ Kind: quota.KindNixCache, Key: key, Identity: id, Resources: quota.Resources{quota.ResourceCacheStorageBytes: amount}, }) if err != nil { t.Fatal(err) } return res}
func commitReservation(t *testing.T, qs *QuotaStore, res quota.Reservation) { t.Helper() if err := qs.BeginCommit(context.Background(), res.ID); err != nil { t.Fatal(err) } if err := qs.Commit(context.Background(), res.ID); err != nil { t.Fatal(err) }}
func committedAmount(t *testing.T, database *DB, id quota.Identity) int64 { t.Helper() var amount int64 if err := database.QueryRow(` SELECT amount FROM quota_allocations WHERE repo_did = ? AND resource = ? AND kind = ? AND key = ? `, id.RepoDID, quota.ResourceCacheStorageBytes, quota.KindNixCache, "nar/object.nar").Scan(&amount); err != nil { t.Fatal(err) } return amount}
func TestRepoDedup(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(0, 100))
res := reserveBytes(t, qs, alice1, "hash123", 60) if !res.Allowed || res.ID == "" { t.Fatalf("expected allowed reservation, got %v", res) } commitReservation(t, qs, res)
res2 := reserveBytes(t, qs, alice1, "hash123", 60) 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) }
if reserveBytes(t, qs, alice1, "hash456", 60).Allowed { t.Fatal("expected reservation to be rejected (exceeds limit)") }}
func TestOwnerUnionDedup(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))
res1 := reserveBytes(t, qs, alice1, "hash123", 60) if !res1.Allowed { t.Fatal("expected res1 to be allowed") } commitReservation(t, qs, res1)
if !reserveBytes(t, qs, alice2, "hash123", 60).Allowed { t.Fatal("expected res2 to be allowed due to owner union dedup") } if reserveBytes(t, qs, alice2, "hash456", 50).Allowed { t.Fatal("expected res3 to be rejected") }}
func TestDifferentOwnerLogicalCharging(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0)) bob := quota.Identity{OwnerDID: "did:web:bob", RepoDID: "did:web:bob/repo1"}
if !reserveBytes(t, qs, alice1, "hash123", 60).Allowed { t.Fatal("expected res1 to be allowed") } if !reserveBytes(t, qs, bob, "hash123", 60).Allowed { t.Fatal("expected res2 to be allowed") }}
func TestBothLimits(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 50))
res := reserveBytes(t, qs, alice1, "hash123", 60) 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() qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 50))
if reserveBytes(t, qs, alice1, "hash123", 60).Allowed { t.Fatal("expected rejection") }
if err := qs.SetLimit(ctx, alice1.RepoDID, quota.ResourceCacheStorageBytes, 200); err != nil { t.Fatal(err) } if !reserveBytes(t, qs, alice1, "hash123", 60).Allowed { t.Fatal("expected allowed after repo override") } if reserveBytes(t, qs, alice1, "hash456", 120).Allowed { t.Fatal("expected rejection due to user limit") }
if err := qs.SetLimit(ctx, alice1.OwnerDID, quota.ResourceCacheStorageBytes, 0); err != nil { t.Fatal(err) } if err := qs.SetLimit(ctx, alice1.RepoDID, quota.ResourceCacheStorageBytes, 0); err != nil { t.Fatal(err) }
res := reserveBytes(t, qs, alice1, "hash456", 120) if !res.Allowed { t.Fatal("expected zero override to make user quota unlimited") } if res.Reason != quota.ReasonUnlimited { t.Fatalf("reason = %q, want %q", res.Reason, quota.ReasonUnlimited) }}
func TestIdempotence(t *testing.T) { ctx := context.Background() qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 100))
res1 := reserveBytes(t, qs, alice1, "hash123", 40) res2 := reserveBytes(t, qs, alice1, "hash123", 40) if res1.ID != res2.ID { t.Fatalf("expected identical reservation IDs, got %q and %q", res1.ID, res2.ID) }
steps := []struct { name string do func(context.Context, string) error }{{"begin commit", qs.BeginCommit}, {"commit", qs.Commit}, {"release", qs.Release}} for _, step := range steps { for _, id := range []string{res1.ID, res1.ID, ""} { if err := step.do(ctx, id); err != nil { t.Fatalf("%s %q: %v", step.name, id, err) } } }}
func TestQuotaTransactionRollsBackAfterContextCancellation(t *testing.T) { ctx := context.Background() qs := NewQuotaStore(newTestDB(t), 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) { database := newTestDB(t) qs := NewQuotaStore(database, nil)
for _, amount := range []int64{60, 80} { res := reserveBytes(t, qs, alice1, "nar/object.nar", amount) if !res.Allowed || res.ID == "" { t.Fatalf("reservation = %+v, want a new allowed reservation", res) } commitReservation(t, qs, res) }
if amount := committedAmount(t, database, alice1); amount != 80 { t.Fatalf("committed amount = %d, want 80", amount) }}
func TestRecoverKeepsTheCommittedAmountOverAnInterruptedRepublish(t *testing.T) { database := newTestDB(t) qs := NewQuotaStore(database, nil)
commitReservation(t, qs, reserveBytes(t, qs, alice1, "nar/object.nar", 60)) interrupted := reserveBytes(t, qs, alice1, "nar/object.nar", 80) if !interrupted.Allowed || interrupted.ID == "" { t.Fatalf("reservation = %+v, want a new allowed reservation", interrupted) } if err := qs.BeginCommit(context.Background(), interrupted.ID); err != nil { t.Fatal(err) } if err := qs.Recover(context.Background(), nil); err != nil { t.Fatal(err) }
if amount := committedAmount(t, database, alice1); amount != 60 { t.Fatalf("recovered amount = %d, want the committed 60", amount) }}
func TestRecoveryPhases(t *testing.T) { ctx := context.Background() database := newTestDB(t)
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}, })
res1 := reserveBytes(t, qs, alice1, "hash_reserved", 10) res2 := reserveBytes(t, qs, alice1, "hash_publishing", 20) if err := qs.BeginCommit(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: alice1, }) if err != nil { t.Fatal(err) } if err := qs.Recover(ctx, []string{res3.ID}); err != nil { t.Fatal(err) }
count := func(query string, args ...any) int { t.Helper() var n int if err := database.QueryRowContext(ctx, query, args...).Scan(&n); err != nil { t.Fatal(err) } return n } if count(`SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_reserved'`) != 0 { t.Fatal("expected reserved object to NOT be committed") } if count(`SELECT COUNT(*) FROM quota_allocations WHERE key = 'hash_publishing'`) != 0 { t.Fatal("expected the publishing reservation to be released, not counted as an object") } if count(`SELECT COUNT(*) FROM quota_reservations WHERE id = ?`, res3.ID) != 1 { t.Fatal("expected preserved reservation to remain in reservations table") } if count(`SELECT COUNT(*) FROM quota_reservations WHERE id IN (?, ?)`, res1.ID, res2.ID) != 0 { t.Fatal("expected unpreserved reservations to be cleaned up") }}
func TestConcurrency(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 100))
var wg sync.WaitGroup const numGoroutines = 10 results := make([]quota.Reservation, numGoroutines) errs := make([]error, numGoroutines)
for i := range numGoroutines { wg.Add(1) go func() { defer wg.Done() results[i], errs[i] = qs.Reserve(context.Background(), quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "hash_" + string(rune('a'+i)), Identity: alice1, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 15}, }) }() } wg.Wait()
allowedCount := 0 for i, res := range results { if errs[i] != nil { t.Fatalf("goroutine %d failed with error: %v", i, errs[i]) } if res.Allowed { allowedCount++ } } 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 := newTestDB(t) qs := NewQuotaStore(database, cacheDefaults(0, 100))
repos := []struct { repo string obj string bytes int64 override int64 }{ {"did:web:alice/repo1", "obj1", 50, -2}, {"did:web:alice/repo2", "obj2", 80, -2}, {"did:web:alice/repo3", "obj3", 90, -2}, {"did:web:alice/repo4", "obj4", 100, -2}, {"did:web:alice/repo5", "obj5", 120, -2}, {"did:web:alice/repo6", "obj6", 200, -1}, }
exec := func(query string, args ...any) { t.Helper() if _, err := database.ExecContext(ctx, query, args...); err != nil { t.Fatal(err) } } for _, r := range repos { exec(`INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?)`, r.repo, "did:web:alice") exec(`INSERT INTO quota_allocations (repo_did, resource, kind, key, amount) VALUES (?, 'cache_storage_bytes', 'nix_cache', ?, ?)`, r.repo, r.obj, r.bytes) if r.override != -2 { var val any if r.override >= 0 { val = r.override } exec(`INSERT INTO quota_limits (did, resource, max_amount) VALUES (?, 'cache_storage_bytes', ?)`, r.repo, val) } } exec(`INSERT INTO quota_repo_owners (repo_did, owner_did) VALUES (?, ?)`, "did:web:bob/idle", "did:web:bob")
snapshot, err := qs.MetricsSnapshot(ctx) if err != nil { t.Fatal(err) }
repoCounts := snapshot.Subjects[quota.ScopeRepo][quota.ResourceCacheStorageBytes] for bucket, want := range map[string]int64{"under_limit": 2, "near_limit": 2, "at_limit": 1, "over_limit": 1, "unlimited": 1} { if repoCounts[bucket] != want { t.Errorf("%s repos = %d, want %d (under_limit includes the idle default-limited repo)", bucket, repoCounts[bucket], want) } }}
func TestUserLimitAllowZeroAdditional(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))
if !reserveBytes(t, qs, alice1, "hash123", 80).Allowed { t.Fatal("expected reservation to be allowed") } if err := qs.SetLimit(context.Background(), alice1.OwnerDID, quota.ResourceCacheStorageBytes, 50); err != nil { t.Fatal(err) } if !reserveBytes(t, qs, alice1, "hash123", 80).Allowed { t.Fatal("expected deduplicated reservation to be allowed even when over limit") }}
func TestUnequalOwnerClaimsRejected(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(100, 0))
res1 := reserveBytes(t, qs, alice1, "hash123", 60) if !res1.Allowed || res1.ID == "" { t.Fatalf("expected allowed reservation, got %v", res1) } commitReservation(t, qs, res1)
if reserveBytes(t, qs, alice2, "hash123", 80).Allowed { t.Fatal("expected unequal larger claim to be rejected at user limit") } if !reserveBytes(t, qs, alice2, "hash123", 60).Allowed { t.Fatal("expected equal-amount cross-repo dedup claim to be allowed") } if !reserveBytes(t, qs, alice2, "hash123", 30).Allowed { t.Fatal("expected smaller unequal claim to be allowed since total is within limit") }}
func TestUnequalOwnerClaimsCharged(t *testing.T) { qs := NewQuotaStore(newTestDB(t), cacheDefaults(200, 0))
res1 := reserveBytes(t, qs, alice1, "hash123", 60) if !res1.Allowed || res1.ID == "" { t.Fatalf("expected allowed reservation, got %v", res1) } commitReservation(t, qs, res1)
res2 := reserveBytes(t, qs, alice2, "hash123", 80) if !res2.Allowed || res2.ID == "" { t.Fatalf("expected allowed unequal claim, got %v", res2) } commitReservation(t, qs, res2)
if reserveBytes(t, qs, alice1, "hash456", 70).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 := newTestDB(t) 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") }}