package adapter import ( "context" "fmt" "log/slog" "net/url" "sync" "time" ) const DefaultDingTalkPerMin = 18 // MinuteWindowStore tracks per-key counters in a fixed UTC-minute window. type MinuteWindowStore interface { // TryIncr increments if under limit; returns (count, true) when acquired, // or (currentOrLimit, false) when the window is full. TryIncr(ctx context.Context, key string, minute int64, limit int) (count int64, ok bool, err error) } // DingTalkLimiter enforces per-robot (access_token) send quota and waits for the next minute when full. type DingTalkLimiter struct { store MinuteWindowStore limit int now func() time.Time sleep func(time.Duration) } func NewDingTalkLimiter(store MinuteWindowStore, limitPerMin int) *DingTalkLimiter { if limitPerMin <= 0 { limitPerMin = DefaultDingTalkPerMin } return &DingTalkLimiter{ store: store, limit: limitPerMin, now: time.Now, sleep: time.Sleep, } } // DingTalkLimitKey returns the rate-limit key for a webhook URL (access_token). func DingTalkLimitKey(webhookURL string) string { if webhookURL == "" { return "unknown" } u, err := url.Parse(webhookURL) if err != nil { return webhookURL } if tok := u.Query().Get("access_token"); tok != "" { return tok } return webhookURL } // Acquire blocks until a send slot is available for key in the current minute window. func (l *DingTalkLimiter) Acquire(ctx context.Context, key string) error { if l == nil || l.store == nil { return nil } for { if err := ctx.Err(); err != nil { return err } now := l.now() minute := now.Unix() / 60 count, ok, err := l.store.TryIncr(ctx, key, minute, l.limit) if err != nil { return fmt.Errorf("dingtalk rate limit: %w", err) } if ok { return nil } next := time.Unix((minute+1)*60, 0) wait := next.Sub(now) if wait < time.Millisecond { wait = time.Millisecond } slog.Info("dingtalk rate limited, waiting for next minute", "key_suffix", maskKey(key), "count", count, "limit", l.limit, "wait", wait, ) if err := l.sleepCtx(ctx, wait); err != nil { return err } } } func (l *DingTalkLimiter) sleepCtx(ctx context.Context, d time.Duration) error { // Prefer injectable sleep for tests; also honor context cancel via polling when possible. done := make(chan struct{}) go func() { l.sleep(d) close(done) }() select { case <-ctx.Done(): return ctx.Err() case <-done: return nil } } func maskKey(key string) string { if len(key) <= 8 { return "****" } return key[:4] + "…" + key[len(key)-4:] } // --- in-memory store (fallback / tests) --- type memoryMinuteStore struct { mu sync.Mutex data map[string]int64 // key:minute -> count } func newMemoryMinuteStore() *memoryMinuteStore { return &memoryMinuteStore{data: make(map[string]int64)} } func (m *memoryMinuteStore) TryIncr(_ context.Context, key string, minute int64, limit int) (int64, bool, error) { m.mu.Lock() defer m.mu.Unlock() k := fmt.Sprintf("%s:%d", key, minute) cur := m.data[k] if cur >= int64(limit) { return cur, false, nil } cur++ m.data[k] = cur return cur, true, nil } // NewMemoryDingTalkLimiter builds a process-local limiter (single-instance only). func NewMemoryDingTalkLimiter(limitPerMin int) *DingTalkLimiter { return NewDingTalkLimiter(newMemoryMinuteStore(), limitPerMin) }