Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions changelog/unreleased/checkin-keepalive-rate-limit.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
### English

- Scheduled check-in and WorkBuddy keepalive now rate-limit across accounts: each 30s maintenance tick handles at most 2 due accounts, with ~2s + small FNV jitter between them, so a large AutoCheckin fleet no longer bursts upstream in one tick after the configured time (or after 22:00 for keepalive).

### 中文

- 定时签到与 WorkBuddy Keepalive 现按跨账号限速:每个 30 秒 maintenance tick 最多处理 2 个到期账号,账号之间约 2 秒并加小幅 FNV jitter,避免大量 AutoCheckin 账号在配置时刻(或 Keepalive 的 22:00)后同一 tick 背靠背打上游。
196 changes: 196 additions & 0 deletions internal/runtime/maintenance_burst_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
package runtime

import (
"context"
"path/filepath"
"sync"
"testing"
"time"

"github.com/caigee-cmd/cli2api/internal/accounts"
"github.com/caigee-cmd/cli2api/internal/providers"
sqlstore "github.com/caigee-cmd/cli2api/internal/store"
)

func TestMaintenanceAccountGapForStableJitter(t *testing.T) {
base := 2 * time.Second
a := maintenanceAccountGapFor("acct-a", base)
b := maintenanceAccountGapFor("acct-a", base)
c := maintenanceAccountGapFor("acct-b", base)
if a != b {
t.Fatalf("gap not stable for same id: %v vs %v", a, b)
}
if a < base || a >= base+time.Duration(maintenanceGapJitterMS)*time.Millisecond {
t.Fatalf("gap out of range: %v", a)
}
if c < base || c >= base+time.Duration(maintenanceGapJitterMS)*time.Millisecond {
t.Fatalf("gap out of range for other id: %v", c)
}
}

func TestSleepMaintenanceGapHonorsCancel(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
start := time.Now()
if sleepMaintenanceGap(ctx, "acct", time.Hour) {
t.Fatal("canceled ctx should not sleep successfully")
}
if time.Since(start) > 200*time.Millisecond {
t.Fatal("canceled gap slept too long")
}
}

type recordingKeepalive struct {
mu sync.Mutex
calls []string
}

func (r *recordingKeepalive) DailyCheckin(context.Context, string) (string, error) {
return "", nil
}

func (r *recordingKeepalive) Keepalive(_ context.Context, accountID string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.calls = append(r.calls, accountID)
return nil
}

func (r *recordingKeepalive) snapshot() []string {
r.mu.Lock()
defer r.mu.Unlock()
out := make([]string, len(r.calls))
copy(out, r.calls)
return out
}

func TestRunScheduledCheckinsCapsPerTick(t *testing.T) {
prevCap, prevGap := maintenancePerTickCap, maintenanceAccountGap
maintenancePerTickCap = 2
maintenanceAccountGap = 0
t.Cleanup(func() {
maintenancePerTickCap = prevCap
maintenanceAccountGap = prevGap
})

var mu sync.Mutex
var calls []string
manager, _ := newCheckinManager(t, checkinFunc(func(_ context.Context, accountID string) (providers.CheckinResult, error) {
mu.Lock()
calls = append(calls, accountID)
mu.Unlock()
return providers.CheckinResult{Status: "success", Message: "ok"}, nil
}))
enabled := true
for i := 0; i < 5; i++ {
if _, err := manager.store.Create(context.Background(), accounts.CreateAccount{
Name: "due", Provider: "qoder", Region: "cn", Enabled: true, AutoCheckin: &enabled,
}); err != nil {
t.Fatal(err)
}
}
// newCheckinManager already created one opted-in account → 6 due total.
location, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
now := time.Now().In(location)
dueAt := time.Date(now.Year(), now.Month(), now.Day(), 23, 59, 0, 0, location)

manager.runScheduledCheckins(context.Background(), dueAt)
mu.Lock()
first := append([]string(nil), calls...)
mu.Unlock()
if len(first) != 2 {
t.Fatalf("first tick calls=%d want 2 ids=%v", len(first), first)
}

manager.runScheduledCheckins(context.Background(), dueAt)
mu.Lock()
second := append([]string(nil), calls...)
mu.Unlock()
if len(second) != 4 {
t.Fatalf("second tick cumulative calls=%d want 4 ids=%v", len(second), second)
}

manager.runScheduledCheckins(context.Background(), dueAt)
manager.runScheduledCheckins(context.Background(), dueAt)
mu.Lock()
total := len(calls)
mu.Unlock()
if total != 6 {
t.Fatalf("after draining calls=%d want 6", total)
}
manager.runScheduledCheckins(context.Background(), dueAt)
mu.Lock()
after := len(calls)
mu.Unlock()
if after != 6 {
t.Fatalf("already-checked accounts re-ran: calls=%d", after)
}
}

