Your images will be stored on the selected hold
diff --git a/pkg/hold/admin/handlers.go b/pkg/hold/admin/handlers.go index f5e51c6..3dfe801 100644 --- a/pkg/hold/admin/handlers.go +++ b/pkg/hold/admin/handlers.go @@ -93,22 +93,17 @@ func (ui *AdminUI) handleStatsAPI(w http.ResponseWriter, r *http.Request) { return } - // Calculate total storage by summing quota for all crew members + // Calculate total storage with a single bulk query var totalSize int64 uniqueDigests := 0 - crew, err := ui.pds.ListCrewMembers(ctx) + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) if err != nil { - slog.Warn("Failed to list crew for stats", "error", err) + slog.Warn("Failed to get all user quotas", "error", err) } else { - // Get usage for each crew member - for _, member := range crew { - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) - if err != nil { - continue - } - totalSize += quotaStats.TotalSize - uniqueDigests += quotaStats.UniqueBlobs + for _, q := range allQuotas { + totalSize += q.TotalSize + uniqueDigests += q.UniqueBlobs } } @@ -152,28 +147,22 @@ func (ui *AdminUI) handleTopUsersAPI(w http.ResponseWriter, r *http.Request) { } } - // Get all crew members and their usage - crew, err := ui.pds.ListCrewMembers(ctx) + // Get all user quotas in a single bulk query + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) if err != nil { - slog.Error("Failed to list crew members for top users", "error", err) + slog.Error("Failed to get all user quotas", "error", err) http.Error(w, "Failed to load top users", http.StatusInternalServerError) return } var users []UserUsage - for _, member := range crew { - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) - if err != nil { - slog.Warn("Failed to get quota for user", "did", member.Record.Member, "error", err) - continue - } - + for did, q := range allQuotas { users = append(users, UserUsage{ - DID: member.Record.Member, - Handle: resolveHandle(ctx, member.Record.Member), - Usage: quotaStats.TotalSize, - UsageHuman: formatHumanBytes(quotaStats.TotalSize), - BlobCount: quotaStats.UniqueBlobs, + DID: did, + Handle: resolveHandle(ctx, did), + Usage: q.TotalSize, + UsageHuman: formatHumanBytes(q.TotalSize), + BlobCount: q.UniqueBlobs, }) } diff --git a/pkg/hold/admin/handlers_crew.go b/pkg/hold/admin/handlers_crew.go index 9361054..9a73eef 100644 --- a/pkg/hold/admin/handlers_crew.go +++ b/pkg/hold/admin/handlers_crew.go @@ -9,6 +9,7 @@ import ( "time" "atcr.io/pkg/atproto" + "atcr.io/pkg/hold/pds" "github.com/go-chi/chi/v5" ) @@ -56,14 +57,11 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) { return nil, err } - userUsage := make(map[string]int64) - for _, member := range crew { - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) - if err != nil { - slog.Warn("Failed to get quota for crew member", "did", member.Record.Member, "error", err) - continue - } - userUsage[member.Record.Member] = quotaStats.TotalSize + // Single bulk query for all user quotas + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) + if err != nil { + slog.Warn("Failed to get all user quotas for crew views", "error", err) + allQuotas = make(map[string]*pds.QuotaStats) } defaultTier := "default" @@ -88,11 +86,16 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) { AddedAt: parseTime(member.Record.AddedAt), } + usage := int64(0) + if q, ok := allQuotas[member.Record.Member]; ok { + usage = q.TotalSize + } + if ui.quotaMgr != nil && ui.quotaMgr.IsEnabled() { if limit := ui.quotaMgr.GetTierLimit(tier); limit != nil { view.TierLimit = formatHumanBytes(*limit) if *limit > 0 { - view.UsagePercent = int(float64(userUsage[view.DID]) / float64(*limit) * 100) + view.UsagePercent = int(float64(usage) / float64(*limit) * 100) } } else { view.TierLimit = "Unlimited" @@ -101,7 +104,7 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) { view.TierLimit = "Unlimited" } - view.CurrentUsage = userUsage[view.DID] + view.CurrentUsage = usage view.UsageHuman = formatHumanBytes(view.CurrentUsage) crewViews = append(crewViews, view) diff --git a/pkg/hold/pds/delete_test.go b/pkg/hold/pds/delete_test.go index 4c77f4e..8da2c53 100644 --- a/pkg/hold/pds/delete_test.go +++ b/pkg/hold/pds/delete_test.go @@ -215,7 +215,7 @@ func TestDeleteAndListBlueskyPosts_WithPosts(t *testing.T) { // Index the record if pds.recordsIndex != nil { - err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID) + err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0) if err != nil { t.Fatalf("Failed to index test post: %v", err) } @@ -304,7 +304,7 @@ func TestDeleteUserData_IncludesPosts(t *testing.T) { // Index the record if pds.recordsIndex != nil { - err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID) + err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0) if err != nil { t.Fatalf("Failed to index test post: %v", err) } diff --git a/pkg/hold/pds/layer.go b/pkg/hold/pds/layer.go index 60e6819..114b3d4 100644 --- a/pkg/hold/pds/layer.go +++ b/pkg/hold/pds/layer.go @@ -89,100 +89,48 @@ type QuotaStats struct { Tier string `json:"tier,omitempty"` // quota tier (e.g., 'deckhand', 'bosun', 'quartermaster') } -// GetQuotaForUser calculates storage quota for a specific user -// It iterates through all layer records, filters by userDid, deduplicates by digest, -// and sums the sizes of unique blobs. +// GetQuotaForUser calculates storage quota for a specific user. +// Uses SQL aggregation over the denormalized digest/size columns in the records index. func (p *HoldPDS) GetQuotaForUser(ctx context.Context, userDID string) (*QuotaStats, error) { if p.recordsIndex == nil { return nil, fmt.Errorf("records index not available") } - // Get session for reading record data - session, err := p.carstore.ReadOnlySession(p.uid) + uniqueBlobs, totalSize, err := p.recordsIndex.QuotaForDID(atproto.LayerCollection, userDID) if err != nil { - return nil, fmt.Errorf("failed to create session: %w", err) - } - - head, err := p.carstore.GetUserRepoHead(ctx, p.uid) - if err != nil { - return nil, fmt.Errorf("failed to get repo head: %w", err) - } - - if !head.Defined() { - // Empty repo - return zero stats - return &QuotaStats{UserDID: userDID}, nil - } - - repoHandle, err := repo.OpenRepo(ctx, session, head) - if err != nil { - return nil, fmt.Errorf("failed to open repo: %w", err) - } - - // Track unique digests and their sizes - digestSizes := make(map[string]int64) - - // Iterate all layer records via the index - cursor := "" - batchSize := 1000 // Process in batches - - for { - records, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true) - if err != nil { - return nil, fmt.Errorf("failed to list layer records: %w", err) - } - - for _, rec := range records { - // Construct record path and get the record data - recordPath := rec.Collection + "/" + rec.Rkey - - _, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath) - if err != nil { - // Skip records we can't read - continue - } - - // Decode the layer record - recordValue, err := lexutil.CborDecodeValue(*recBytes) - if err != nil { - continue - } - - layerRecord, ok := recordValue.(*atproto.LayerRecord) - if !ok { - continue - } - - // Filter by userDID - if layerRecord.UserDID != userDID { - continue - } - - // Deduplicate by digest - keep the size (could be different pushes of same blob) - // Store the size - we only count each unique digest once - if _, exists := digestSizes[layerRecord.Digest]; !exists { - digestSizes[layerRecord.Digest] = layerRecord.Size - } - } - - if nextCursor == "" { - break - } - cursor = nextCursor - } - - // Calculate totals - var totalSize int64 - for _, size := range digestSizes { - totalSize += size + return nil, fmt.Errorf("failed to query quota: %w", err) } return &QuotaStats{ UserDID: userDID, - UniqueBlobs: len(digestSizes), + UniqueBlobs: uniqueBlobs, TotalSize: totalSize, }, nil } +// GetAllUserQuotas returns quota stats for all users in a single SQL query. +// Used by admin endpoints to avoid N+1 per-user quota lookups. +func (p *HoldPDS) GetAllUserQuotas(ctx context.Context) (map[string]*QuotaStats, error) { + if p.recordsIndex == nil { + return nil, fmt.Errorf("records index not available") + } + + quotas, err := p.recordsIndex.QuotasByDID(atproto.LayerCollection) + if err != nil { + return nil, fmt.Errorf("failed to query all quotas: %w", err) + } + + result := make(map[string]*QuotaStats, len(quotas)) + for did, q := range quotas { + result[did] = &QuotaStats{ + UserDID: did, + UniqueBlobs: q.UniqueBlobs, + TotalSize: q.TotalSize, + } + } + return result, nil +} + // GetQuotaForUserWithTier calculates quota with tier-aware limits // It returns the base quota stats plus the tier limit and tier name. // Captain (owner) always has unlimited quota. @@ -262,27 +210,24 @@ func (p *HoldPDS) ListLayerRecordsForUser(ctx context.Context, userDID string) ( var records []*atproto.LayerRecord - // Iterate all layer records via the index + // Iterate layer records for this user via the index (filtered by DID in SQL) cursor := "" - batchSize := 1000 // Process in batches + batchSize := 1000 for { - indexRecords, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true) + indexRecords, nextCursor, err := p.recordsIndex.ListRecordsByDID(atproto.LayerCollection, userDID, batchSize, cursor) if err != nil { return nil, fmt.Errorf("failed to list layer records: %w", err) } for _, rec := range indexRecords { - // Construct record path and get the record data recordPath := rec.Collection + "/" + rec.Rkey _, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath) if err != nil { - // Skip records we can't read continue } - // Decode the layer record recordValue, err := lexutil.CborDecodeValue(*recBytes) if err != nil { continue @@ -293,11 +238,6 @@ func (p *HoldPDS) ListLayerRecordsForUser(ctx context.Context, userDID string) ( continue } - // Filter by userDID - if layerRecord.UserDID != userDID { - continue - } - records = append(records, layerRecord) } diff --git a/pkg/hold/pds/records.go b/pkg/hold/pds/records.go index d26e6bb..e73ef40 100644 --- a/pkg/hold/pds/records.go +++ b/pkg/hold/pds/records.go @@ -36,6 +36,8 @@ CREATE TABLE IF NOT EXISTS records ( rkey TEXT NOT NULL, cid TEXT NOT NULL, did TEXT, + digest TEXT, + size INTEGER, PRIMARY KEY (collection, rkey) ); CREATE INDEX IF NOT EXISTS idx_records_collection_rkey ON records(collection, rkey); @@ -54,18 +56,26 @@ func NewRecordsIndex(dbPath string) (*RecordsIndex, error) { return nil, fmt.Errorf("failed to open records database: %w", err) } - // Check if table exists and has the did column + // Check if table exists and has required columns needsRebuild := false var tableName string err = db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName) if err == nil { - // Table exists, check for did column var colCount int + // Check for did column err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount) if err != nil || colCount == 0 { needsRebuild = true slog.Info("Records index schema outdated, rebuilding with did column") } + // Check for digest column (added for SQL-based quota queries) + if !needsRebuild { + err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount) + if err != nil || colCount == 0 { + needsRebuild = true + slog.Info("Records index schema outdated, rebuilding with digest/size columns") + } + } } if needsRebuild { @@ -91,17 +101,26 @@ func NewRecordsIndex(dbPath string) (*RecordsIndex, error) { // NewRecordsIndexWithDB creates a records index using an existing *sql.DB connection. // The caller is responsible for the DB lifecycle. func NewRecordsIndexWithDB(db *sql.DB) (*RecordsIndex, error) { - // Check if table exists and has the did column + // Check if table exists and has required columns needsRebuild := false var tableName string err := db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName) if err == nil { var colCount int + // Check for did column err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount) if err != nil || colCount == 0 { needsRebuild = true slog.Info("Records index schema outdated, rebuilding with did column") } + // Check for digest column (added for SQL-based quota queries) + if !needsRebuild { + err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount) + if err != nil || colCount == 0 { + needsRebuild = true + slog.Info("Records index schema outdated, rebuilding with digest/size columns") + } + } } if needsRebuild { @@ -129,12 +148,15 @@ func (ri *RecordsIndex) Close() error { } // IndexRecord adds or updates a record in the index -// did parameter is optional - pass empty string if not applicable -func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did string) error { +// did, digest, size parameters are optional - pass empty/zero if not applicable +func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did, digest string, size int64) error { _, err := ri.db.Exec(` - INSERT OR REPLACE INTO records (collection, rkey, cid, did) - VALUES (?, ?, ?, ?) - `, collection, rkey, cidStr, sql.NullString{String: did, Valid: did != ""}) + INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size) + VALUES (?, ?, ?, ?, ?, ?) + `, collection, rkey, cidStr, + sql.NullString{String: did, Valid: did != ""}, + sql.NullString{String: digest, Valid: digest != ""}, + sql.NullInt64{Int64: size, Valid: size > 0}) return err } @@ -328,8 +350,8 @@ func (ri *RecordsIndex) BackfillFromRepo(ctx context.Context, repoHandle *repo.R defer tx.Rollback() stmt, err := tx.Prepare(` - INSERT OR REPLACE INTO records (collection, rkey, cid, did) - VALUES (?, ?, ?, ?) + INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size) + VALUES (?, ?, ?, ?, ?, ?) `) if err != nil { return fmt.Errorf("failed to prepare statement: %w", err) @@ -345,14 +367,19 @@ func (ri *RecordsIndex) BackfillFromRepo(ctx context.Context, repoHandle *repo.R } collection, rkey := parts[0], parts[1] - // Extract DID from record content based on collection type - var did string + // Extract fields from record content based on collection type + var did, digest string + var size int64 _, recBytes, err := repoHandle.GetRecordBytes(ctx, key) if err == nil && recBytes != nil { did = extractDIDFromRecord(collection, *recBytes) + digest, size = extractLayerFieldsFromRecord(collection, *recBytes) } - _, err = stmt.Exec(collection, rkey, c.String(), sql.NullString{String: did, Valid: did != ""}) + _, err = stmt.Exec(collection, rkey, c.String(), + sql.NullString{String: did, Valid: did != ""}, + sql.NullString{String: digest, Valid: digest != ""}, + sql.NullInt64{Int64: size, Valid: size > 0}) if err != nil { return fmt.Errorf("failed to index record %s: %w", key, err) } @@ -392,6 +419,60 @@ func splitStatements(sql string) []string { return out } +// QuotaForDID returns unique blob count and total size for a single DID in a collection. +// Uses SQL aggregation over the denormalized digest/size columns. +func (ri *RecordsIndex) QuotaForDID(collection, did string) (uniqueBlobs int, totalSize int64, err error) { + err = ri.db.QueryRow(` + SELECT COUNT(*), COALESCE(SUM(size), 0) + FROM (SELECT DISTINCT digest, size FROM records WHERE collection = ? AND did = ? AND digest IS NOT NULL) + `, collection, did).Scan(&uniqueBlobs, &totalSize) + return +} + +// QuotasByDID returns unique blob count and total size grouped by DID. +// Single query replaces N individual quota lookups. +func (ri *RecordsIndex) QuotasByDID(collection string) (map[string]QuotaResult, error) { + rows, err := ri.db.Query(` + SELECT did, COUNT(*), COALESCE(SUM(size), 0) + FROM (SELECT DISTINCT did, digest, size FROM records WHERE collection = ? AND did IS NOT NULL AND digest IS NOT NULL) + GROUP BY did + `, collection) + if err != nil { + return nil, fmt.Errorf("failed to query quotas by DID: %w", err) + } + defer rows.Close() + + result := make(map[string]QuotaResult) + for rows.Next() { + var did string + var qr QuotaResult + if err := rows.Scan(&did, &qr.UniqueBlobs, &qr.TotalSize); err != nil { + return nil, fmt.Errorf("failed to scan quota row: %w", err) + } + result[did] = qr + } + return result, rows.Err() +} + +// QuotaResult holds aggregated quota data from SQL queries +type QuotaResult struct { + UniqueBlobs int + TotalSize int64 +} + +// extractLayerFieldsFromRecord extracts digest and size from layer records. +// Returns empty/zero for non-layer records. +func extractLayerFieldsFromRecord(collection string, recBytes []byte) (string, int64) { + if collection != atproto.LayerCollection { + return "", 0 + } + var rec atproto.LayerRecord + if err := rec.UnmarshalCBOR(bytes.NewReader(recBytes)); err != nil { + return "", 0 + } + return rec.Digest, rec.Size +} + // extractDIDFromRecord extracts the associated DID from a record based on its collection type func extractDIDFromRecord(collection string, recBytes []byte) string { switch collection { diff --git a/pkg/hold/pds/records_test.go b/pkg/hold/pds/records_test.go index f4cb9de..fc41810 100644 --- a/pkg/hold/pds/records_test.go +++ b/pkg/hold/pds/records_test.go @@ -50,7 +50,7 @@ func TestRecordsIndex_IndexRecord(t *testing.T) { defer ri.Close() // Index a record - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) if err != nil { t.Fatalf("IndexRecord() error = %v", err) } @@ -75,13 +75,13 @@ func TestRecordsIndex_IndexRecord_Upsert(t *testing.T) { defer ri.Close() // Index a record - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) if err != nil { t.Fatalf("IndexRecord() first call error = %v", err) } // Update the same record with new CID - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "") + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "", "", 0) if err != nil { t.Fatalf("IndexRecord() second call error = %v", err) } @@ -118,7 +118,7 @@ func TestRecordsIndex_DeleteRecord(t *testing.T) { defer ri.Close() // Index a record - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) if err != nil { t.Fatalf("IndexRecord() error = %v", err) } @@ -217,7 +217,7 @@ func TestRecordsIndex_ListRecords_Basic(t *testing.T) { {"ccc", "cid3"}, } for _, r := range records { - if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -248,7 +248,7 @@ func TestRecordsIndex_ListRecords_DefaultOrder(t *testing.T) { // Add records with different rkeys (TIDs are lexicographically ordered by time) rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"} for _, rkey := range rkeys { - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -286,7 +286,7 @@ func TestRecordsIndex_ListRecords_ReverseOrder(t *testing.T) { // Add records rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"} for _, rkey := range rkeys { - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -324,7 +324,7 @@ func TestRecordsIndex_ListRecords_Limit(t *testing.T) { // Add 5 records for i := range 5 { rkey := string(rune('a' + i)) - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -355,7 +355,7 @@ func TestRecordsIndex_ListRecords_Cursor(t *testing.T) { // Add 5 records rkeys := []string{"a", "b", "c", "d", "e"} for _, rkey := range rkeys { - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -430,7 +430,7 @@ func TestRecordsIndex_ListRecords_CursorReverse(t *testing.T) { // Add 5 records rkeys := []string{"a", "b", "c", "d", "e"} for _, rkey := range rkeys { - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { t.Fatalf("IndexRecord() error = %v", err) } } @@ -474,10 +474,10 @@ func TestRecordsIndex_Count(t *testing.T) { // Add records to two collections for i := range 3 { - ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "") + ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "", "", 0) } for i := range 5 { - ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "") + ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "", "", 0) } // Count crew @@ -527,10 +527,10 @@ func TestRecordsIndex_TotalCount(t *testing.T) { defer ri.Close() // Add records to multiple collections - ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "") - ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "") - ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "") - ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "") + ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "", "", 0) + ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "", "", 0) + ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "", "", 0) + ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "", "", 0) count, err := ri.TotalCount() if err != nil { @@ -581,9 +581,9 @@ func TestRecordsIndex_MultipleCollections(t *testing.T) { defer ri.Close() // Add records to different collections with same rkeys - ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "") - ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "") - ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "") + ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "", "", 0) + ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "", "", 0) + ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "", "", 0) // Listing should only return records from requested collection records, _, err := ri.ListRecords("io.atcr.hold.crew", 10, "", false) @@ -605,3 +605,138 @@ func TestRecordsIndex_MultipleCollections(t *testing.T) { t.Errorf("Expected captain count 1 after deleting crew, got %d", count) } } + +// TestRecordsIndex_QuotaForDID tests single-user quota aggregation +func TestRecordsIndex_QuotaForDID(t *testing.T) { + tmpDir := t.TempDir() + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) + if err != nil { + t.Fatalf("NewRecordsIndex() error = %v", err) + } + defer ri.Close() + + // Add layer records for two users, with some duplicate digests + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200) + ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:alice", "sha256:aaa", 100) // duplicate digest + ri.IndexRecord("io.atcr.hold.layer", "r4", "cid4", "did:plc:bob", "sha256:ccc", 300) + + // Alice should have 2 unique blobs, total 300 bytes + uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice") + if err != nil { + t.Fatalf("QuotaForDID() error = %v", err) + } + if uniqueBlobs != 2 { + t.Errorf("Expected 2 unique blobs for alice, got %d", uniqueBlobs) + } + if totalSize != 300 { + t.Errorf("Expected total size 300 for alice, got %d", totalSize) + } + + // Bob should have 1 unique blob, total 300 bytes + uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:bob") + if err != nil { + t.Fatalf("QuotaForDID() error = %v", err) + } + if uniqueBlobs != 1 { + t.Errorf("Expected 1 unique blob for bob, got %d", uniqueBlobs) + } + if totalSize != 300 { + t.Errorf("Expected total size 300 for bob, got %d", totalSize) + } + + // Unknown user should have 0 + uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:unknown") + if err != nil { + t.Fatalf("QuotaForDID() error = %v", err) + } + if uniqueBlobs != 0 || totalSize != 0 { + t.Errorf("Expected 0/0 for unknown user, got %d/%d", uniqueBlobs, totalSize) + } +} + +// TestRecordsIndex_QuotasByDID tests bulk quota aggregation +func TestRecordsIndex_QuotasByDID(t *testing.T) { + tmpDir := t.TempDir() + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) + if err != nil { + t.Fatalf("NewRecordsIndex() error = %v", err) + } + defer ri.Close() + + // Add layer records for multiple users + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200) + ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:bob", "sha256:ccc", 500) + // Non-layer record should be excluded + ri.IndexRecord("io.atcr.hold.crew", "r4", "cid4", "did:plc:alice", "", 0) + + quotas, err := ri.QuotasByDID("io.atcr.hold.layer") + if err != nil { + t.Fatalf("QuotasByDID() error = %v", err) + } + + if len(quotas) != 2 { + t.Fatalf("Expected 2 users in quotas, got %d", len(quotas)) + } + + alice := quotas["did:plc:alice"] + if alice.UniqueBlobs != 2 { + t.Errorf("Expected 2 unique blobs for alice, got %d", alice.UniqueBlobs) + } + if alice.TotalSize != 300 { + t.Errorf("Expected total size 300 for alice, got %d", alice.TotalSize) + } + + bob := quotas["did:plc:bob"] + if bob.UniqueBlobs != 1 { + t.Errorf("Expected 1 unique blob for bob, got %d", bob.UniqueBlobs) + } + if bob.TotalSize != 500 { + t.Errorf("Expected total size 500 for bob, got %d", bob.TotalSize) + } +} + +// TestRecordsIndex_QuotasByDID_Empty tests bulk quota with no records +func TestRecordsIndex_QuotasByDID_Empty(t *testing.T) { + tmpDir := t.TempDir() + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) + if err != nil { + t.Fatalf("NewRecordsIndex() error = %v", err) + } + defer ri.Close() + + quotas, err := ri.QuotasByDID("io.atcr.hold.layer") + if err != nil { + t.Fatalf("QuotasByDID() error = %v", err) + } + if len(quotas) != 0 { + t.Errorf("Expected empty quotas map, got %d entries", len(quotas)) + } +} + +// TestRecordsIndex_QuotaForDID_IgnoresNullDigest tests that records without digest are excluded +func TestRecordsIndex_QuotaForDID_IgnoresNullDigest(t *testing.T) { + tmpDir := t.TempDir() + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) + if err != nil { + t.Fatalf("NewRecordsIndex() error = %v", err) + } + defer ri.Close() + + // Record with digest+size + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) + // Record without digest (e.g. old data before migration) + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "", 0) + + uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice") + if err != nil { + t.Fatalf("QuotaForDID() error = %v", err) + } + if uniqueBlobs != 1 { + t.Errorf("Expected 1 unique blob (ignoring null digest), got %d", uniqueBlobs) + } + if totalSize != 100 { + t.Errorf("Expected total size 100, got %d", totalSize) + } +} diff --git a/pkg/hold/pds/server.go b/pkg/hold/pds/server.go index f2da17f..faf26aa 100644 --- a/pkg/hold/pds/server.go +++ b/pkg/hold/pds/server.go @@ -411,9 +411,10 @@ func (p *HoldPDS) CreateRecordsIndexEventHandler(broadcasterHandler func(context if op.RecCid != nil { cidStr = op.RecCid.String() } - // Extract DID from record based on collection type + // Extract fields from record based on collection type did := extractDIDFromOp(op) - if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did); err != nil { + digest, size := extractLayerFieldsFromOp(op) + if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did, digest, size); err != nil { slog.Warn("Failed to index record", "collection", op.Collection, "rkey", op.Rkey, "error", err) } case EvtKindDeleteRecord: @@ -454,6 +455,17 @@ func extractDIDFromOp(op RepoOp) string { return "" } +// extractLayerFieldsFromOp extracts digest and size from a layer record operation +func extractLayerFieldsFromOp(op RepoOp) (string, int64) { + if op.Record == nil || op.Collection != atproto.LayerCollection { + return "", 0 + } + if rec, ok := op.Record.(*atproto.LayerRecord); ok { + return rec.Digest, rec.Size + } + return "", 0 +} + // BackfillRecordsIndex populates the records index from existing MST data func (p *HoldPDS) BackfillRecordsIndex(ctx context.Context) error { if p.recordsIndex == nil {