diff --git a/atproto/client.go b/atproto/client.go index e53e4b8..d8e31d6 100644 --- a/atproto/client.go +++ b/atproto/client.go @@ -126,10 +126,6 @@ func ResolveMiniDoc(ctx context.Context, identifier string, opts *ClientOptions) return result.DID, result.PDS, result.SigningKey, nil } -func hasJSONContent(header http.Header) bool { - return len(header.Get("Content-Type")) > 0 && header.Get("Content-Type")[0:19] == "application/json" -} - func ResolveIdentity(ctx context.Context, handle string, opts *ClientOptions) (resolvedIdentity, error) { if handle == "" { return resolvedIdentity{}, fmt.Errorf("handle cannot be empty") diff --git a/atproto/repo_test.go b/atproto/repo_test.go index d50a5f2..49d3e11 100644 --- a/atproto/repo_test.go +++ b/atproto/repo_test.go @@ -897,10 +897,10 @@ func TestNewRateClient(t *testing.T) { } } -func TestRepoClient_Interfaces(t *testing.T) { - var _ RepoClient[map[string]any] = (*RateClient[map[string]any])(nil) - var _ RepoClient[map[string]any] = (*RepoClientFuncs[map[string]any])(nil) -} +var ( + _ RepoClient[map[string]any] = (*RateClient[map[string]any])(nil) + _ RepoClient[map[string]any] = (*RepoClientFuncs[map[string]any])(nil) +) func TestRateClient_ListRecords_WithTypedRecords(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/main.go b/main.go index 86d22bf..48815ec 100644 --- a/main.go +++ b/main.go @@ -1,10 +1,12 @@ package main import ( + "archive/zip" "context" "encoding/json" "fmt" "io" + "io/fs" "log/slog" "os" "slices" @@ -475,31 +477,16 @@ func (a *App) runExport(ctx context.Context, cmd *cli.Command) error { reverse := cmd.Bool("reverse") tolerance := cmd.Duration("tolerance") - lastfmRecords, err := sync.ParseInput(ctx, lastfmPath, lastfm.Parser{}) + records, _, err := loadRecordsMerge(ctx, lastfmPath, spotifyPath, tolerance) if err != nil { - return fmt.Errorf("parse lastfm: %w", err) + return fmt.Errorf("failed to deduplicate records: %w", err) } - a.log.Info("Loaded Last.fm records", slog.Int("count", len(lastfmRecords))) - - spotifyRecords, err := sync.ParseInput(ctx, spotifyPath, spotify.Parser{}) - if err != nil { - return fmt.Errorf("parse spotify: %w", err) - } - a.log.Info("Loaded Spotify records", slog.Int("count", len(spotifyRecords))) - - mergedRecords := kway.Merge([][]sync.PlayRecord{lastfmRecords, spotifyRecords}, tolerance) - - a.log.Info( - "Merged records", - slog.Int("merged_total", len(mergedRecords)), - slog.Int("duplicates_removed", len(lastfmRecords)+len(spotifyRecords)-len(mergedRecords)), - ) if reverse { - slices.Reverse(mergedRecords) + slices.Reverse(records) } - return a.outputRecords(mergedRecords, outputPath) + return a.outputRecords(records, outputPath) } func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { @@ -516,7 +503,7 @@ func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { reverse := cmd.Bool("reverse") fresh := cmd.Bool("fresh") clearCache := cmd.Bool("clear-cache") - batchSize := int(cmd.Int("batch-size")) + batchSize := cmd.Int("batch-size") tolerance := cmd.Duration("tolerance") if clearCache { @@ -527,13 +514,7 @@ func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { } } - records, totalCount, err := sync.LoadRecordsForImport(ctx, sync.ImportOptions{ - LastFMPath: lastfmPath, - SpotifyPath: spotifyPath, - Tolerance: tolerance, - LastFMParser: lastfm.Parser{}, - SpotifyParser: spotify.Parser{}, - }) + records, totalCount, err := loadRecordsMerge(ctx, lastfmPath, spotifyPath, tolerance) if err != nil { return fmt.Errorf("load records: %w", err) } @@ -588,13 +569,10 @@ func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { } } - cfg := sync.DefaultConfig - cfg.BatchSize = batchSize - progressLog := a.createProgressLogger() publishOpts := sync.PublishOptions{ - BatchSize: cfg.BatchSize, + BatchSize: batchSize, DryRun: dryRun, ATProtoClient: repoClient, ProgressLog: progressLog, @@ -626,7 +604,7 @@ func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { } } - if result.Errored() { + if result.ErrorCount > 0 { return fmt.Errorf("import completed with %d errors", result.ErrorCount) } @@ -948,3 +926,64 @@ var dedupeFlags = []cli.Flag{ Sources: cli.EnvVars(EnvYes), }, } + +type Parser interface { + ParseFile(ctx context.Context, r io.Reader) ([]sync.PlayRecord, error) + ParseFS(ctx context.Context, fsys fs.FS) ([]sync.PlayRecord, error) +} + +func parseInput(ctx context.Context, path string, parser Parser) ([]sync.PlayRecord, error) { + info, err := os.Stat(path) + if err != nil { + return nil, fmt.Errorf("stat path: %w", err) + } + + if info.IsDir() { + return parser.ParseFS(ctx, os.DirFS(path)) + } + + if strings.HasSuffix(path, ".zip") { + zf, err := zip.OpenReader(path) + if err != nil { + return nil, fmt.Errorf("open zip: %w", err) + } + + defer zf.Close() + + return parser.ParseFS(ctx, zf) + } + + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("open file: %w", err) + } + + defer file.Close() + + return parser.ParseFile(ctx, file) +} + +func loadRecordsMerge(ctx context.Context, lastFMPath, spotifyPath string, tolerance time.Duration) ([]sync.PlayRecord, int, error) { + var lastfmRecords, spotifyRecords []sync.PlayRecord + var err error + + if lastFMPath != "" { + lastfmRecords, err = parseInput(ctx, lastFMPath, lastfm.Parser{}) + if err != nil { + return nil, 0, fmt.Errorf("parse lastfm: %w", err) + } + } + + if spotifyPath != "" { + spotifyRecords, err = parseInput(ctx, spotifyPath, spotify.Parser{}) + if err != nil { + return nil, 0, fmt.Errorf("parse spotify: %w", err) + } + } + + totalInput := len(lastfmRecords) + len(spotifyRecords) + + mergedRecords := kway.Merge([][]sync.PlayRecord{lastfmRecords, spotifyRecords}, tolerance) + + return mergedRecords, totalInput, nil +} diff --git a/sync/batch_test.go b/sync/batch_test.go index 7d2d83e..18d192f 100644 --- a/sync/batch_test.go +++ b/sync/batch_test.go @@ -10,9 +10,11 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/bluesky-social/indigo/atproto/atclient" + "tangled.org/karitham.dev/lazuli/atproto" "tangled.org/karitham.dev/lazuli/cache" ) @@ -359,36 +361,34 @@ func TestPublish_Iterative(t *testing.T) { rec2, _ := json.Marshal(PlayRecord{TrackName: "Song 2"}) t.Run("Retry on transient error", func(t *testing.T) { - storage := newMockStorage() - storage.SaveRecords(did, map[string][]byte{"k1": rec1, "k2": rec2}) - - var attempts int32 - client := &mockATProtoClient{ - applyWritesFunc: func(ctx context.Context, collection string, records []PlayRecord) error { - if atomic.AddInt32(&attempts, 1) <= 2 { - return &atclient.APIError{StatusCode: 503} - } - return nil - }, - } + synctest.Test(t, func(t *testing.T) { + storage := newMockStorage() + storage.SaveRecords(did, map[string][]byte{"k1": rec1, "k2": rec2}) + + var attempts int32 + client := &mockATProtoClient{ + applyWritesFunc: func(ctx context.Context, collection string, records []PlayRecord) error { + if atomic.AddInt32(&attempts, 1) <= 2 { + return &atclient.APIError{StatusCode: 503} + } + return nil + }, + } - oldBase := BaseRetryDelay - BaseRetryDelay = time.Millisecond - defer func() { BaseRetryDelay = oldBase }() + res := Publish(ctx, &mockAuthClient{did: did}, PublishOptions{ + BatchSize: 1, + ATProtoClient: client, + Storage: storage, + ClientAgent: clientAgent, + }) - res := Publish(ctx, &mockAuthClient{did: did}, PublishOptions{ - BatchSize: 1, - ATProtoClient: client, - Storage: storage, - ClientAgent: clientAgent, + if res.SuccessCount != 2 { + t.Errorf("expected 2 successes, got %d", res.SuccessCount) + } + if atomic.LoadInt32(&attempts) < 3 { + t.Errorf("expected at least 3 attempts (2 fails + 1 success), got %d", attempts) + } }) - - if res.SuccessCount != 2 { - t.Errorf("expected 2 successes, got %d", res.SuccessCount) - } - if atomic.LoadInt32(&attempts) < 3 { - t.Errorf("expected at least 3 attempts (2 fails + 1 success), got %d", attempts) - } }) t.Run("Fail fast on non-transient error", func(t *testing.T) { diff --git a/sync/config.go b/sync/config.go deleted file mode 100644 index c992f79..0000000 --- a/sync/config.go +++ /dev/null @@ -1,55 +0,0 @@ -package sync - -import ( - "time" -) - -const ( - RecordType = "fm.teal.alpha.feed.play" - DefaultBatchSize = 20 - DefaultCrossSourceTolerance = 5 * time.Minute - CrossSourceTolerance = DefaultCrossSourceTolerance - CacheTTL = 24 * time.Hour - CacheVersion = 1 - SlingshotResolverURL = "https://slingshot.microcosm.blue/xrpc/com.bad-example.identity.resolveMiniDoc" - MaxRetryDelay = 15 * time.Minute - MaxRetries = 1000 -) - -var BaseRetryDelay = 2 * time.Second - -type Config struct { - RecordType string `json:"recordType"` - ClientAgent string `json:"clientAgent"` - BatchSize int `json:"batchSize"` - CrossSourceTolerance time.Duration `json:"crossSourceTolerance"` - CacheTTL time.Duration `json:"cacheTTL"` - CacheVersion int `json:"cacheVersion"` - SlingshotResolverURL string `json:"slingshotResolverURL"` - UserAgent string `json:"userAgent"` -} - -var DefaultConfig = Config{ - RecordType: RecordType, - ClientAgent: DefaultClientAgent, - BatchSize: DefaultBatchSize, - CrossSourceTolerance: CrossSourceTolerance, - CacheTTL: CacheTTL, - CacheVersion: CacheVersion, - SlingshotResolverURL: SlingshotResolverURL, -} - -type PublishResult struct { - SuccessCount int `json:"successCount"` - ErrorCount int `json:"errorCount"` - Cancelled bool `json:"cancelled"` - Duration time.Duration `json:"duration"` - TotalRecords int `json:"totalRecords"` - RecordsPerMinute float64 `json:"recordsPerMinute"` - FirstRecordTime time.Time `json:"firstRecordTime"` - LastRecordTime time.Time `json:"lastRecordTime"` -} - -func (r *PublishResult) Errored() bool { - return r.ErrorCount > 0 -} diff --git a/sync/import.go b/sync/import.go deleted file mode 100644 index 26c1b36..0000000 --- a/sync/import.go +++ /dev/null @@ -1,84 +0,0 @@ -package sync - -import ( - "archive/zip" - "context" - "fmt" - "io" - "io/fs" - "os" - "strings" - "time" - - "tangled.org/karitham.dev/lazuli/kway" -) - -type Parser interface { - ParseFile(ctx context.Context, r io.Reader) ([]PlayRecord, error) - ParseFS(ctx context.Context, fsys fs.FS) ([]PlayRecord, error) -} - -func ParseInput(ctx context.Context, path string, parser Parser) ([]PlayRecord, error) { - if path == "" { - return nil, nil - } - - info, err := os.Stat(path) - if err != nil { - return nil, fmt.Errorf("stat path: %w", err) - } - - if info.IsDir() { - return parser.ParseFS(ctx, os.DirFS(path)) - } - - if strings.HasSuffix(path, ".zip") { - zf, err := zip.OpenReader(path) - if err != nil { - return nil, fmt.Errorf("open zip: %w", err) - } - defer zf.Close() - return parser.ParseFS(ctx, zf) - } - - file, err := os.Open(path) - if err != nil { - return nil, fmt.Errorf("open file: %w", err) - } - defer file.Close() - return parser.ParseFile(ctx, file) -} - -type ImportOptions struct { - LastFMPath string - SpotifyPath string - Tolerance time.Duration - - LastFMParser Parser - SpotifyParser Parser -} - -func LoadRecordsForImport(ctx context.Context, opts ImportOptions) ([]PlayRecord, int, error) { - var lastfmRecords, spotifyRecords []PlayRecord - var err error - - if opts.LastFMPath != "" { - lastfmRecords, err = ParseInput(ctx, opts.LastFMPath, opts.LastFMParser) - if err != nil { - return nil, 0, fmt.Errorf("parse lastfm: %w", err) - } - } - - if opts.SpotifyPath != "" { - spotifyRecords, err = ParseInput(ctx, opts.SpotifyPath, opts.SpotifyParser) - if err != nil { - return nil, 0, fmt.Errorf("parse spotify: %w", err) - } - } - - totalInput := len(lastfmRecords) + len(spotifyRecords) - - mergedRecords := kway.Merge([][]PlayRecord{lastfmRecords, spotifyRecords}, opts.Tolerance) - - return mergedRecords, totalInput, nil -} diff --git a/sync/import_test.go b/sync/import_test.go deleted file mode 100644 index ed2b858..0000000 --- a/sync/import_test.go +++ /dev/null @@ -1,192 +0,0 @@ -package sync_test - -import ( - "context" - "encoding/json" - "fmt" - "io" - "io/fs" - "testing" - "time" - - "github.com/bluesky-social/indigo/atproto/atclient" - - "tangled.org/karitham.dev/lazuli/cache" - "tangled.org/karitham.dev/lazuli/sync" -) - -type mockParser struct { - records []sync.PlayRecord -} - -func (m *mockParser) ParseFile(ctx context.Context, r io.Reader) ([]sync.PlayRecord, error) { - return m.records, nil -} - -func (m *mockParser) ParseFS(ctx context.Context, fsys fs.FS) ([]sync.PlayRecord, error) { - return m.records, nil -} - -type mockRepoClient struct { - records []sync.RecordRef - deleted []string - applied []sync.PlayRecord -} - -func (m *mockRepoClient) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]sync.RecordRef, string, error) { - return m.records, "", nil -} - -func (m *mockRepoClient) ApplyWrites(ctx context.Context, collection string, records []sync.PlayRecord) error { - if len(records) > 200 { - return fmt.Errorf("too many records") - } - m.applied = append(m.applied, records...) - return nil -} - -func (m *mockRepoClient) DeleteRecord(ctx context.Context, collection, rkey string) error { - m.deleted = append(m.deleted, rkey) - return nil -} - -type mockAuthClient struct { - did string -} - -func (m *mockAuthClient) APIClient() *atclient.APIClient { return nil } -func (m *mockAuthClient) DID() string { return m.did } - -type mockKV struct { - data map[string]int -} - -func (m *mockKV) GetMulti(keys []string) (map[string]int, error) { - out := make(map[string]int) - for _, k := range keys { - out[k] = m.data[k] - } - return out, nil -} - -func (m *mockKV) IncrByMulti(counts map[string]int) error { - for k, v := range counts { - m.data[k] += v - } - return nil -} - -func (m *mockKV) Get(key string) (int, error) { - return m.data[key], nil -} - -func (m *mockKV) Set(key string, val int) error { - m.data[key] = val - return nil -} - -func (m *mockKV) IncrBy(key string, n int) (int, error) { - m.data[key] += n - return m.data[key], nil -} - -func TestImportE2E(t *testing.T) { - ctx := context.Background() - did := "did:plc:test" - - // 1. Setup Storage - storage, err := cache.NewBoltStorage() - if err != nil { - t.Fatal(err) - } - defer storage.Close() - defer storage.ClearAll() - - // 2. Mock Data - t1 := time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC) - t2 := time.Date(2023, 1, 1, 12, 0, 5, 0, time.UTC) // Within tolerance (5s) - t3 := time.Date(2023, 1, 2, 12, 0, 0, 0, time.UTC) // New record - - rec1 := sync.PlayRecord{TrackName: "Song A", PlayedTime: sync.Timestamp{Time: t1}} - rec2 := sync.PlayRecord{TrackName: "Song A", PlayedTime: sync.Timestamp{Time: t2}} - rec3 := sync.PlayRecord{TrackName: "Song B", PlayedTime: sync.Timestamp{Time: t3}} - - // 3. Load Records - opts := sync.ImportOptions{ - Tolerance: 10 * time.Second, - LastFMParser: &mockParser{records: []sync.PlayRecord{rec1}}, - SpotifyParser: &mockParser{records: []sync.PlayRecord{rec2, rec3}}, - LastFMPath: "import_test.go", // Use existing file to pass Stat - SpotifyPath: "import_test.go", - } - - records, total, err := sync.LoadRecordsForImport(ctx, opts) - if err != nil { - t.Fatal(err) - } - - if total != 3 { - t.Errorf("expected 3 total records, got %d", total) - } - if len(records) != 2 { - t.Errorf("expected 2 merged records, got %d", len(records)) - } - - // 4. Save to storage (as if we just imported them) - newEntries := make(map[string][]byte) - for _, rec := range records { - key := sync.CreateRecordKey(rec) - val, _ := json.Marshal(rec) - newEntries[key] = val - } - if err := storage.SaveRecords(did, newEntries); err != nil { - t.Fatal(err) - } - - // 5. Mock ATProto Client - mockRepo := &mockRepoClient{ - records: []sync.RecordRef{ - {Value: rec1}, // Already exists on remote - }, - } - - // 6. Fetch Existing (Deduplicate) - existing, err := sync.FetchExisting(ctx, mockRepo, did, storage, false) - if err != nil { - t.Fatal(err) - } - - if len(existing) != 1 { - t.Errorf("expected 1 existing record, got %d", len(existing)) - } - - // 7. Publish - kv := &mockKV{data: make(map[string]int)} - limiter := sync.NewRateLimiter(kv, 1) - publishOpts := sync.PublishOptions{ - BatchSize: 10, - ATProtoClient: mockRepo, - Storage: storage, - Limiter: limiter, - ClientAgent: sync.DefaultClientAgent, - } - - auth := &mockAuthClient{did: did} - result := sync.Publish(ctx, auth, publishOpts) - - if result.SuccessCount != 1 { - t.Errorf("expected 1 successful publish, got %d", result.SuccessCount) - } - if len(mockRepo.applied) != 1 { - t.Errorf("expected 1 record applied to repo, got %d", len(mockRepo.applied)) - } - if mockRepo.applied[0].TrackName != "Song B" { - t.Errorf("expected Song B to be published, got %s", mockRepo.applied[0].TrackName) - } - - // 8. Verify storage state - stats, _ := storage.Stats() - if stats.UnpublishedCount != 0 { - t.Errorf("expected 0 unpublished records, got %d", stats.UnpublishedCount) - } -} diff --git a/sync/publish.go b/sync/publish.go index 69af561..30e1875 100644 --- a/sync/publish.go +++ b/sync/publish.go @@ -1,6 +1,7 @@ package sync import ( + "cmp" "context" "encoding/json" "fmt" @@ -17,15 +18,51 @@ import ( "tangled.org/karitham.dev/lazuli/cache" ) -type ( - ATProtoClient = atproto.RepoClient[PlayRecord] - AuthClient = atproto.AuthClient - RateLimiter = atproto.RateLimiter -) - -var ( +const ( WriteLimitDay = atproto.WriteLimitDay GlobalLimitDay = atproto.GlobalLimitDay + + RecordType = "fm.teal.alpha.feed.play" + DefaultBatchSize = 20 + DefaultCrossSourceTolerance = 5 * time.Minute + CrossSourceTolerance = DefaultCrossSourceTolerance + CacheTTL = 24 * time.Hour + CacheVersion = 1 + SlingshotResolverURL = "https://slingshot.microcosm.blue/xrpc/com.bad-example.identity.resolveMiniDoc" + MaxRetryDelay = 15 * time.Minute + MaxRetries = 1000 + BaseRetryDelay = 2 * time.Second +) + +type ( + ATProtoClient = atproto.RepoClient[PlayRecord] + AuthClient = atproto.AuthClient + RateLimiter = atproto.RateLimiter + Client = atproto.Client + RepoClient[T any] = atproto.RepoClient[T] + RecordRef = atproto.RecordRef[PlayRecord] + + PublishOptions struct { + BatchSize int + DryRun bool + ATProtoClient ATProtoClient + ProgressLog func(ProgressReport) + Storage cache.Storage + Limiter RateLimiter + ClientAgent string + RetryDelay time.Duration + } + + PublishResult struct { + SuccessCount int `json:"successCount"` + ErrorCount int `json:"errorCount"` + Cancelled bool `json:"cancelled"` + Duration time.Duration `json:"duration"` + TotalRecords int `json:"totalRecords"` + RecordsPerMinute float64 `json:"recordsPerMinute"` + FirstRecordTime time.Time `json:"firstRecordTime"` + LastRecordTime time.Time `json:"lastRecordTime"` + } ) func NewRateLimiter(kv atproto.KVStore, maxPercent float32) RateLimiter { @@ -40,30 +77,15 @@ func IsTransientError(err error) bool { return atproto.IsTransientError(err) } -type Client = atproto.Client - func NewClient(ctx context.Context, handle, password string, opts ...func(*atproto.ClientOptions)) (*Client, error) { return atproto.NewClient(ctx, handle, password, opts...) } -type RepoClient[T any] = atproto.RepoClient[T] - -type RecordRef = atproto.RecordRef[PlayRecord] - -type PublishOptions struct { - BatchSize int - DryRun bool - ATProtoClient ATProtoClient - ProgressLog func(ProgressReport) - Storage cache.Storage - Limiter RateLimiter - ClientAgent string -} - func Publish(ctx context.Context, client AuthClient, opts PublishOptions) PublishResult { startTime := time.Now() - batchSize := defaultBatchSize(opts.BatchSize) + retryDelay := cmp.Or(opts.RetryDelay, BaseRetryDelay) + batchSize := cmp.Or(opts.BatchSize, DefaultBatchSize) atprotoClient, err := atproto.BuildClient(client, opts.ATProtoClient) if err != nil { @@ -131,7 +153,7 @@ func Publish(ctx context.Context, client AuthClient, opts PublishOptions) Publis did := client.DID() retryPolicy := retrypolicy.NewBuilder[any](). WithMaxRetries(10). - WithBackoff(BaseRetryDelay, 5*time.Minute). + WithBackoff(retryDelay, 5*time.Minute). HandleIf(func(_ any, err error) bool { return isTransientError(err) }). @@ -246,13 +268,6 @@ func defaultProgressLog(f func(ProgressReport)) func(ProgressReport) { } } -func defaultBatchSize(size int) int { - if size > 0 { - return size - } - return DefaultBatchSize -} - func newPublishResult(success, errors, total int, start time.Time, cancelled bool) PublishResult { return PublishResult{ SuccessCount: success, diff --git a/sync/rate_test.go b/sync/rate_test.go deleted file mode 100644 index d1c57a3..0000000 --- a/sync/rate_test.go +++ /dev/null @@ -1,315 +0,0 @@ -package sync - -import ( - "context" - "encoding/json" - "fmt" - "testing" - "time" - - "github.com/bluesky-social/indigo/atproto/atclient" - - "tangled.org/karitham.dev/lazuli/atproto" -) - -type mockKV struct { - data map[string]int -} - -func (m *mockKV) GetMulti(keys []string) (map[string]int, error) { - out := make(map[string]int) - for _, k := range keys { - out[k] = m.data[k] - } - return out, nil -} - -func (m *mockKV) IncrByMulti(counts map[string]int) error { - for k, v := range counts { - m.data[k] += v - } - return nil -} - -type testClock struct { - now time.Time -} - -func (m *testClock) Now() time.Time { return m.now } - -func TestRateLimiter_Refunds(t *testing.T) { - kv := &mockKV{data: make(map[string]int)} - clock := &testClock{now: time.Date(2026, 1, 22, 12, 0, 0, 0, time.UTC)} - limiter := &testRateLimiter{ - kv: kv, - prefix: "quota", - clock: clock, - rlQuota: 1, - } - ctx := context.Background() - - chargedAt, _ := limiter.AllowBulkWrite(ctx, 10) - limiter.RefundBulkWrite(ctx, 10, chargedAt) - - w, g, _ := limiter.Stats() - if w != 0 || g != 0 { - t.Errorf("BulkWrite refund failed: w=%d, g=%d", w, g) - } - - chargedAt, _ = limiter.AllowRead(ctx) - limiter.RefundRead(ctx, chargedAt) - - _, g, _ = limiter.Stats() - if g != 0 { - t.Errorf("Read refund failed: g=%d", g) - } -} - -type testRateLimiter struct { - kv *mockKV - prefix string - clock *testClock - rlQuota float32 -} - -func (l *testRateLimiter) Stats() (int, int, error) { - wd, gd, _, _, _, _ := l.getKeys(l.clock.now) - vals, err := l.kv.GetMulti([]string{wd, gd}) - if err != nil { - return 0, 0, err - } - return vals[wd], vals[gd], nil -} - -func (l *testRateLimiter) AllowBulkWrite(ctx context.Context, n int) (time.Time, error) { - wCost := n * atproto.WriteOnlyCost - gCost := n * atproto.WriteGlobalCost - - now := l.clock.now - wKeys, gKeys := l.getAllKeys(now) - - maxWait, err := l.checkQuota(now, wKeys, gKeys, wCost, gCost) - if err != nil { - return now, err - } - - if maxWait > 0 { - return now, context.DeadlineExceeded - } - - err = l.charge(wKeys, gKeys, wCost, gCost) - if err != nil { - return now, err - } - return now, nil -} - -func (l *testRateLimiter) AllowRead(ctx context.Context) (time.Time, error) { - gCost := atproto.ReadGlobalCost - - now := l.clock.now - _, gKeys := l.getAllKeys(now) - - maxWait, err := l.checkQuota(now, nil, gKeys, 0, gCost) - if err != nil { - return now, err - } - - if maxWait > 0 { - return now, context.DeadlineExceeded - } - - err = l.charge(nil, gKeys, 0, gCost) - if err != nil { - return now, err - } - return now, nil -} - -func (l *testRateLimiter) RefundBulkWrite(ctx context.Context, n int, chargedAt time.Time) { - wKeys, gKeys := l.getAllKeys(chargedAt) - wCost := n * atproto.WriteOnlyCost - gCost := n * atproto.WriteGlobalCost - - updates := make(map[string]int, len(wKeys)+len(gKeys)) - for _, k := range wKeys { - updates[k] = -wCost - } - for _, k := range gKeys { - updates[k] = -gCost - } - - l.kv.IncrByMulti(updates) -} - -func (l *testRateLimiter) RefundRead(ctx context.Context, chargedAt time.Time) { - _, gKeys := l.getAllKeys(chargedAt) - gCost := atproto.ReadGlobalCost - - updates := make(map[string]int, len(gKeys)) - for _, k := range gKeys { - updates[k] = -gCost - } - - l.kv.IncrByMulti(updates) -} - -func (l *testRateLimiter) getKeys(t time.Time) (string, string, string, string, string, string) { - day := t.Format("2006-01-02") - hour := t.Format("2006-01-02-15") - minute := t.Format("2006-01-02-15-04") - return fmt.Sprintf("%s:writes:d:%s", l.prefix, day), fmt.Sprintf("%s:global:d:%s", l.prefix, day), - fmt.Sprintf("%s:writes:h:%s", l.prefix, hour), fmt.Sprintf("%s:global:h:%s", l.prefix, hour), - fmt.Sprintf("%s:writes:m:%s", l.prefix, minute), fmt.Sprintf("%s:global:m:%s", l.prefix, minute) -} - -func (l *testRateLimiter) getAllKeys(t time.Time) ([]string, []string) { - wd, gd, wh, gh, wm, gm := l.getKeys(t) - return []string{wm, wh, wd}, []string{gm, gh, gd} -} - -func (l *testRateLimiter) checkQuota(now time.Time, wKeys, gKeys []string, wCost, gCost int) (time.Duration, error) { - wLimits := []int{atproto.WriteLimitMinute, atproto.WriteLimitHour, atproto.WriteLimitDay} - gLimits := []int{atproto.GlobalLimitMinute, atproto.GlobalLimitHour, atproto.GlobalLimitDay} - - allKeys := make([]string, 0, len(wKeys)+len(gKeys)) - allKeys = append(allKeys, wKeys...) - allKeys = append(allKeys, gKeys...) - - if len(allKeys) == 0 { - return 0, nil - } - - values, err := l.kv.GetMulti(allKeys) - if err != nil { - return 0, err - } - - maxWait := time.Duration(0) - - for i, k := range wKeys { - curr := values[k] - if curr+wCost > int(float32(wLimits[i])*l.rlQuota) { - maxWait = max(l.untilNextWindow(now, i), maxWait) - } - } - - for i, k := range gKeys { - curr := values[k] - if curr+gCost > int(float32(gLimits[i])*l.rlQuota) { - maxWait = max(l.untilNextWindow(now, i), maxWait) - } - } - - return maxWait, nil -} - -func (l *testRateLimiter) charge(wKeys, gKeys []string, wCost, gCost int) error { - updates := make(map[string]int, len(wKeys)+len(gKeys)) - for _, k := range wKeys { - updates[k] = wCost - } - for _, k := range gKeys { - updates[k] = gCost - } - - if len(updates) == 0 { - return nil - } - - return l.kv.IncrByMulti(updates) -} - -func (l *testRateLimiter) untilNextWindow(now time.Time, tier int) time.Duration { - switch tier { - case 0: - return now.Truncate(time.Minute).Add(time.Minute).Sub(now) - case 1: - return now.Truncate(time.Hour).Add(time.Hour).Sub(now) - case 2: - return now.Truncate(24 * time.Hour).Add(24 * time.Hour).Sub(now) - default: - return time.Minute - } -} - -func TestRateLimiter_Weighting(t *testing.T) { - kv := &mockKV{data: make(map[string]int)} - limiter := atproto.NewRateLimiter(kv, 1) - ctx := context.Background() - - _, err := limiter.AllowRead(ctx) - if err != nil { - t.Fatal(err) - } - _, g, err := limiter.Stats() - if err != nil { - t.Fatal(err) - } - if g != 1 { - t.Errorf("expected 1 global unit, got %d", g) - } - - _, err = limiter.AllowBulkWrite(ctx, 1) - if err != nil { - t.Fatal(err) - } - w, g, err := limiter.Stats() - if err != nil { - t.Fatal(err) - } - if w != 1 { - t.Errorf("expected 1 write unit, got %d", w) - } - if g != 4 { - t.Errorf("expected 4 global units, got %d", g) - } - - _, err = limiter.AllowBulkWrite(ctx, 10) - if err != nil { - t.Fatal(err) - } - w, g, err = limiter.Stats() - if err != nil { - t.Fatal(err) - } - if w != 11 { - t.Errorf("expected 11 write units, got %d", w) - } - if g != 34 { - t.Errorf("expected 34 global units, got %d", g) - } -} - -func TestRetryExhaustionMarkFailed(t *testing.T) { - ctx := context.Background() - did := "did:example:123" - clientAgent := "test-agent" - storage := newMockStorage() - rec1, _ := json.Marshal(PlayRecord{TrackName: "Song 1"}) - storage.SaveRecords(did, map[string][]byte{"k1": rec1}) - - client := &mockATProtoClient{ - applyWritesFunc: func(ctx context.Context, collection string, records []PlayRecord) error { - return &atclient.APIError{StatusCode: 503} - }, - } - - oldBase := BaseRetryDelay - BaseRetryDelay = time.Nanosecond - defer func() { BaseRetryDelay = oldBase }() - - res := Publish(ctx, &mockAuthClient{did: did}, PublishOptions{ - BatchSize: 1, - ATProtoClient: client, - Storage: storage, - ClientAgent: clientAgent, - }) - - if res.SuccessCount != 0 { - t.Errorf("expected 0 successes, got %d", res.SuccessCount) - } - if res.ErrorCount != 1 { - t.Errorf("expected 1 error, got %d", res.ErrorCount) - } -}