func TestKeepaliveWorkBuddyCapsPerTickAndResumes(t *testing.T) {
prevCap, prevGap := maintenancePerTickCap, maintenanceAccountGap
maintenancePerTickCap = 2
maintenanceAccountGap = 0
t.Cleanup(func() {
maintenancePerTickCap = prevCap
maintenanceAccountGap = prevGap
})

store, err := sqlstore.OpenStore(filepath.Join(t.TempDir(), "keepalive.db"))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { store.Close() })
manager := NewManager(ManagerConfig{DataDir: t.TempDir()}, store, nil)
t.Cleanup(func() { manager.Close() })
rec := &recordingKeepalive{}
manager.SetWorkBuddy(rec)

enabled := true
for i := 0; i < 5; i++ {
if _, err := store.Create(context.Background(), accounts.CreateAccount{
Name: "wb", Provider: "workbuddy", Region: "cn", Enabled: true, WorkBuddyAutoCheckin: &enabled,
}); err != nil {
t.Fatal(err)
}
}

done := map[string]struct{}{}
more := manager.KeepaliveWorkBuddy(context.Background(), true, done)
if !more {
t.Fatal("expected more accounts after first capped tick")
}
if got := rec.snapshot(); len(got) != 2 {
t.Fatalf("first tick keepalive calls=%d want 2 ids=%v", len(got), got)
}
if len(done) != 2 {
t.Fatalf("done map size=%d want 2", len(done))
}

more = manager.KeepaliveWorkBuddy(context.Background(), true, done)
if !more {
t.Fatal("expected more after second tick")
}
if got := rec.snapshot(); len(got) != 4 {
t.Fatalf("second tick cumulative=%d want 4", len(got))
}

more = manager.KeepaliveWorkBuddy(context.Background(), true, done)
if more {
t.Fatal("expected no more after final tick")
}
if got := rec.snapshot(); len(got) != 5 {
t.Fatalf("final cumulative=%d want 5", len(got))
}

more = manager.KeepaliveWorkBuddy(context.Background(), true, done)
if more {
t.Fatal("repeat with full done should report no more")
}
if got := rec.snapshot(); len(got) != 5 {
t.Fatalf("repeat re-ran keepalive: %v", got)
}
}
113 changes: 98 additions & 15 deletions internal/runtime/manager_maintenance.go
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,41 @@ func recordedToday(at string, now time.Time) bool {
return local.Year() == now.Year() && local.YearDay() == now.YearDay()
}

// Cross-account burst limits for scheduled check-in / WorkBuddy keepalive.
// One 30s maintenance tick processes at most maintenancePerTickCap due
// accounts, with maintenanceAccountGap (+ FNV jitter) between them, so a
// large AutoCheckin fleet is spread across successive ticks instead of
// hammering upstream in one burst.
const (
keepaliveTickTimeout = 90 * time.Second
maintenanceGapJitterMS = 1000
)

// Defaults are conservative; tests may override the vars to skip real sleeps.
var (
maintenancePerTickCap = 2
maintenanceAccountGap = 2 * time.Second
)

// maintenanceAccountGapFor returns base delay plus a stable 0..(jitter-1)ms
// offset derived from accountID (same FNV style as checkinSlot).
func maintenanceAccountGapFor(accountID string, base time.Duration) time.Duration {
hash := fnv.New32a()
_, _ = hash.Write([]byte(accountID))
return base + time.Duration(hash.Sum32()%uint32(maintenanceGapJitterMS))*time.Millisecond
}

func sleepMaintenanceGap(ctx context.Context, accountID string, base time.Duration) bool {
timer := time.NewTimer(maintenanceAccountGapFor(accountID, base))
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}

