aiaa-notification-server/internal/adapter/dingtalk_limit.go

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)
}