From 78a7519de54cc05302d95d9eaaf0a0c4f02b00b7 Mon Sep 17 00:00:00 2001 From: karitham Date: Thu, 22 Jan 2026 20:21:33 +0100 Subject: [PATCH] feat: rate limit quotas per period --- flake.nix | 2 +- main.go | 22 +++++++++++----------- sync/import_test.go | 3 ++- sync/rate.go | 39 ++++++++++++++++++--------------------- sync/rate_test.go | 23 +++++++++++++---------- 5 files changed, 45 insertions(+), 44 deletions(-) diff --git a/flake.nix b/flake.nix index 4eefad6..40074a3 100644 --- a/flake.nix +++ b/flake.nix @@ -17,7 +17,7 @@ let lazuli = pkgs.buildGoModule rec { name = "lazuli"; - version = "0.1.2"; + version = "0.1.3"; src = pkgs.nix-gitignore.gitignoreSource [ "*.csv" "*.zip" "*.json" ] ./.; vendorHash = "sha256-MfBPv/L7wHuUGXx4BDd+DFq0RB11KuMHCzPjFv6FMgs="; ldflags = [ diff --git a/main.go b/main.go index 66b25b7..be814d5 100644 --- a/main.go +++ b/main.go @@ -160,7 +160,7 @@ func (a *App) runStats(ctx context.Context, cmd *cli.Command) error { return fmt.Errorf("failed to get database stats: %w", err) } - limiter := sync.NewRateLimiter(a.storage) + limiter := sync.NewRateLimiter(a.storage, 1) writes, global, err := limiter.Stats() if err != nil { return fmt.Errorf("failed to get rate limit stats: %w", err) @@ -247,7 +247,7 @@ func (a *App) runRetry(ctx context.Context, cmd *cli.Command) error { did := authClient.GetDID() dryRun := cmd.Bool("dry-run") - limiter := sync.NewRateLimiter(a.storage) + limiter := sync.NewRateLimiter(a.storage, 0.9) repoClient := sync.NewRateClient(authClient.GetAPIClient(), did, limiter) var failedRecords []struct { @@ -383,14 +383,14 @@ func (a *App) versionCommand() *cli.Command { func (a *App) runDebugFetch(ctx context.Context, cmd *cli.Command) error { authClient, err := a.prepareAuth(ctx, cmd) if err != nil { - return fmt.Errorf("authentication failed: %w\nHint: Make sure your credentials are correct and you have network access.", err) + return fmt.Errorf("authentication failed: %w", err) } repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), nil) records, _, err := repoClient.ListRecords(ctx, sync.RecordType, 10, "") if err != nil { - return fmt.Errorf("failed to fetch records from Bluesky: %w\nHint: Check your network connection and try again.", err) + return fmt.Errorf("failed to fetch records from Bluesky: %w", err) } enc := json.NewEncoder(os.Stdout) @@ -435,7 +435,7 @@ func (a *App) getCredentials(cmd *cli.Command) (string, string, error) { password := cmd.String("password") if handle == "" { - return "", "", fmt.Errorf("Bluesky handle is required (set --handle or set the LAZULI_HANDLE environment variable)") + return "", "", fmt.Errorf("bluesky handle is required (set --handle or set the LAZULI_HANDLE environment variable)") } if password == "" { return "", "", fmt.Errorf("app password is required (set --password or set the LAZULI_PASSWORD environment variable)") @@ -551,7 +551,7 @@ func (a *App) runImport(ctx context.Context, cmd *cli.Command) error { } a.log.Info("Authenticated", logutil.DID(authClient.GetDID()), slog.String("pds", authClient.GetPDS())) - limiter := sync.NewRateLimiter(a.storage) + limiter := sync.NewRateLimiter(a.storage, 0.9) repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.GetDID(), a.storage, fresh) @@ -663,7 +663,7 @@ func (a *App) runSync(ctx context.Context, cmd *cli.Command) error { fresh := cmd.Bool("fresh") a.log.Info("Starting sync operation", logutil.DID(authClient.GetDID()), slog.Bool("fresh", fresh)) - limiter := sync.NewRateLimiter(a.storage) + limiter := sync.NewRateLimiter(a.storage, 0.85) repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) if fresh { @@ -687,7 +687,7 @@ func (a *App) runSync(ctx context.Context, cmd *cli.Command) error { func (a *App) runDedupe(ctx context.Context, cmd *cli.Command) error { authClient, err := a.prepareAuth(ctx, cmd) if err != nil { - return fmt.Errorf("authentication failed: %w\nHint: Make sure your credentials are correct.", err) + return fmt.Errorf("authentication failed: %w", err) } dryRun := cmd.Bool("dry-run") @@ -698,7 +698,7 @@ func (a *App) runDedupe(ctx context.Context, cmd *cli.Command) error { slog.Bool("dry_run", dryRun), slog.Bool("fresh", fresh)) - limiter := sync.NewRateLimiter(a.storage) + limiter := sync.NewRateLimiter(a.storage, 0.9) repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) if fresh { @@ -711,7 +711,7 @@ func (a *App) runDedupe(ctx context.Context, cmd *cli.Command) error { existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.GetDID(), a.storage, fresh) if err != nil { - return fmt.Errorf("failed to fetch existing records: %w\nHint: Check your network connection and try again.", err) + return fmt.Errorf("failed to fetch existing records: %w", err) } duplicates := sync.FindDuplicates(existingRecords) @@ -777,7 +777,7 @@ func (a *App) runDedupe(ctx context.Context, cmd *cli.Command) error { }). Build() - err := failsafe.With[any](retryPolicy).WithContext(ctx).Run(func() error { + err := failsafe.With(retryPolicy).WithContext(ctx).Run(func() error { return repoClient.DeleteRecord(ctx, sync.RecordType, rkey) }) diff --git a/sync/import_test.go b/sync/import_test.go index 03c63b0..143af00 100644 --- a/sync/import_test.go +++ b/sync/import_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/bluesky-social/indigo/atproto/atclient" + "tangled.org/karitham.dev/lazuli/cache" "tangled.org/karitham.dev/lazuli/sync" ) @@ -162,7 +163,7 @@ func TestImportE2E(t *testing.T) { // 7. Publish kv := &mockKV{data: make(map[string]int)} - limiter := sync.NewRateLimiter(kv) + limiter := sync.NewRateLimiter(kv, 1) publishOpts := sync.PublishOptions{ BatchSize: 10, ATProtoClient: mockRepo, diff --git a/sync/rate.go b/sync/rate.go index fee8616..ab5ddbc 100644 --- a/sync/rate.go +++ b/sync/rate.go @@ -17,7 +17,7 @@ const ( // Limits WriteLimitMinute = 100 WriteLimitHour = 1000 - WriteLimitDay = 9000 + WriteLimitDay = 10000 GlobalLimitMinute = 300 GlobalLimitHour = 3000 @@ -57,10 +57,11 @@ type realClock struct{} func (realClock) Now() time.Time { return time.Now().UTC() } type quotaLimiter struct { - kv KVStore - prefix string - clock Clock - mu sync.Mutex + mu sync.Mutex + kv KVStore + prefix string + clock Clock + rlQuota float32 } func (l *quotaLimiter) Stats() (int, int, error) { @@ -72,11 +73,12 @@ func (l *quotaLimiter) Stats() (int, int, error) { return vals[wd], vals[gd], nil } -func NewRateLimiter(kv KVStore) RateLimiter { +func NewRateLimiter(kv KVStore, maxPercent float32) RateLimiter { return "aLimiter{ - kv: kv, - prefix: "quota", - clock: realClock{}, + kv: kv, + prefix: "quota", + clock: realClock{}, + rlQuota: maxPercent, } } @@ -228,23 +230,18 @@ func (l *quotaLimiter) checkQuota(now time.Time, wKeys, gKeys []string, wCost, g for i, k := range wKeys { curr := values[k] - if curr+wCost > wLimits[i] { - wait := l.untilNextWindow(now, i) - if wait > maxWait { - maxWait = wait - } + 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 > gLimits[i] { - wait := l.untilNextWindow(now, i) - if wait > maxWait { - maxWait = wait - } + if curr+gCost > int(float32(gLimits[i])*l.rlQuota) { + maxWait = max(l.untilNextWindow(now, i), maxWait) } } + return maxWait, nil } @@ -289,8 +286,8 @@ func addJitter(d time.Duration) time.Duration { var b [8]byte _, _ = rand.Read(b[:]) n := binary.LittleEndian.Uint64(b[:]) - // Add 0-20% jitter - jitter := float64(d) * 0.2 * (float64(n) / math.MaxUint64) + // Add 0-10% jitter + jitter := float64(d) * 0.1 * (float64(n) / math.MaxUint64) return d + time.Duration(jitter) } diff --git a/sync/rate_test.go b/sync/rate_test.go index ff98821..fdedb55 100644 --- a/sync/rate_test.go +++ b/sync/rate_test.go @@ -38,9 +38,10 @@ func TestRateLimiter_Refunds(t *testing.T) { kv := &mockKV{data: make(map[string]int)} clock := &mockClock{now: time.Date(2026, 1, 22, 12, 0, 0, 0, time.UTC)} limiter := "aLimiter{ - kv: kv, - prefix: "quota", - clock: clock, + kv: kv, + prefix: "quota", + clock: clock, + rlQuota: 1, } ctx := context.Background() @@ -65,7 +66,7 @@ func TestRateLimiter_Refunds(t *testing.T) { func TestRateLimiter_Weighting(t *testing.T) { kv := &mockKV{data: make(map[string]int)} - limiter := NewRateLimiter(kv) + limiter := NewRateLimiter(kv, 1) ctx := context.Background() // 1 Read = 1 Global @@ -119,9 +120,10 @@ func TestRateLimiter_Smoothing(t *testing.T) { // Ensure we are at the very beginning of the minute to avoid window edge issues in test clock := &mockClock{now: time.Date(2026, 1, 22, 1, 0, 0, 0, time.UTC)} limiter := "aLimiter{ - kv: kv, - prefix: "quota", - clock: clock, + kv: kv, + prefix: "quota", + clock: clock, + rlQuota: 1, } wd, gd, wh, gh, wm, gm := limiter.getKeys(clock.now) @@ -198,9 +200,10 @@ func TestRateLimiter_MidnightRollover(t *testing.T) { kv := &mockKV{data: make(map[string]int)} clock := &mockClock{now: time.Date(2026, 1, 22, 23, 59, 59, 0, time.UTC)} limiter := "aLimiter{ - kv: kv, - prefix: "quota", - clock: clock, + kv: kv, + prefix: "quota", + clock: clock, + rlQuota: 1, } ctx := context.Background() -- 2.51.2