func checkinSlot(value, accountID string, now time.Time) (time.Time, error) {
parsed, err := time.Parse("15:04", value)
if err != nil {
Expand Down Expand Up @@ -177,6 +212,7 @@ func (manager *Manager) runScheduledCheckins(ctx context.Context, now time.Time)
return
}
allowDisabled := manager.allowDisabledCheckin(ctx)
processed := 0
for _, account := range items {
if ctx.Err() != nil {
return
Expand All @@ -196,11 +232,20 @@ func (manager *Manager) runScheduledCheckins(ctx context.Context, now time.Time)
log.Printf("checkin settings account=%s: %v", account.ID, err)
continue
}
if checkinDueFor(account, configuredTime, now.In(location), allowDisabled) {
if _, err := manager.checkinAccount(ctx, account.ID, allowDisabled); err != nil {
log.Printf("checkin account=%s: %v", account.ID, err)
}
if !checkinDueFor(account, configuredTime, now.In(location), allowDisabled) {
continue
}
if processed >= maintenancePerTickCap {
// Remaining due accounts stay due; the next 30s tick picks them up.
return
}
if processed > 0 && !sleepMaintenanceGap(ctx, account.ID, maintenanceAccountGap) {
return
}
if _, err := manager.checkinAccount(ctx, account.ID, allowDisabled); err != nil {
log.Printf("checkin account=%s: %v", account.ID, err)
}
processed++
}
}

Expand Down Expand Up @@ -251,26 +296,54 @@ func (manager *Manager) allowDisabledCheckin(ctx context.Context) bool {
}
}

func (manager *Manager) KeepaliveWorkBuddy(ctx context.Context, onlyOptIn bool) {
// KeepaliveWorkBuddy refreshes WorkBuddy sessions. At most
// maintenancePerTickCap accounts are processed per call, with
// maintenanceAccountGap between them. done tracks account IDs already
// refreshed for the current local day; pass the same map across successive
// maintenance ticks. Returns true when more eligible accounts remain.
func (manager *Manager) KeepaliveWorkBuddy(ctx context.Context, onlyOptIn bool, done map[string]struct{}) (more bool) {
if manager == nil || manager.workbuddy == nil {
return
return false
}
if done == nil {
done = map[string]struct{}{}
}
items, err := manager.store.List(ctx)
if err != nil {
log.Printf("workbuddy keepalive list: %v", err)
return
return false
}
var pending []string
for _, account := range items {
if ctx.Err() != nil {
return
}
if account.Provider != "workbuddy" || !account.Enabled || (onlyOptIn && !account.AutoCheckin) {
continue
}
if err := manager.workbuddy.Keepalive(ctx, account.ID); err != nil {
log.Printf("workbuddy keepalive account=%s: %v", account.ID, err)
if _, ok := done[account.ID]; ok {
continue
}
pending = append(pending, account.ID)
}
if len(pending) == 0 {
return false
}
batch := pending
if len(batch) > maintenancePerTickCap {
batch = batch[:maintenancePerTickCap]
more = true
}
for i, accountID := range batch {
if ctx.Err() != nil {
return true
}
if i > 0 && !sleepMaintenanceGap(ctx, accountID, maintenanceAccountGap) {
return true
}
if err := manager.workbuddy.Keepalive(ctx, accountID); err != nil {
log.Printf("workbuddy keepalive account=%s: %v", accountID, err)
}
done[accountID] = struct{}{}
}
return more
}

func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) {
Expand All @@ -286,6 +359,8 @@ func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
lastKeepaliveDay := ""
keepaliveDone := map[string]struct{}{}
keepaliveProgressDay := ""
// Seed the throttle so the first pass does not duplicate the boot-time
// RefreshAll: the next stale-quota sweep runs 5 minutes after startup.
lastStaleQuota := time.Now()
Expand All @@ -300,11 +375,19 @@ func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) {
now := time.Now()
manager.runScheduledCheckins(ctx, now)
day := now.Format("2006-01-02")
if day != keepaliveProgressDay {
keepaliveProgressDay = day
keepaliveDone = map[string]struct{}{}
}
if now.Hour() >= 22 && lastKeepaliveDay != day {
keepaliveCtx, stopKeepalive := context.WithTimeout(ctx, 2*time.Minute)
manager.KeepaliveWorkBuddy(keepaliveCtx, true)
// Cap + gap per tick; unfinished accounts continue on later ticks
// instead of one 2-minute burst after 22:00.
keepaliveCtx, stopKeepalive := context.WithTimeout(ctx, keepaliveTickTimeout)
more := manager.KeepaliveWorkBuddy(keepaliveCtx, true, keepaliveDone)
stopKeepalive()
lastKeepaliveDay = day
if !more {
lastKeepaliveDay = day
}
}
manager.sampleResources()
if now.Sub(lastLiveness) >= 30*time.Second {
Expand Down
Loading