From f19dfa2716dce031c9d33f3c0f7bb77474954d65 Mon Sep 17 00:00:00 2001 From: Evan Jarrett Date: Wed, 31 Dec 2025 11:04:15 -0600 Subject: [PATCH] move download stats to the hold account so it can persist across different appviews --- cmd/appview/serve.go | 19 +- pkg/appview/db/queries.go | 57 +-- pkg/appview/db/stats_migration.go | 231 +++++++++++++ pkg/appview/jetstream/backfill.go | 4 +- pkg/appview/jetstream/processor.go | 69 +++- pkg/appview/jetstream/processor_test.go | 18 +- pkg/appview/jetstream/stats_cache.go | 100 ++++++ pkg/appview/jetstream/worker.go | 41 ++- pkg/appview/middleware/registry.go | 20 +- pkg/appview/storage/context.go | 8 +- pkg/appview/storage/context_test.go | 54 +-- pkg/appview/storage/manifest_store.go | 143 ++++---- pkg/appview/storage/manifest_store_test.go | 102 +----- pkg/atproto/cbor_gen.go | 381 +++++++++++++++++++++ pkg/atproto/endpoints.go | 6 + pkg/atproto/generate.go | 3 +- pkg/atproto/lexicon.go | 45 +++ pkg/hold/oci/xrpc.go | 236 +++++++++---- pkg/hold/pds/server.go | 3 +- pkg/hold/pds/stats.go | 218 ++++++++++++ 20 files changed, 1397 insertions(+), 361 deletions(-) create mode 100644 pkg/appview/db/stats_migration.go create mode 100644 pkg/appview/jetstream/stats_cache.go create mode 100644 pkg/hold/pds/stats.go diff --git a/cmd/appview/serve.go b/cmd/appview/serve.go index d5d4291..bb267ac 100644 --- a/cmd/appview/serve.go +++ b/cmd/appview/serve.go @@ -149,9 +149,9 @@ func serveRegistry(cmd *cobra.Command, args []string) error { // Set global refresher for middleware middleware.SetGlobalRefresher(refresher) - // Set global database for pull/push metrics tracking - metricsDB := db.NewMetricsDB(uiDatabase) - middleware.SetGlobalDatabase(metricsDB) + // Set global database for hold DID lookups (used by blob routing) + holdDIDDB := db.NewHoldDIDDB(uiDatabase) + middleware.SetGlobalDatabase(holdDIDDB) // Create RemoteHoldAuthorizer for hold authorization with caching holdAuthorizer := auth.NewRemoteHoldAuthorizer(uiDatabase, testMode) @@ -161,6 +161,19 @@ func serveRegistry(cmd *cobra.Command, args []string) error { // Initialize Jetstream workers (background services before HTTP routes) initializeJetstream(uiDatabase, &cfg.Jetstream, defaultHoldDID, testMode, refresher) + // Run stats migration to holds (one-time migration, skipped if already done) + go func() { + // Wait for services to be ready (Docker startup race condition) + time.Sleep(10 * time.Second) + // Create service token getter callback that uses auth.GetOrFetchServiceToken + getServiceToken := func(ctx context.Context, userDID, holdDID, pdsEndpoint string) (string, error) { + return auth.GetOrFetchServiceToken(ctx, refresher, userDID, holdDID, pdsEndpoint) + } + if err := db.MigrateStatsToHolds(context.Background(), uiDatabase, getServiceToken); err != nil { + slog.Warn("Stats migration failed", "error", err) + } + }() + // Create main chi router mainRouter := chi.NewRouter() diff --git a/pkg/appview/db/queries.go b/pkg/appview/db/queries.go index 37c7ed7..5609d9d 100644 --- a/pkg/appview/db/queries.go +++ b/pkg/appview/db/queries.go @@ -1446,17 +1446,17 @@ func GetRepositoryStats(db *sql.DB, did, repository string) (*RepositoryStats, e } // UpsertRepositoryStats inserts or updates repository stats +// Note: star_count is calculated dynamically from the stars table, not stored here func UpsertRepositoryStats(db *sql.DB, stats *RepositoryStats) error { _, err := db.Exec(` - INSERT INTO repository_stats (did, repository, star_count, pull_count, last_pull, push_count, last_push) - VALUES (?, ?, ?, ?, ?, ?, ?) + INSERT INTO repository_stats (did, repository, pull_count, last_pull, push_count, last_push) + VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(did, repository) DO UPDATE SET - star_count = excluded.star_count, pull_count = excluded.pull_count, last_pull = excluded.last_pull, push_count = excluded.push_count, last_push = excluded.last_push - `, stats.DID, stats.Repository, stats.StarCount, stats.PullCount, stats.LastPull, stats.PushCount, stats.LastPush) + `, stats.DID, stats.Repository, stats.PullCount, stats.LastPull, stats.PushCount, stats.LastPush) return err } @@ -1593,30 +1593,6 @@ func DeleteStarsNotInList(db *sql.DB, starrerDID string, foundStars map[string]t return nil } -// IncrementPullCount increments the pull count for a repository -func IncrementPullCount(db *sql.DB, did, repository string) error { - _, err := db.Exec(` - INSERT INTO repository_stats (did, repository, pull_count, last_pull) - VALUES (?, ?, 1, datetime('now')) - ON CONFLICT(did, repository) DO UPDATE SET - pull_count = pull_count + 1, - last_pull = datetime('now') - `, did, repository) - return err -} - -// IncrementPushCount increments the push count for a repository -func IncrementPushCount(db *sql.DB, did, repository string) error { - _, err := db.Exec(` - INSERT INTO repository_stats (did, repository, push_count, last_push) - VALUES (?, ?, 1, datetime('now')) - ON CONFLICT(did, repository) DO UPDATE SET - push_count = push_count + 1, - last_push = datetime('now') - `, did, repository) - return err -} - // parseTimestamp parses a timestamp string with multiple format attempts func parseTimestamp(s string) (time.Time, error) { formats := []string{ @@ -1634,29 +1610,20 @@ func parseTimestamp(s string) (time.Time, error) { return time.Time{}, fmt.Errorf("unable to parse timestamp: %s", s) } -// MetricsDB wraps a sql.DB and implements the metrics interface for middleware -type MetricsDB struct { +// HoldDIDDB wraps a sql.DB and implements the HoldDIDLookup interface for middleware +// This is a minimal wrapper that only provides hold DID lookups for blob routing +type HoldDIDDB struct { db *sql.DB } -// NewMetricsDB creates a new metrics database wrapper -func NewMetricsDB(db *sql.DB) *MetricsDB { - return &MetricsDB{db: db} -} - -// IncrementPullCount increments the pull count for a repository -func (m *MetricsDB) IncrementPullCount(did, repository string) error { - return IncrementPullCount(m.db, did, repository) -} - -// IncrementPushCount increments the push count for a repository -func (m *MetricsDB) IncrementPushCount(did, repository string) error { - return IncrementPushCount(m.db, did, repository) +// NewHoldDIDDB creates a new hold DID database wrapper +func NewHoldDIDDB(db *sql.DB) *HoldDIDDB { + return &HoldDIDDB{db: db} } // GetLatestHoldDIDForRepo returns the hold DID from the most recent manifest for a repository -func (m *MetricsDB) GetLatestHoldDIDForRepo(did, repository string) (string, error) { - return GetLatestHoldDIDForRepo(m.db, did, repository) +func (h *HoldDIDDB) GetLatestHoldDIDForRepo(did, repository string) (string, error) { + return GetLatestHoldDIDForRepo(h.db, did, repository) } // GetFeaturedRepositories fetches top repositories sorted by stars and pulls diff --git a/pkg/appview/db/stats_migration.go b/pkg/appview/db/stats_migration.go new file mode 100644 index 0000000..58b121e --- /dev/null +++ b/pkg/appview/db/stats_migration.go @@ -0,0 +1,231 @@ +package db + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "time" + + "atcr.io/pkg/atproto" +) + +// ServiceTokenGetter is a function type for getting service tokens. +// This avoids importing auth from db (which would create import cycles with tests). +type ServiceTokenGetter func(ctx context.Context, userDID, holdDID, pdsEndpoint string) (string, error) + +// MigrateStatsToHolds migrates existing repository_stats data to hold services. +// This is a one-time migration that runs on startup. +// +// The migration: +// 1. Checks if migration has already completed +// 2. Reads all repository_stats entries +// 3. For each entry, looks up the hold DID from manifests table +// 4. Gets a service token for the user and calls the hold's setStats endpoint +// 5. Marks migration complete after all entries are processed +// +// If a hold is offline, the migration logs a warning and continues. +// The hold will receive real-time stats updates via Jetstream once online. +// +// The getServiceToken parameter is a callback to avoid import cycles with pkg/auth. +func MigrateStatsToHolds(ctx context.Context, db *sql.DB, getServiceToken ServiceTokenGetter) error { + // Check if migration already done + var migrationDone bool + err := db.QueryRowContext(ctx, ` + SELECT EXISTS( + SELECT 1 FROM schema_migrations WHERE version = 1000 + ) + `).Scan(&migrationDone) + + // Table might not exist yet on fresh install + if err == sql.ErrNoRows { + migrationDone = false + } else if err != nil { + // Check if it's a "no such table" error (fresh install) + if err.Error() != "no such table: schema_migrations" { + return fmt.Errorf("failed to check migration status: %w", err) + } + migrationDone = false + } + + if migrationDone { + slog.Debug("Stats migration already complete, skipping", "component", "migration") + return nil + } + + slog.Info("Starting stats migration to holds", "component", "migration") + + // Get all repository_stats entries + rows, err := db.QueryContext(ctx, ` + SELECT did, repository, pull_count, last_pull, push_count, last_push + FROM repository_stats + WHERE pull_count > 0 OR push_count > 0 + `) + if err != nil { + // Table might not exist on fresh install + if err.Error() == "no such table: repository_stats" { + slog.Info("No repository_stats table found, skipping migration", "component", "migration") + return markMigrationComplete(db) + } + return fmt.Errorf("failed to query repository_stats: %w", err) + } + defer rows.Close() + + var stats []struct { + DID string + Repository string + PullCount int64 + LastPull sql.NullString + PushCount int64 + LastPush sql.NullString + } + + for rows.Next() { + var stat struct { + DID string + Repository string + PullCount int64 + LastPull sql.NullString + PushCount int64 + LastPush sql.NullString + } + if err := rows.Scan(&stat.DID, &stat.Repository, &stat.PullCount, &stat.LastPull, &stat.PushCount, &stat.LastPush); err != nil { + return fmt.Errorf("failed to scan stat: %w", err) + } + stats = append(stats, stat) + } + + if len(stats) == 0 { + slog.Info("No stats to migrate", "component", "migration") + return markMigrationComplete(db) + } + + slog.Info("Found stats entries to migrate", "component", "migration", "count", len(stats)) + + // Process each stat + successCount := 0 + skipCount := 0 + errorCount := 0 + + for _, stat := range stats { + // Look up hold DID from manifests table + holdDID, err := GetLatestHoldDIDForRepo(db, stat.DID, stat.Repository) + if err != nil || holdDID == "" { + slog.Debug("No hold DID found for repo, skipping", "component", "migration", + "did", stat.DID, "repository", stat.Repository) + skipCount++ + continue + } + + // Get user's PDS endpoint + user, err := GetUserByDID(db, stat.DID) + if err != nil || user == nil { + slog.Debug("User not found in database, skipping", "component", "migration", + "did", stat.DID, "repository", stat.Repository) + skipCount++ + continue + } + + // Get service token for the user + serviceToken, err := getServiceToken(ctx, stat.DID, holdDID, user.PDSEndpoint) + if err != nil { + slog.Warn("Failed to get service token, skipping", "component", "migration", + "did", stat.DID, "repository", stat.Repository, "error", err) + errorCount++ + continue + } + + // Resolve hold DID to HTTP URL + holdURL := atproto.ResolveHoldURL(holdDID) + if holdURL == "" { + slog.Warn("Failed to resolve hold DID, skipping", "component", "migration", + "hold_did", holdDID) + errorCount++ + continue + } + + // Call hold's setStats endpoint + err = callSetStats(ctx, holdURL, serviceToken, stat.DID, stat.Repository, + stat.PullCount, stat.PushCount, stat.LastPull.String, stat.LastPush.String) + if err != nil { + slog.Warn("Failed to migrate stats to hold, continuing", "component", "migration", + "did", stat.DID, "repository", stat.Repository, "hold", holdDID, "error", err) + errorCount++ + continue + } + + successCount++ + slog.Debug("Migrated stats", "component", "migration", + "did", stat.DID, "repository", stat.Repository, "hold", holdDID, + "pull_count", stat.PullCount, "push_count", stat.PushCount) + } + + slog.Info("Stats migration completed", "component", "migration", + "success", successCount, "skipped", skipCount, "errors", errorCount, "total", len(stats)) + + // Mark migration complete (even if some failed - they'll get updates via Jetstream) + return markMigrationComplete(db) +} + +// markMigrationComplete records that the stats migration has been done +func markMigrationComplete(db *sql.DB) error { + _, err := db.Exec(` + INSERT INTO schema_migrations (version, applied_at) + VALUES (1000, datetime('now')) + ON CONFLICT(version) DO NOTHING + `) + if err != nil { + return fmt.Errorf("failed to mark migration complete: %w", err) + } + return nil +} + +// callSetStats calls the hold's io.atcr.hold.setStats endpoint +func callSetStats(ctx context.Context, holdURL, serviceToken, ownerDID, repository string, pullCount, pushCount int64, lastPull, lastPush string) error { + // Build request + reqBody := map[string]any{ + "ownerDid": ownerDID, + "repository": repository, + "pullCount": pullCount, + "pushCount": pushCount, + } + if lastPull != "" { + reqBody["lastPull"] = lastPull + } + if lastPush != "" { + reqBody["lastPush"] = lastPush + } + + body, err := json.Marshal(reqBody) + if err != nil { + return fmt.Errorf("failed to marshal request: %w", err) + } + + // Create HTTP request + req, err := http.NewRequestWithContext(ctx, "POST", holdURL+atproto.HoldSetStats, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+serviceToken) + + // Send request with timeout + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("setStats failed: status %d, body: %s", resp.StatusCode, body) + } + + return nil +} diff --git a/pkg/appview/jetstream/backfill.go b/pkg/appview/jetstream/backfill.go index cde7d84..80b0574 100644 --- a/pkg/appview/jetstream/backfill.go +++ b/pkg/appview/jetstream/backfill.go @@ -48,8 +48,8 @@ func NewBackfillWorker(database *sql.DB, relayEndpoint, defaultHoldDID string, t return &BackfillWorker{ db: database, - client: client, // This points to the relay - processor: NewProcessor(database, false), // No cache for batch processing + client: client, // This points to the relay + processor: NewProcessor(database, false, nil), // No cache for batch processing, no stats defaultHoldDID: defaultHoldDID, testMode: testMode, refresher: refresher, diff --git a/pkg/appview/jetstream/processor.go b/pkg/appview/jetstream/processor.go index ce06957..ab728c4 100644 --- a/pkg/appview/jetstream/processor.go +++ b/pkg/appview/jetstream/processor.go @@ -16,17 +16,20 @@ import ( // Processor handles shared database operations for both Worker (live) and Backfill (sync) // This eliminates code duplication between the two data ingestion paths type Processor struct { - db *sql.DB - userCache *UserCache // Optional - enabled for Worker, disabled for Backfill - useCache bool + db *sql.DB + userCache *UserCache // Optional - enabled for Worker, disabled for Backfill + statsCache *StatsCache // In-memory cache for per-hold stats aggregation + useCache bool } // NewProcessor creates a new shared processor // useCache: true for Worker (live streaming), false for Backfill (batch processing) -func NewProcessor(database *sql.DB, useCache bool) *Processor { +// statsCache: shared stats cache for aggregating across holds (nil to skip stats processing) +func NewProcessor(database *sql.DB, useCache bool, statsCache *StatsCache) *Processor { p := &Processor{ - db: database, - useCache: useCache, + db: database, + useCache: useCache, + statsCache: statsCache, } if useCache { @@ -369,6 +372,60 @@ func (p *Processor) ProcessIdentity(ctx context.Context, did string, newHandle s return nil } +// ProcessStats handles stats record events from hold PDSes +// This is called when Jetstream receives a stats create/update/delete event from a hold +// The holdDID is the DID of the hold PDS (event.DID), and the record contains ownerDID + repository +func (p *Processor) ProcessStats(ctx context.Context, holdDID string, recordData []byte, isDelete bool) error { + // Skip if no stats cache configured + if p.statsCache == nil { + return nil + } + + // Unmarshal stats record + var statsRecord atproto.StatsRecord + if err := json.Unmarshal(recordData, &statsRecord); err != nil { + return fmt.Errorf("failed to unmarshal stats record: %w", err) + } + + if isDelete { + // Delete from in-memory cache + p.statsCache.Delete(holdDID, statsRecord.OwnerDID, statsRecord.Repository) + } else { + // Parse timestamps + var lastPull, lastPush *time.Time + if statsRecord.LastPull != "" { + t, err := time.Parse(time.RFC3339, statsRecord.LastPull) + if err == nil { + lastPull = &t + } + } + if statsRecord.LastPush != "" { + t, err := time.Parse(time.RFC3339, statsRecord.LastPush) + if err == nil { + lastPush = &t + } + } + + // Update in-memory cache + p.statsCache.Update(holdDID, statsRecord.OwnerDID, statsRecord.Repository, + statsRecord.PullCount, statsRecord.PushCount, lastPull, lastPush) + } + + // Get aggregated stats across all holds + totalPull, totalPush, latestPull, latestPush := p.statsCache.GetAggregated( + statsRecord.OwnerDID, statsRecord.Repository) + + // Upsert aggregated stats to repository_stats + return db.UpsertRepositoryStats(p.db, &db.RepositoryStats{ + DID: statsRecord.OwnerDID, + Repository: statsRecord.Repository, + PullCount: int(totalPull), + PushCount: int(totalPush), + LastPull: latestPull, + LastPush: latestPush, + }) +} + // ProcessAccount handles account status events (deactivation/reactivation) // This is called when Jetstream receives an account event indicating status changes. // diff --git a/pkg/appview/jetstream/processor_test.go b/pkg/appview/jetstream/processor_test.go index c54237d..81ef82a 100644 --- a/pkg/appview/jetstream/processor_test.go +++ b/pkg/appview/jetstream/processor_test.go @@ -115,7 +115,7 @@ func TestNewProcessor(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - p := NewProcessor(database, tt.useCache) + p := NewProcessor(database, tt.useCache, nil) if p == nil { t.Fatal("NewProcessor returned nil") } @@ -139,7 +139,7 @@ func TestProcessManifest_ImageManifest(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() // Create test manifest record @@ -238,7 +238,7 @@ func TestProcessManifest_ManifestList(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() // Create test manifest list record @@ -322,7 +322,7 @@ func TestProcessTag(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() // Create test tag record (using ManifestDigest field for simplicity) @@ -403,7 +403,7 @@ func TestProcessStar(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() // Create test star record @@ -463,7 +463,7 @@ func TestProcessManifest_Duplicate(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() manifestRecord := &atproto.ManifestRecord{ @@ -514,7 +514,7 @@ func TestProcessManifest_EmptyAnnotations(t *testing.T) { database := setupTestDB(t) defer database.Close() - p := NewProcessor(database, false) + p := NewProcessor(database, false, nil) ctx := context.Background() // Manifest with nil annotations @@ -555,7 +555,7 @@ func TestProcessIdentity(t *testing.T) { db := setupTestDB(t) defer db.Close() - processor := NewProcessor(db, false) + processor := NewProcessor(db, false, nil) // Setup: Create test user testDID := "did:plc:alice123" @@ -621,7 +621,7 @@ func TestProcessAccount(t *testing.T) { db := setupTestDB(t) defer db.Close() - processor := NewProcessor(db, false) + processor := NewProcessor(db, false, nil) // Setup: Create test user testDID := "did:plc:bob456" diff --git a/pkg/appview/jetstream/stats_cache.go b/pkg/appview/jetstream/stats_cache.go new file mode 100644 index 0000000..5fdd7ce --- /dev/null +++ b/pkg/appview/jetstream/stats_cache.go @@ -0,0 +1,100 @@ +package jetstream + +import ( + "sync" + "time" +) + +// HoldRepoStats represents stats for a single owner+repo from a specific hold +type HoldRepoStats struct { + OwnerDID string + Repository string + PullCount int64 + PushCount int64 + LastPull *time.Time + LastPush *time.Time +} + +// StatsCache provides in-memory caching of per-hold stats with aggregation +// This allows summing stats across multiple holds for the same owner+repo +type StatsCache struct { + mu sync.RWMutex + // holdDID -> (ownerDID/repo -> stats) + holds map[string]map[string]*HoldRepoStats +} + +// NewStatsCache creates a new in-memory stats cache +func NewStatsCache() *StatsCache { + return &StatsCache{ + holds: make(map[string]map[string]*HoldRepoStats), + } +} + +// makeKey creates a cache key from ownerDID and repository +func makeKey(ownerDID, repo string) string { + return ownerDID + "/" + repo +} + +// Update stores or updates stats for a hold+owner+repo combination +func (c *StatsCache) Update(holdDID, ownerDID, repo string, pullCount, pushCount int64, lastPull, lastPush *time.Time) { + c.mu.Lock() + defer c.mu.Unlock() + + // Ensure hold map exists + if c.holds[holdDID] == nil { + c.holds[holdDID] = make(map[string]*HoldRepoStats) + } + + key := makeKey(ownerDID, repo) + c.holds[holdDID][key] = &HoldRepoStats{ + OwnerDID: ownerDID, + Repository: repo, + PullCount: pullCount, + PushCount: pushCount, + LastPull: lastPull, + LastPush: lastPush, + } +} + +// Delete removes stats for a hold+owner+repo combination +func (c *StatsCache) Delete(holdDID, ownerDID, repo string) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.holds[holdDID] != nil { + key := makeKey(ownerDID, repo) + delete(c.holds[holdDID], key) + } +} + +// GetAggregated returns aggregated stats for an owner+repo by summing across all holds +// Returns (pullCount, pushCount, lastPull, lastPush) +func (c *StatsCache) GetAggregated(ownerDID, repo string) (int64, int64, *time.Time, *time.Time) { + c.mu.RLock() + defer c.mu.RUnlock() + + key := makeKey(ownerDID, repo) + var totalPull, totalPush int64 + var latestPull, latestPush *time.Time + + for _, holdStats := range c.holds { + if stats, ok := holdStats[key]; ok { + totalPull += stats.PullCount + totalPush += stats.PushCount + + // Track latest timestamps + if stats.LastPull != nil { + if latestPull == nil || stats.LastPull.After(*latestPull) { + latestPull = stats.LastPull + } + } + if stats.LastPush != nil { + if latestPush == nil || stats.LastPush.After(*latestPush) { + latestPush = stats.LastPush + } + } + } + } + + return totalPull, totalPush, latestPull, latestPush +} diff --git a/pkg/appview/jetstream/worker.go b/pkg/appview/jetstream/worker.go index 8a11867..1a280d2 100644 --- a/pkg/appview/jetstream/worker.go +++ b/pkg/appview/jetstream/worker.go @@ -34,7 +34,8 @@ type Worker struct { startCursor int64 wantedCollections []string debugCollectionCount int - processor *Processor // Shared processor for DB operations + processor *Processor // Shared processor for DB operations + statsCache *StatsCache // In-memory cache for stats aggregation across holds eventCallback EventCallback connStartTime time.Time // Track when connection started for debugging @@ -56,6 +57,9 @@ func NewWorker(database *sql.DB, jetstreamURL string, startCursor int64) *Worker jetstreamURL = "wss://jetstream2.us-west.bsky.network/subscribe" } + // Create shared stats cache for aggregating across holds + statsCache := NewStatsCache() + return &Worker{ db: database, jetstreamURL: jetstreamURL, @@ -63,7 +67,8 @@ func NewWorker(database *sql.DB, jetstreamURL string, startCursor int64) *Worker wantedCollections: []string{ "io.atcr.*", // Subscribe to all ATCR collections }, - processor: NewProcessor(database, true), // Use cache for live streaming + statsCache: statsCache, + processor: NewProcessor(database, true, statsCache), // Use cache for live streaming } } @@ -313,6 +318,9 @@ func (w *Worker) processMessage(message []byte) error { case atproto.RepoPageCollection: slog.Info("Jetstream processing repo page event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey) return w.processRepoPage(commit) + case atproto.StatsCollection: + slog.Info("Jetstream processing stats event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey) + return w.processStats(commit) default: // Ignore other collections return nil @@ -472,6 +480,35 @@ func (w *Worker) processRepoPage(commit *CommitEvent) error { return w.processor.ProcessRepoPage(context.Background(), commit.DID, commit.RKey, recordBytes, false) } +// processStats processes a stats commit event from a hold PDS +func (w *Worker) processStats(commit *CommitEvent) error { + isDelete := commit.Operation == "delete" + + if isDelete { + // For delete events, we need to parse the rkey to get ownerDID + repository + // The rkey is deterministic: base32(sha256(ownerDID + "/" + repository)[:16]) + // Unfortunately, we can't reverse this - we need the record data + // Delete events don't include record data, so we can't delete from cache + // This is acceptable - stats will be refreshed on next update from hold + slog.Debug("Jetstream ignoring stats delete event (cannot reverse rkey)", "did", commit.DID, "rkey", commit.RKey) + return nil + } + + // Parse stats record + if commit.Record == nil { + return nil + } + + // Marshal map to bytes for processing + recordBytes, err := json.Marshal(commit.Record) + if err != nil { + return fmt.Errorf("failed to marshal record: %w", err) + } + + // Use shared processor - commit.DID is the hold's DID + return w.processor.ProcessStats(context.Background(), commit.DID, recordBytes, false) +} + // processIdentity processes an identity event (handle change) func (w *Worker) processIdentity(event *JetstreamEvent) error { if event.Identity == nil { diff --git a/pkg/appview/middleware/registry.go b/pkg/appview/middleware/registry.go index 96b3482..a9d4e90 100644 --- a/pkg/appview/middleware/registry.go +++ b/pkg/appview/middleware/registry.go @@ -174,7 +174,7 @@ func (vc *validationCache) getOrFetch(ctx context.Context, cacheKey string, fetc // After initialization, request handling uses the NamespaceResolver's instance fields. var ( globalRefresher *oauth.Refresher - globalDatabase storage.DatabaseMetrics + globalDatabase storage.HoldDIDLookup globalAuthorizer auth.HoldAuthorizer ) @@ -186,7 +186,7 @@ func SetGlobalRefresher(refresher *oauth.Refresher) { // SetGlobalDatabase sets the database instance during initialization // Must be called before the registry starts serving requests -func SetGlobalDatabase(database storage.DatabaseMetrics) { +func SetGlobalDatabase(database storage.HoldDIDLookup) { globalDatabase = database } @@ -204,14 +204,14 @@ func init() { // NamespaceResolver wraps a namespace and resolves names type NamespaceResolver struct { distribution.Namespace - defaultHoldDID string // Default hold DID (e.g., "did:web:hold01.atcr.io") - baseURL string // Base URL for error messages (e.g., "https://atcr.io") - testMode bool // If true, fallback to default hold when user's hold is unreachable - refresher *oauth.Refresher // OAuth session manager (copied from global on init) - database storage.DatabaseMetrics // Metrics database (copied from global on init) - authorizer auth.HoldAuthorizer // Hold authorization (copied from global on init) - validationCache *validationCache // Request-level service token cache - readmeFetcher *readme.Fetcher // README fetcher for repo pages + defaultHoldDID string // Default hold DID (e.g., "did:web:hold01.atcr.io") + baseURL string // Base URL for error messages (e.g., "https://atcr.io") + testMode bool // If true, fallback to default hold when user's hold is unreachable + refresher *oauth.Refresher // OAuth session manager (copied from global on init) + database storage.HoldDIDLookup // Database for hold DID lookups (copied from global on init) + authorizer auth.HoldAuthorizer // Hold authorization (copied from global on init) + validationCache *validationCache // Request-level service token cache + readmeFetcher *readme.Fetcher // README fetcher for repo pages } // initATProtoResolver initializes the name resolution middleware diff --git a/pkg/appview/storage/context.go b/pkg/appview/storage/context.go index dda615e..28fe81e 100644 --- a/pkg/appview/storage/context.go +++ b/pkg/appview/storage/context.go @@ -7,10 +7,8 @@ import ( "atcr.io/pkg/auth/oauth" ) -// DatabaseMetrics interface for tracking pull/push counts and querying hold DIDs -type DatabaseMetrics interface { - IncrementPullCount(did, repository string) error - IncrementPushCount(did, repository string) error +// HoldDIDLookup interface for querying hold DIDs from manifests +type HoldDIDLookup interface { GetLatestHoldDIDForRepo(did, repository string) (string, error) } @@ -32,7 +30,7 @@ type RegistryContext struct { PullerPDSEndpoint string // Puller's PDS endpoint URL // Shared services (same for all requests) - Database DatabaseMetrics // Metrics tracking database + Database HoldDIDLookup // Database for hold DID lookups Authorizer auth.HoldAuthorizer // Hold access authorization Refresher *oauth.Refresher // OAuth session manager ReadmeFetcher *readme.Fetcher // README fetcher for repo pages diff --git a/pkg/appview/storage/context_test.go b/pkg/appview/storage/context_test.go index 2ff123e..c26b0f1 100644 --- a/pkg/appview/storage/context_test.go +++ b/pkg/appview/storage/context_test.go @@ -1,48 +1,20 @@ package storage import ( - "sync" "testing" "atcr.io/pkg/atproto" ) // Mock implementations for testing -type mockDatabaseMetrics struct { - mu sync.Mutex - pullCount int - pushCount int + +// mockHoldDIDLookup implements HoldDIDLookup for testing +type mockHoldDIDLookup struct { + holdDID string // Return value for GetLatestHoldDIDForRepo } -func (m *mockDatabaseMetrics) IncrementPullCount(did, repository string) error { - m.mu.Lock() - defer m.mu.Unlock() - m.pullCount++ - return nil -} - -func (m *mockDatabaseMetrics) IncrementPushCount(did, repository string) error { - m.mu.Lock() - defer m.mu.Unlock() - m.pushCount++ - return nil -} - -func (m *mockDatabaseMetrics) GetLatestHoldDIDForRepo(did, repository string) (string, error) { - // Return empty string for mock - tests can override if needed - return "", nil -} - -func (m *mockDatabaseMetrics) getPullCount() int { - m.mu.Lock() - defer m.mu.Unlock() - return m.pullCount -} - -func (m *mockDatabaseMetrics) getPushCount() int { - m.mu.Lock() - defer m.mu.Unlock() - return m.pushCount +func (m *mockHoldDIDLookup) GetLatestHoldDIDForRepo(did, repository string) (string, error) { + return m.holdDID, nil } type mockHoldAuthorizer struct{} @@ -63,7 +35,7 @@ func TestRegistryContext_Fields(t *testing.T) { ATProtoClient: &atproto.Client{ // Mock client - would need proper initialization in real tests }, - Database: &mockDatabaseMetrics{}, + Database: &mockHoldDIDLookup{holdDID: "did:web:hold01.atcr.io"}, } // Verify fields are accessible @@ -88,20 +60,18 @@ func TestRegistryContext_Fields(t *testing.T) { } func TestRegistryContext_DatabaseInterface(t *testing.T) { - db := &mockDatabaseMetrics{} + db := &mockHoldDIDLookup{holdDID: "did:web:test-hold.example.com"} ctx := &RegistryContext{ Database: db, } - // Test that interface methods are callable - err := ctx.Database.IncrementPullCount("did:plc:test", "repo") + // Test that interface method is callable + holdDID, err := ctx.Database.GetLatestHoldDIDForRepo("did:plc:test", "repo") if err != nil { t.Errorf("Unexpected error: %v", err) } - - err = ctx.Database.IncrementPushCount("did:plc:test", "repo") - if err != nil { - t.Errorf("Unexpected error: %v", err) + if holdDID != "did:web:test-hold.example.com" { + t.Errorf("Expected holdDID %q, got %q", "did:web:test-hold.example.com", holdDID) } } diff --git a/pkg/appview/storage/manifest_store.go b/pkg/appview/storage/manifest_store.go index 76b6392..f832cae 100644 --- a/pkg/appview/storage/manifest_store.go +++ b/pkg/appview/storage/manifest_store.go @@ -73,14 +73,20 @@ func (s *ManifestStore) Get(ctx context.Context, dgst digest.Digest, options ... } } - // Track pull count (increment asynchronously to avoid blocking the response) + // Notify hold about manifest pull (for stats tracking) // Only count GET requests (actual downloads), not HEAD requests (existence checks) - if s.ctx.Database != nil { - // Check HTTP method from context (distribution library stores it as "http.request.method") - if method, ok := ctx.Value("http.request.method").(string); ok && method == "GET" { + // Check HTTP method from context (distribution library stores it as "http.request.method") + if method, ok := ctx.Value("http.request.method").(string); ok && method == "GET" { + // Do this asynchronously to avoid blocking the response + if s.ctx.ServiceToken != "" && s.ctx.Handle != "" { go func() { - if err := s.ctx.Database.IncrementPullCount(s.ctx.DID, s.ctx.Repository); err != nil { - slog.Warn("Failed to increment pull count", "did", s.ctx.DID, "repository", s.ctx.Repository, "error", err) + defer func() { + if r := recover(); r != nil { + slog.Error("Panic in notifyHoldAboutManifest (pull)", "panic", r) + } + }() + if err := s.notifyHoldAboutManifest(context.Background(), nil, "", "", "pull"); err != nil { + slog.Warn("Failed to notify hold about manifest pull", "error", err) } }() } @@ -190,15 +196,6 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, return "", fmt.Errorf("failed to store manifest record in ATProto: %w", err) } - // Track push count (increment asynchronously to avoid blocking the response) - if s.ctx.Database != nil { - go func() { - if err := s.ctx.Database.IncrementPushCount(s.ctx.DID, s.ctx.Repository); err != nil { - slog.Warn("Failed to increment push count", "did", s.ctx.DID, "repository", s.ctx.Repository, "error", err) - } - }() - } - // Also handle tag if specified var tag string for _, option := range options { @@ -213,7 +210,7 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, } } - // Notify hold about manifest upload (for layer tracking and Bluesky posts) + // Notify hold about manifest push (for layer tracking, Bluesky posts, and stats) // Do this asynchronously to avoid blocking the push if tag != "" && s.ctx.ServiceToken != "" && s.ctx.Handle != "" { go func() { @@ -222,8 +219,8 @@ func (s *ManifestStore) Put(ctx context.Context, manifest distribution.Manifest, slog.Error("Panic in notifyHoldAboutManifest", "panic", r) } }() - if err := s.notifyHoldAboutManifest(context.Background(), manifestRecord, tag, dgst.String()); err != nil { - slog.Warn("Failed to notify hold about manifest", "error", err) + if err := s.notifyHoldAboutManifest(context.Background(), manifestRecord, tag, dgst.String(), "push"); err != nil { + slog.Warn("Failed to notify hold about manifest push", "error", err) } }() } @@ -298,9 +295,11 @@ func (s *ManifestStore) extractConfigLabels(ctx context.Context, configDigestStr return configJSON.Config.Labels, nil } -// notifyHoldAboutManifest notifies the hold service about a manifest upload -// This enables the hold to create layer records and Bluesky posts -func (s *ManifestStore) notifyHoldAboutManifest(ctx context.Context, manifestRecord *atproto.ManifestRecord, tag, manifestDigest string) error { +// notifyHoldAboutManifest notifies the hold service about a manifest operation +// For push: Creates layer records and optionally posts to Bluesky +// For pull: Just increments stats (no layer records or posts) +// operation should be "push" or "pull" +func (s *ManifestStore) notifyHoldAboutManifest(ctx context.Context, manifestRecord *atproto.ManifestRecord, tag, manifestDigest, operation string) error { // Skip if no service token configured (e.g., anonymous pulls) if s.ctx.ServiceToken == "" { return nil @@ -314,57 +313,63 @@ func (s *ManifestStore) notifyHoldAboutManifest(ctx context.Context, manifestRec serviceToken := s.ctx.ServiceToken // Build notification request - manifestData := map[string]any{ - "mediaType": manifestRecord.MediaType, - } - - // Add config if present (not present in manifest lists/indexes) - if manifestRecord.Config != nil { - manifestData["config"] = map[string]any{ - "digest": manifestRecord.Config.Digest, - "size": manifestRecord.Config.Size, - } - } - - // Add layers if present - if len(manifestRecord.Layers) > 0 { - layers := make([]map[string]any, len(manifestRecord.Layers)) - for i, layer := range manifestRecord.Layers { - layers[i] = map[string]any{ - "digest": layer.Digest, - "size": layer.Size, - "mediaType": layer.MediaType, - } - } - manifestData["layers"] = layers - } - - // Add manifests if present (for multi-arch images / manifest lists) - if len(manifestRecord.Manifests) > 0 { - manifests := make([]map[string]any, len(manifestRecord.Manifests)) - for i, m := range manifestRecord.Manifests { - mData := map[string]any{ - "digest": m.Digest, - "size": m.Size, - "mediaType": m.MediaType, - } - if m.Platform != nil { - mData["platform"] = map[string]any{ - "os": m.Platform.OS, - "architecture": m.Platform.Architecture, - } - } - manifests[i] = mData - } - manifestData["manifests"] = manifests - } - notifyReq := map[string]any{ "repository": s.ctx.Repository, - "tag": tag, "userDid": s.ctx.DID, "userHandle": s.ctx.Handle, - "manifest": manifestData, + "operation": operation, + } + + // For push operations, include full manifest data + if operation == "push" && manifestRecord != nil { + notifyReq["tag"] = tag + + manifestData := map[string]any{ + "mediaType": manifestRecord.MediaType, + } + + // Add config if present (not present in manifest lists/indexes) + if manifestRecord.Config != nil { + manifestData["config"] = map[string]any{ + "digest": manifestRecord.Config.Digest, + "size": manifestRecord.Config.Size, + } + } + + // Add layers if present + if len(manifestRecord.Layers) > 0 { + layers := make([]map[string]any, len(manifestRecord.Layers)) + for i, layer := range manifestRecord.Layers { + layers[i] = map[string]any{ + "digest": layer.Digest, + "size": layer.Size, + "mediaType": layer.MediaType, + } + } + manifestData["layers"] = layers + } + + // Add manifests if present (for multi-arch images / manifest lists) + if len(manifestRecord.Manifests) > 0 { + manifests := make([]map[string]any, len(manifestRecord.Manifests)) + for i, m := range manifestRecord.Manifests { + mData := map[string]any{ + "digest": m.Digest, + "size": m.Size, + "mediaType": m.MediaType, + } + if m.Platform != nil { + mData["platform"] = map[string]any{ + "os": m.Platform.OS, + "architecture": m.Platform.Architecture, + } + } + manifests[i] = mData + } + manifestData["manifests"] = manifests + } + + notifyReq["manifest"] = manifestData } // Marshal request @@ -401,7 +406,7 @@ func (s *ManifestStore) notifyHoldAboutManifest(ctx context.Context, manifestRec // Parse response (optional logging) var notifyResp map[string]any if err := json.NewDecoder(resp.Body).Decode(¬ifyResp); err == nil { - slog.Info("Hold notification successful", "repository", s.ctx.Repository, "tag", tag, "response", notifyResp) + slog.Debug("Hold notification successful", "repository", s.ctx.Repository, "operation", operation, "response", notifyResp) } return nil diff --git a/pkg/appview/storage/manifest_store_test.go b/pkg/appview/storage/manifest_store_test.go index 73291f6..1b2afb3 100644 --- a/pkg/appview/storage/manifest_store_test.go +++ b/pkg/appview/storage/manifest_store_test.go @@ -8,14 +8,13 @@ import ( "net/http" "net/http/httptest" "testing" - "time" "atcr.io/pkg/atproto" "github.com/distribution/distribution/v3" "github.com/opencontainers/go-digest" ) -// mockDatabaseMetrics removed - using the one from context_test.go +// mockHoldDIDLookup defined in context_test.go // mockBlobStore is a minimal mock of distribution.BlobStore for testing type mockBlobStore struct { @@ -73,7 +72,7 @@ func (m *mockBlobStore) Open(ctx context.Context, dgst digest.Digest) (io.ReadSe } // mockRegistryContext creates a mock RegistryContext for testing -func mockRegistryContext(client *atproto.Client, repository, holdDID, did, handle string, database DatabaseMetrics) *RegistryContext { +func mockRegistryContext(client *atproto.Client, repository, holdDID, did, handle string, database HoldDIDLookup) *RegistryContext { return &RegistryContext{ ATProtoClient: client, Repository: repository, @@ -117,7 +116,7 @@ func TestDigestToRKey(t *testing.T) { func TestNewManifestStore(t *testing.T) { client := atproto.NewClient("https://pds.example.com", "did:plc:test123", "token") blobStore := newMockBlobStore() - db := &mockDatabaseMetrics{} + db := &mockHoldDIDLookup{} ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:alice123", "alice.test", db) store := NewManifestStore(ctx, blobStore) @@ -274,9 +273,9 @@ func TestExtractConfigLabels_InvalidJSON(t *testing.T) { } } -// TestManifestStore_WithMetrics tests that metrics are tracked -func TestManifestStore_WithMetrics(t *testing.T) { - db := &mockDatabaseMetrics{} +// TestManifestStore_WithDatabase tests that database is wired up +func TestManifestStore_WithDatabase(t *testing.T) { + db := &mockHoldDIDLookup{holdDID: "did:web:test-hold.example.com"} client := atproto.NewClient("https://pds.example.com", "did:plc:test123", "token") ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:alice123", "alice.test", db) store := NewManifestStore(ctx, nil) @@ -285,12 +284,11 @@ func TestManifestStore_WithMetrics(t *testing.T) { t.Error("ManifestStore should store database reference") } - // Note: Actual metrics tracking happens in Put() and Get() which require - // full mock setup. The important thing is that the database is wired up. + // Database is used for hold DID lookups during blob routing } -// TestManifestStore_WithoutMetrics tests that nil database is acceptable -func TestManifestStore_WithoutMetrics(t *testing.T) { +// TestManifestStore_WithoutDatabase tests that nil database is acceptable +func TestManifestStore_WithoutDatabase(t *testing.T) { client := atproto.NewClient("https://pds.example.com", "did:plc:test123", "token") ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:alice123", "alice.test", nil) store := NewManifestStore(ctx, nil) @@ -464,7 +462,7 @@ func TestManifestStore_Get(t *testing.T) { defer server.Close() client := atproto.NewClient(server.URL, "did:plc:test123", "token") - db := &mockDatabaseMetrics{} + db := &mockHoldDIDLookup{} ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:test123", "test.handle", db) store := NewManifestStore(ctx, nil) @@ -487,82 +485,6 @@ func TestManifestStore_Get(t *testing.T) { } } -// TestManifestStore_Get_OnlyCountsGETRequests verifies that HEAD requests don't increment pull count -func TestManifestStore_Get_OnlyCountsGETRequests(t *testing.T) { - ociManifest := []byte(`{"schemaVersion":2}`) - - tests := []struct { - name string - httpMethod string - expectPullIncrement bool - }{ - { - name: "GET request increments pull count", - httpMethod: "GET", - expectPullIncrement: true, - }, - { - name: "HEAD request does not increment pull count", - httpMethod: "HEAD", - expectPullIncrement: false, - }, - { - name: "POST request does not increment pull count", - httpMethod: "POST", - expectPullIncrement: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == atproto.SyncGetBlob { - w.Write(ociManifest) - return - } - w.Write([]byte(`{ - "uri": "at://did:plc:test123/io.atcr.manifest/abc123", - "value": { - "$type":"io.atcr.manifest", - "holdDid":"did:web:hold01.atcr.io", - "mediaType":"application/vnd.oci.image.manifest.v1+json", - "manifestBlob":{"ref":{"$link":"bafytest"},"size":100} - } - }`)) - })) - defer server.Close() - - client := atproto.NewClient(server.URL, "did:plc:test123", "token") - mockDB := &mockDatabaseMetrics{} - ctx := mockRegistryContext(client, "myapp", "did:web:hold01.atcr.io", "did:plc:test123", "test.handle", mockDB) - store := NewManifestStore(ctx, nil) - - // Create a context with the HTTP method stored (as distribution library does) - testCtx := context.WithValue(context.Background(), "http.request.method", tt.httpMethod) - - _, err := store.Get(testCtx, "sha256:abc123") - if err != nil { - t.Fatalf("Get() error = %v", err) - } - - // Wait for async goroutine to complete (metrics are incremented asynchronously) - time.Sleep(50 * time.Millisecond) - - if tt.expectPullIncrement { - // Check that IncrementPullCount was called - if mockDB.getPullCount() == 0 { - t.Error("Expected pull count to be incremented for GET request, but it wasn't") - } - } else { - // Check that IncrementPullCount was NOT called - if mockDB.getPullCount() > 0 { - t.Errorf("Expected pull count NOT to be incremented for %s request, but it was (count=%d)", tt.httpMethod, mockDB.getPullCount()) - } - } - }) - } -} - // TestManifestStore_Put tests storing manifests func TestManifestStore_Put(t *testing.T) { ociManifest := []byte(`{ @@ -655,7 +577,7 @@ func TestManifestStore_Put(t *testing.T) { defer server.Close() client := atproto.NewClient(server.URL, "did:plc:test123", "token") - db := &mockDatabaseMetrics{} + db := &mockHoldDIDLookup{} ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:test123", "test.handle", db) store := NewManifestStore(ctx, nil) @@ -939,7 +861,7 @@ func TestManifestStore_Put_ManifestListValidation(t *testing.T) { defer server.Close() client := atproto.NewClient(server.URL, "did:plc:test123", "token") - db := &mockDatabaseMetrics{} + db := &mockHoldDIDLookup{} ctx := mockRegistryContext(client, "myapp", "did:web:hold.example.com", "did:plc:test123", "test.handle", db) store := NewManifestStore(ctx, nil) diff --git a/pkg/atproto/cbor_gen.go b/pkg/atproto/cbor_gen.go index 80fcd9e..e2aa5b3 100644 --- a/pkg/atproto/cbor_gen.go +++ b/pkg/atproto/cbor_gen.go @@ -1419,3 +1419,384 @@ func (t *TangledProfileRecord) UnmarshalCBOR(r io.Reader) (err error) { return nil } +func (t *StatsRecord) MarshalCBOR(w io.Writer) error { + if t == nil { + _, err := w.Write(cbg.CborNull) + return err + } + + cw := cbg.NewCborWriter(w) + fieldCount := 8 + + if t.LastPull == "" { + fieldCount-- + } + + if t.LastPush == "" { + fieldCount-- + } + + if _, err := cw.Write(cbg.CborEncodeMajorType(cbg.MajMap, uint64(fieldCount))); err != nil { + return err + } + + // t.Type (string) (string) + if len("$type") > 8192 { + return xerrors.Errorf("Value in field \"$type\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("$type"))); err != nil { + return err + } + if _, err := cw.WriteString(string("$type")); err != nil { + return err + } + + if len(t.Type) > 8192 { + return xerrors.Errorf("Value in field t.Type was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.Type))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.Type)); err != nil { + return err + } + + // t.LastPull (string) (string) + if t.LastPull != "" { + + if len("lastPull") > 8192 { + return xerrors.Errorf("Value in field \"lastPull\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("lastPull"))); err != nil { + return err + } + if _, err := cw.WriteString(string("lastPull")); err != nil { + return err + } + + if len(t.LastPull) > 8192 { + return xerrors.Errorf("Value in field t.LastPull was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.LastPull))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.LastPull)); err != nil { + return err + } + } + + // t.LastPush (string) (string) + if t.LastPush != "" { + + if len("lastPush") > 8192 { + return xerrors.Errorf("Value in field \"lastPush\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("lastPush"))); err != nil { + return err + } + if _, err := cw.WriteString(string("lastPush")); err != nil { + return err + } + + if len(t.LastPush) > 8192 { + return xerrors.Errorf("Value in field t.LastPush was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.LastPush))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.LastPush)); err != nil { + return err + } + } + + // t.OwnerDID (string) (string) + if len("ownerDid") > 8192 { + return xerrors.Errorf("Value in field \"ownerDid\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("ownerDid"))); err != nil { + return err + } + if _, err := cw.WriteString(string("ownerDid")); err != nil { + return err + } + + if len(t.OwnerDID) > 8192 { + return xerrors.Errorf("Value in field t.OwnerDID was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.OwnerDID))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.OwnerDID)); err != nil { + return err + } + + // t.PullCount (int64) (int64) + if len("pullCount") > 8192 { + return xerrors.Errorf("Value in field \"pullCount\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("pullCount"))); err != nil { + return err + } + if _, err := cw.WriteString(string("pullCount")); err != nil { + return err + } + + if t.PullCount >= 0 { + if err := cw.WriteMajorTypeHeader(cbg.MajUnsignedInt, uint64(t.PullCount)); err != nil { + return err + } + } else { + if err := cw.WriteMajorTypeHeader(cbg.MajNegativeInt, uint64(-t.PullCount-1)); err != nil { + return err + } + } + + // t.PushCount (int64) (int64) + if len("pushCount") > 8192 { + return xerrors.Errorf("Value in field \"pushCount\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("pushCount"))); err != nil { + return err + } + if _, err := cw.WriteString(string("pushCount")); err != nil { + return err + } + + if t.PushCount >= 0 { + if err := cw.WriteMajorTypeHeader(cbg.MajUnsignedInt, uint64(t.PushCount)); err != nil { + return err + } + } else { + if err := cw.WriteMajorTypeHeader(cbg.MajNegativeInt, uint64(-t.PushCount-1)); err != nil { + return err + } + } + + // t.UpdatedAt (string) (string) + if len("updatedAt") > 8192 { + return xerrors.Errorf("Value in field \"updatedAt\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("updatedAt"))); err != nil { + return err + } + if _, err := cw.WriteString(string("updatedAt")); err != nil { + return err + } + + if len(t.UpdatedAt) > 8192 { + return xerrors.Errorf("Value in field t.UpdatedAt was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.UpdatedAt))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.UpdatedAt)); err != nil { + return err + } + + // t.Repository (string) (string) + if len("repository") > 8192 { + return xerrors.Errorf("Value in field \"repository\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("repository"))); err != nil { + return err + } + if _, err := cw.WriteString(string("repository")); err != nil { + return err + } + + if len(t.Repository) > 8192 { + return xerrors.Errorf("Value in field t.Repository was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.Repository))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.Repository)); err != nil { + return err + } + return nil +} + +func (t *StatsRecord) UnmarshalCBOR(r io.Reader) (err error) { + *t = StatsRecord{} + + cr := cbg.NewCborReader(r) + + maj, extra, err := cr.ReadHeader() + if err != nil { + return err + } + defer func() { + if err == io.EOF { + err = io.ErrUnexpectedEOF + } + }() + + if maj != cbg.MajMap { + return fmt.Errorf("cbor input should be of type map") + } + + if extra > cbg.MaxLength { + return fmt.Errorf("StatsRecord: map struct too large (%d)", extra) + } + + n := extra + + nameBuf := make([]byte, 10) + for i := uint64(0); i < n; i++ { + nameLen, ok, err := cbg.ReadFullStringIntoBuf(cr, nameBuf, 8192) + if err != nil { + return err + } + + if !ok { + // Field doesn't exist on this type, so ignore it + if err := cbg.ScanForLinks(cr, func(cid.Cid) {}); err != nil { + return err + } + continue + } + + switch string(nameBuf[:nameLen]) { + // t.Type (string) (string) + case "$type": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.Type = string(sval) + } + // t.LastPull (string) (string) + case "lastPull": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.LastPull = string(sval) + } + // t.LastPush (string) (string) + case "lastPush": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.LastPush = string(sval) + } + // t.OwnerDID (string) (string) + case "ownerDid": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.OwnerDID = string(sval) + } + // t.PullCount (int64) (int64) + case "pullCount": + { + maj, extra, err := cr.ReadHeader() + if err != nil { + return err + } + var extraI int64 + switch maj { + case cbg.MajUnsignedInt: + extraI = int64(extra) + if extraI < 0 { + return fmt.Errorf("int64 positive overflow") + } + case cbg.MajNegativeInt: + extraI = int64(extra) + if extraI < 0 { + return fmt.Errorf("int64 negative overflow") + } + extraI = -1 - extraI + default: + return fmt.Errorf("wrong type for int64 field: %d", maj) + } + + t.PullCount = int64(extraI) + } + // t.PushCount (int64) (int64) + case "pushCount": + { + maj, extra, err := cr.ReadHeader() + if err != nil { + return err + } + var extraI int64 + switch maj { + case cbg.MajUnsignedInt: + extraI = int64(extra) + if extraI < 0 { + return fmt.Errorf("int64 positive overflow") + } + case cbg.MajNegativeInt: + extraI = int64(extra) + if extraI < 0 { + return fmt.Errorf("int64 negative overflow") + } + extraI = -1 - extraI + default: + return fmt.Errorf("wrong type for int64 field: %d", maj) + } + + t.PushCount = int64(extraI) + } + // t.UpdatedAt (string) (string) + case "updatedAt": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.UpdatedAt = string(sval) + } + // t.Repository (string) (string) + case "repository": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.Repository = string(sval) + } + + default: + // Field doesn't exist on this type, so ignore it + if err := cbg.ScanForLinks(r, func(cid.Cid) {}); err != nil { + return err + } + } + } + + return nil +} diff --git a/pkg/atproto/endpoints.go b/pkg/atproto/endpoints.go index 3d78978..9148153 100644 --- a/pkg/atproto/endpoints.go +++ b/pkg/atproto/endpoints.go @@ -45,6 +45,12 @@ const ( // Request: {"repository": "...", "tag": "...", "userDid": "...", "userHandle": "...", "manifest": {...}} // Response: {"success": true, "layersCreated": 5, "postCreated": true, "postUri": "at://..."} HoldNotifyManifest = "/xrpc/io.atcr.hold.notifyManifest" + + // HoldSetStats sets absolute stats values for a repository (used by migration). + // Method: POST + // Request: {"ownerDid": "...", "repository": "...", "pullCount": 10, "pushCount": 5, "lastPull": "...", "lastPush": "..."} + // Response: {"success": true} + HoldSetStats = "/xrpc/io.atcr.hold.setStats" ) // Hold service crew management endpoints (io.atcr.hold.*) diff --git a/pkg/atproto/generate.go b/pkg/atproto/generate.go index c29fd79..626a1f1 100644 --- a/pkg/atproto/generate.go +++ b/pkg/atproto/generate.go @@ -25,12 +25,13 @@ import ( ) func main() { - // Generate map-style encoders for CrewRecord, CaptainRecord, LayerRecord, and TangledProfileRecord + // Generate map-style encoders if err := cbg.WriteMapEncodersToFile("cbor_gen.go", "atproto", atproto.CrewRecord{}, atproto.CaptainRecord{}, atproto.LayerRecord{}, atproto.TangledProfileRecord{}, + atproto.StatsRecord{}, ); err != nil { fmt.Printf("Failed to generate CBOR encoders: %v\n", err) os.Exit(1) diff --git a/pkg/atproto/lexicon.go b/pkg/atproto/lexicon.go index 5cfa042..1fc58c8 100644 --- a/pkg/atproto/lexicon.go +++ b/pkg/atproto/lexicon.go @@ -3,6 +3,8 @@ package atproto //go:generate go run generate.go import ( + "crypto/sha256" + "encoding/base32" "encoding/base64" "encoding/json" "fmt" @@ -35,6 +37,10 @@ const ( // Stored in hold's embedded PDS to track which layers are stored LayerCollection = "io.atcr.hold.layer" + // StatsCollection is the collection name for repository statistics + // Stored in hold's embedded PDS to track pull/push counts per owner+repo + StatsCollection = "io.atcr.hold.stats" + // TangledProfileCollection is the collection name for tangled profiles // Stored in hold's embedded PDS (singleton record at rkey "self") TangledProfileCollection = "sh.tangled.actor.profile" @@ -620,6 +626,45 @@ func NewLayerRecord(digest string, size int64, mediaType, repository, userDID, u } } +// StatsRecord represents repository statistics stored in the hold's PDS +// Collection: io.atcr.hold.stats +// Stored in the hold's embedded PDS for tracking manifest pull/push counts +// Uses CBOR encoding for efficient storage in hold's carstore +// RKey is deterministic: base32(sha256(ownerDID + "/" + repository)[:16]) +type StatsRecord struct { + Type string `json:"$type" cborgen:"$type"` + OwnerDID string `json:"ownerDid" cborgen:"ownerDid"` // DID of the image owner (e.g., "did:plc:xyz123") + Repository string `json:"repository" cborgen:"repository"` // Repository name (e.g., "myapp") + PullCount int64 `json:"pullCount" cborgen:"pullCount"` // Number of manifest downloads + LastPull string `json:"lastPull,omitempty" cborgen:"lastPull,omitempty"` + PushCount int64 `json:"pushCount" cborgen:"pushCount"` // Number of manifest uploads + LastPush string `json:"lastPush,omitempty" cborgen:"lastPush,omitempty"` + UpdatedAt string `json:"updatedAt" cborgen:"updatedAt"` // RFC3339 timestamp +} + +// NewStatsRecord creates a new stats record +func NewStatsRecord(ownerDID, repository string) *StatsRecord { + return &StatsRecord{ + Type: StatsCollection, + OwnerDID: ownerDID, + Repository: repository, + PullCount: 0, + PushCount: 0, + UpdatedAt: time.Now().Format(time.RFC3339), + } +} + +// StatsRecordKey generates a deterministic record key for stats +// Uses base32 encoding of first 16 bytes of SHA-256 hash of "ownerDID/repository" +// This ensures same owner+repo always maps to same rkey +func StatsRecordKey(ownerDID, repository string) string { + combined := ownerDID + "/" + repository + hash := sha256.Sum256([]byte(combined)) + // Use first 16 bytes (128 bits) for collision resistance + // Encode with base32 (alphanumeric, lowercase, no padding) for ATProto rkey compatibility + return strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(hash[:16])) +} + // TangledProfileRecord represents a Tangled profile for the hold // Collection: sh.tangled.actor.profile (singleton record at rkey "self") // Stored in the hold's embedded PDS diff --git a/pkg/hold/oci/xrpc.go b/pkg/hold/oci/xrpc.go index 8a4b68a..8ca5d52 100644 --- a/pkg/hold/oci/xrpc.go +++ b/pkg/hold/oci/xrpc.go @@ -50,6 +50,7 @@ func (h *XRPCHandler) RegisterHandlers(r chi.Router) { r.Post(atproto.HoldCompleteUpload, h.HandleCompleteUpload) r.Post(atproto.HoldAbortUpload, h.HandleAbortUpload) r.Post(atproto.HoldNotifyManifest, h.HandleNotifyManifest) + r.Post(atproto.HoldSetStats, h.HandleSetStats) }) } @@ -201,8 +202,10 @@ func (h *XRPCHandler) HandleAbortUpload(w http.ResponseWriter, r *http.Request) }) } -// HandleNotifyManifest handles manifest upload notifications from AppView -// Creates layer records and optionally posts to Bluesky +// HandleNotifyManifest handles manifest notifications from AppView +// For pushes: Creates layer records and optionally posts to Bluesky +// For pulls: Just increments stats (no layer records or posts) +// Always increments stats (pull or push counts) func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Request) { ctx := r.Context() @@ -219,6 +222,7 @@ func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Reques Tag string `json:"tag"` UserDID string `json:"userDid"` UserHandle string `json:"userHandle"` + Operation string `json:"operation"` // "push" or "pull", defaults to "push" for backward compatibility Manifest struct { MediaType string `json:"mediaType"` Config struct { @@ -253,102 +257,182 @@ func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Reques return } - // Check if manifest posts are enabled - // Read from captain record (which is synced with HOLD_BLUESKY_POSTS_ENABLED env var) - postsEnabled := false - _, captain, err := h.pds.GetCaptainRecord(ctx) - if err == nil { - postsEnabled = captain.EnableBlueskyPosts - } else { - // Fallback to env var if captain record doesn't exist (shouldn't happen in normal operation) - postsEnabled = h.enableBlueskyPosts + // Default operation to "push" for backward compatibility + operation := req.Operation + if operation == "" { + operation = "push" } - // Create layer records for each blob - layersCreated := 0 - for _, layer := range req.Manifest.Layers { - record := atproto.NewLayerRecord( - layer.Digest, - layer.Size, - layer.MediaType, - req.Repository, - req.UserDID, - req.UserHandle, - ) + // Validate operation + if operation != "push" && operation != "pull" { + RespondError(w, http.StatusBadRequest, fmt.Sprintf("invalid operation: %s (must be 'push' or 'pull')", operation)) + return + } - _, _, err := h.pds.CreateLayerRecord(ctx, record) - if err != nil { - slog.Error("Failed to create layer record", "error", err) - // Continue creating other records + var layersCreated int + var postCreated bool + var postURI string + + // Only create layer records and Bluesky posts for pushes + if operation == "push" { + // Check if manifest posts are enabled + // Read from captain record (which is synced with HOLD_BLUESKY_POSTS_ENABLED env var) + postsEnabled := false + _, captain, err := h.pds.GetCaptainRecord(ctx) + if err == nil { + postsEnabled = captain.EnableBlueskyPosts } else { - layersCreated++ + // Fallback to env var if captain record doesn't exist (shouldn't happen in normal operation) + postsEnabled = h.enableBlueskyPosts } - } - // Check if this is a multi-arch image (has manifests instead of layers) - isMultiArch := len(req.Manifest.Manifests) > 0 + // Create layer records for each blob + for _, layer := range req.Manifest.Layers { + record := atproto.NewLayerRecord( + layer.Digest, + layer.Size, + layer.MediaType, + req.Repository, + req.UserDID, + req.UserHandle, + ) - // Calculate total size from all layers (for single-arch images) - var totalSize int64 - for _, layer := range req.Manifest.Layers { - totalSize += layer.Size - } - totalSize += req.Manifest.Config.Size // Add config blob size + _, _, err := h.pds.CreateLayerRecord(ctx, record) + if err != nil { + slog.Error("Failed to create layer record", "error", err) + // Continue creating other records + } else { + layersCreated++ + } + } - // Extract platforms for multi-arch images - var platforms []string - if isMultiArch { - for _, m := range req.Manifest.Manifests { - if m.Platform != nil { - platforms = append(platforms, m.Platform.OS+"/"+m.Platform.Architecture) + // Check if this is a multi-arch image (has manifests instead of layers) + isMultiArch := len(req.Manifest.Manifests) > 0 + + // Calculate total size from all layers (for single-arch images) + var totalSize int64 + for _, layer := range req.Manifest.Layers { + totalSize += layer.Size + } + totalSize += req.Manifest.Config.Size // Add config blob size + + // Extract platforms for multi-arch images + var platforms []string + if isMultiArch { + for _, m := range req.Manifest.Manifests { + if m.Platform != nil { + platforms = append(platforms, m.Platform.OS+"/"+m.Platform.Architecture) + } + } + } + + // Create Bluesky post if enabled + if postsEnabled { + // Extract manifest digest from first layer (or use config digest as fallback) + manifestDigest := req.Manifest.Config.Digest + if len(req.Manifest.Layers) > 0 { + manifestDigest = req.Manifest.Layers[0].Digest + } + + postURI, err = h.pds.CreateManifestPost( + ctx, + h.driver, + req.Repository, + req.Tag, + req.UserHandle, + req.UserDID, + manifestDigest, + totalSize, + platforms, + ) + if err != nil { + slog.Error("Failed to create manifest post", "error", err) + } else { + postCreated = true } } } - // Create Bluesky post if enabled - var postURI string - postCreated := false - if postsEnabled { - // Extract manifest digest from first layer (or use config digest as fallback) - manifestDigest := req.Manifest.Config.Digest - if len(req.Manifest.Layers) > 0 { - manifestDigest = req.Manifest.Layers[0].Digest - } - - postURI, err = h.pds.CreateManifestPost( - ctx, - h.driver, - req.Repository, - req.Tag, - req.UserHandle, - req.UserDID, - manifestDigest, - totalSize, - platforms, - ) - if err != nil { - slog.Error("Failed to create manifest post", "error", err) - } else { - postCreated = true - } + // ALWAYS increment stats (even if Bluesky posts disabled, even for pulls) + statsUpdated := false + if err := h.pds.IncrementStats(ctx, req.UserDID, req.Repository, operation); err != nil { + slog.Error("Failed to increment stats", "operation", operation, "error", err) + } else { + statsUpdated = true } // Return response resp := map[string]any{ - "success": layersCreated > 0 || postCreated, - "layersCreated": layersCreated, - "postCreated": postCreated, + "success": statsUpdated || layersCreated > 0 || postCreated, + "operation": operation, + "statsUpdated": statsUpdated, } - if postURI != "" { - resp["postUri"] = postURI - } - if err != nil && layersCreated == 0 && !postCreated { - resp["error"] = err.Error() + + // Only include push-specific fields for push operations + if operation == "push" { + resp["layersCreated"] = layersCreated + resp["postCreated"] = postCreated + if postURI != "" { + resp["postUri"] = postURI + } } RespondJSON(w, http.StatusOK, resp) } +// HandleSetStats sets absolute stats values for a repository (used by migration) +// This is a migration-only endpoint that allows AppView to sync existing stats to holds +func (h *XRPCHandler) HandleSetStats(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + // Validate service token (same auth as blob:write endpoints) + validatedUser, err := pds.ValidateBlobWriteAccess(r, h.pds, h.httpClient) + if err != nil { + RespondError(w, http.StatusForbidden, fmt.Sprintf("authorization failed: %v", err)) + return + } + + // Parse request + var req struct { + OwnerDID string `json:"ownerDid"` + Repository string `json:"repository"` + PullCount int64 `json:"pullCount"` + PushCount int64 `json:"pushCount"` + LastPull string `json:"lastPull,omitempty"` + LastPush string `json:"lastPush,omitempty"` + } + + if err := DecodeJSON(r, &req); err != nil { + RespondError(w, http.StatusBadRequest, err.Error()) + return + } + + // Verify user DID matches token (user can only set stats for their own repos) + if req.OwnerDID != validatedUser.DID { + RespondError(w, http.StatusForbidden, "owner DID mismatch") + return + } + + // Validate required fields + if req.OwnerDID == "" || req.Repository == "" { + RespondError(w, http.StatusBadRequest, "ownerDid and repository are required") + return + } + + // Set stats using the SetStats method + if err := h.pds.SetStats(ctx, req.OwnerDID, req.Repository, req.PullCount, req.PushCount, req.LastPull, req.LastPush); err != nil { + slog.Error("Failed to set stats", "error", err) + RespondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to set stats: %v", err)) + return + } + + slog.Info("Stats set via migration", "owner_did", req.OwnerDID, "repository", req.Repository, "pull_count", req.PullCount, "push_count", req.PushCount) + + RespondJSON(w, http.StatusOK, map[string]any{ + "success": true, + }) +} + // requireBlobWriteAccess middleware - validates DPoP + OAuth and checks for blob:write permission func (h *XRPCHandler) requireBlobWriteAccess(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/pkg/hold/pds/server.go b/pkg/hold/pds/server.go index a5c187f..d6cfa68 100644 --- a/pkg/hold/pds/server.go +++ b/pkg/hold/pds/server.go @@ -22,12 +22,13 @@ import ( // init registers our custom ATProto types with indigo's lexutil type registry // This allows repomgr.GetRecord to automatically unmarshal our types func init() { - // Register captain, crew, tangled profile, and layer record types + // Register captain, crew, tangled profile, layer, and stats record types // These must match the $type field in the records lexutil.RegisterType(atproto.CaptainCollection, &atproto.CaptainRecord{}) lexutil.RegisterType(atproto.CrewCollection, &atproto.CrewRecord{}) lexutil.RegisterType(atproto.LayerCollection, &atproto.LayerRecord{}) lexutil.RegisterType(atproto.TangledProfileCollection, &atproto.TangledProfileRecord{}) + lexutil.RegisterType(atproto.StatsCollection, &atproto.StatsRecord{}) } // HoldPDS is a minimal ATProto PDS implementation for a hold service diff --git a/pkg/hold/pds/stats.go b/pkg/hold/pds/stats.go new file mode 100644 index 0000000..59a95e6 --- /dev/null +++ b/pkg/hold/pds/stats.go @@ -0,0 +1,218 @@ +package pds + +import ( + "bytes" + "context" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "atcr.io/pkg/atproto" + "github.com/bluesky-social/indigo/repo" + "github.com/ipfs/go-cid" +) + +// IncrementStats increments the pull or push count for a repository +// operation should be "pull" or "push" +// Creates a new record if none exists, updates existing record otherwise +func (p *HoldPDS) IncrementStats(ctx context.Context, ownerDID, repository, operation string) error { + if operation != "pull" && operation != "push" { + return fmt.Errorf("invalid operation: %s (must be 'pull' or 'push')", operation) + } + + rkey := atproto.StatsRecordKey(ownerDID, repository) + now := time.Now().Format(time.RFC3339) + + // Try to get existing record + _, existing, err := p.GetStats(ctx, ownerDID, repository) + if err != nil { + // Record doesn't exist - create new one + record := atproto.NewStatsRecord(ownerDID, repository) + if operation == "pull" { + record.PullCount = 1 + record.LastPull = now + } else { + record.PushCount = 1 + record.LastPush = now + } + record.UpdatedAt = now + + _, _, err := p.repomgr.PutRecord(ctx, p.uid, atproto.StatsCollection, rkey, record) + if err != nil { + return fmt.Errorf("failed to create stats record: %w", err) + } + + slog.Debug("Created stats record", + "ownerDID", ownerDID, + "repository", repository, + "operation", operation) + return nil + } + + // Record exists - update it + if operation == "pull" { + existing.PullCount++ + existing.LastPull = now + } else { + existing.PushCount++ + existing.LastPush = now + } + existing.UpdatedAt = now + + _, err = p.repomgr.UpdateRecord(ctx, p.uid, atproto.StatsCollection, rkey, existing) + if err != nil { + return fmt.Errorf("failed to update stats record: %w", err) + } + + slog.Debug("Updated stats record", + "ownerDID", ownerDID, + "repository", repository, + "operation", operation, + "pullCount", existing.PullCount, + "pushCount", existing.PushCount) + return nil +} + +// GetStats retrieves the stats record for a repository +// Returns nil, nil if no stats record exists +func (p *HoldPDS) GetStats(ctx context.Context, ownerDID, repository string) (cid.Cid, *atproto.StatsRecord, error) { + rkey := atproto.StatsRecordKey(ownerDID, repository) + + recordCID, val, err := p.repomgr.GetRecord(ctx, p.uid, atproto.StatsCollection, rkey, cid.Undef) + if err != nil { + return cid.Undef, nil, err + } + + statsRecord, ok := val.(*atproto.StatsRecord) + if !ok { + return cid.Undef, nil, fmt.Errorf("unexpected type for stats record: %T", val) + } + + return recordCID, statsRecord, nil +} + +// SetStats directly sets the stats for a repository (used for migration) +// Creates or updates the stats record with the specified counts +func (p *HoldPDS) SetStats(ctx context.Context, ownerDID, repository string, pullCount, pushCount int64, lastPull, lastPush string) error { + rkey := atproto.StatsRecordKey(ownerDID, repository) + now := time.Now().Format(time.RFC3339) + + // Try to get existing record + _, existing, err := p.GetStats(ctx, ownerDID, repository) + if err != nil { + // Record doesn't exist - create new one + record := &atproto.StatsRecord{ + Type: atproto.StatsCollection, + OwnerDID: ownerDID, + Repository: repository, + PullCount: pullCount, + PushCount: pushCount, + LastPull: lastPull, + LastPush: lastPush, + UpdatedAt: now, + } + + _, _, err := p.repomgr.PutRecord(ctx, p.uid, atproto.StatsCollection, rkey, record) + if err != nil { + return fmt.Errorf("failed to create stats record: %w", err) + } + return nil + } + + // Record exists - update it + existing.PullCount = pullCount + existing.PushCount = pushCount + existing.LastPull = lastPull + existing.LastPush = lastPush + existing.UpdatedAt = now + + _, err = p.repomgr.UpdateRecord(ctx, p.uid, atproto.StatsCollection, rkey, existing) + if err != nil { + return fmt.Errorf("failed to update stats record: %w", err) + } + + return nil +} + +// ListStats returns all stats records in the hold's PDS +// This is used by AppView to aggregate stats from all holds +func (p *HoldPDS) ListStats(ctx context.Context) ([]*atproto.StatsRecord, error) { + // Get read-only session from carstore + session, err := p.carstore.ReadOnlySession(p.uid) + if err != nil { + return nil, fmt.Errorf("failed to get read-only session: %w", err) + } + + // Get repo head + 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() { + // No repo yet, return empty list + return []*atproto.StatsRecord{}, nil + } + + // Open repo + r, err := repo.OpenRepo(ctx, session, head) + if err != nil { + return nil, fmt.Errorf("failed to open repo: %w", err) + } + + var stats []*atproto.StatsRecord + + // Iterate over all stats records + err = r.ForEach(ctx, atproto.StatsCollection, func(k string, v cid.Cid) error { + // Extract collection and rkey from full path (k is like "io.atcr.hold.stats/abcd1234...") + parts := strings.Split(k, "/") + if len(parts) < 2 { + return nil // Skip invalid keys + } + + // Extract actual collection + actualCollection := strings.Join(parts[:len(parts)-1], "/") + + // MST keys are sorted, so once we hit a different collection, stop walking + if actualCollection != atproto.StatsCollection { + return repo.ErrDoneIterating + } + + // Get record bytes + _, recBytes, err := r.GetRecordBytes(ctx, k) + if err != nil { + slog.Warn("Failed to get stats record bytes", "key", k, "error", err) + return nil // Continue with other records + } + + if recBytes == nil { + return nil + } + + // Unmarshal the CBOR bytes + var statsRecord atproto.StatsRecord + if err := statsRecord.UnmarshalCBOR(bytes.NewReader(*recBytes)); err != nil { + slog.Warn("Failed to unmarshal stats record", "key", k, "error", err) + return nil // Continue with other records + } + + stats = append(stats, &statsRecord) + return nil + }) + + if err != nil { + // ErrDoneIterating is expected when we stop walking early + if errors.Is(err, repo.ErrDoneIterating) { + // Successfully stopped at collection boundary + } else if strings.Contains(err.Error(), "not found") { + // Collection doesn't exist yet - return empty list + return []*atproto.StatsRecord{}, nil + } else { + return nil, fmt.Errorf("failed to iterate stats records: %w", err) + } + } + + return stats, nil +}