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

package cache

import (
	"context"
	"fmt"

	"github.com/redis/go-redis/v9"
)

// TryIncr implements adapter.MinuteWindowStore for DingTalk per-minute quotas.
func (c *Cache) TryIncr(ctx context.Context, key string, minute int64, limit int) (int64, bool, error) {
	if c == nil || c.rdb == nil {
		return 0, false, fmt.Errorf("redis unavailable")
	}
	redisKey := fmt.Sprintf("dingtalk:rl:%s:%d", key, minute)
	script := redis.NewScript(`
		local key = KEYS[1]
		local limit = tonumber(ARGV[1])
		local cur = tonumber(redis.call('GET', key) or '0')
		if cur >= limit then
			return {cur, 0}
		end
		local n = redis.call('INCR', key)
		if n == 1 then
			redis.call('EXPIRE', key, 120)
		end
		return {n, 1}
	`)
	res, err := script.Run(ctx, c.rdb, []string{redisKey}, limit).Slice()
	if err != nil {
		return 0, false, fmt.Errorf("dingtalk TryIncr: %w", err)
	}
	if len(res) != 2 {
		return 0, false, fmt.Errorf("dingtalk TryIncr: unexpected result %#v", res)
	}
	count, err := toInt64(res[0])
	if err != nil {
		return 0, false, err
	}
	okFlag, err := toInt64(res[1])
	if err != nil {
		return 0, false, err
	}
	return count, okFlag == 1, nil
}

func toInt64(v interface{}) (int64, error) {
	switch n := v.(type) {
	case int64:
		return n, nil
	case int:
		return int64(n), nil
	case string:
		var x int64
		_, err := fmt.Sscan(n, &x)
		return x, err
	default:
		return 0, fmt.Errorf("cannot convert %T to int64", v)
	}
}