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 } // ListStatsRecordsForUser returns all stats records where the user is the repository owner // Used for GDPR data export to return all stats for repositories owned by the user func (p *HoldPDS) ListStatsRecordsForUser(ctx context.Context, userDID string) ([]*atproto.StatsRecord, error) { // Get all stats records and filter by ownerDID allStats, err := p.ListStats(ctx) if err != nil { return nil, err } var userStats []*atproto.StatsRecord for _, stat := range allStats { if stat.OwnerDID == userDID { userStats = append(userStats, stat) } } if userStats == nil { userStats = []*atproto.StatsRecord{} } return userStats, nil }