매일메일 AWS SES 처리율 제한 장치 구현

@Component
public class RateLimiter {

    private final String bucketKey;
    private final Duration waitTimeout;
    private final BucketConfiguration bucketConfiguration;
    private final MySQLSelectForUpdateBasedProxyManager<Long> proxyManager;
    private final ConcurrentMap<String, BucketProxy> buckets = new ConcurrentHashMap<>();

    public RateLimiter(
            DataSource dataSource,
            @Value("${mail.ses.rate-limit.bucket-key}") String bucketKey,
            @Value("${mail.ses.rate-limit.capacity}") int capacity,
            @Value("${mail.ses.rate-limit.refill-amount}") int refillAmount,
            @Value("${mail.ses.rate-limit.refill-seconds}") long refillSeconds,
            @Value("${mail.ses.rate-limit.wait-timeout-millis}") long waitTimeoutMillis
    ) {
        this.bucketKey = bucketKey;
        this.waitTimeout = Duration.ofMillis(waitTimeoutMillis);
        this.bucketConfiguration = createConfiguration(capacity, refillAmount, Duration.ofSeconds(refillSeconds));

        SQLProxyConfiguration<Long> sqlProxyConfiguration = SQLProxyConfiguration.builder()
                .withTableSettings(BucketTableSettings.getDefault())
                .build(dataSource);
        this.proxyManager = new MySQLSelectForUpdateBasedProxyManager<>(sqlProxyConfiguration);
    }

    private BucketConfiguration createConfiguration(int capacity, int refillAmount, Duration refillDuration) {
        return BucketConfiguration.builder()
                .addLimit(Bandwidth.builder()
                        .capacity(capacity)
                        .refillIntervally(refillAmount, refillDuration)
                        .build())
                .build();
    }

    public void tryConsume() {
        if (!tryConsume(bucketKey, bucketConfiguration, waitTimeout)) {
            throw new IllegalStateException("SES 처리율 제한을 초과했습니다.");
        }
    }

    private boolean tryConsume(String key, BucketConfiguration configuration, Duration waitTimeout) {
        BucketProxy bucket = getOrCreateBucket(key, configuration);
        try {
            return bucket.asBlocking().tryConsume(1, waitTimeout);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            return false;
        }
    }

    private BucketProxy getOrCreateBucket(String key, BucketConfiguration configuration) {
        return buckets.computeIfAbsent(key, bucketKey -> {
            long bucketId = bucketKey.hashCode();
            return proxyManager.builder().build(bucketId, configuration);
        });
    }
}

@Component
public class DistributedRateLimitSupport {

    private final DistributedTokenLeaseService leaseService;
    private final int leaseAmount;

    private long leasedTokens = 0;
    private long leaseExpiresAtMillis = 0;

    public synchronized boolean tryConsume() {
        // 로컬에 임대한 토큰은 다음 refill 경계까지만 유효하므로, 소비 전에 먼저 만료 여부를 확인한다.
        expireLeasedTokens(Instant.now());

        if (!ensureLeasedTokens()) {
            return false;
        }

        return consume();
    }

    private void expireLeasedTokens(Instant now) {
        if (leasedTokens > 0 && (now.toEpochMilli() >= leaseExpiresAtMillis)) {
            leasedTokens = 0;
            leaseExpiresAtMillis = 0;
        }
    }

    private boolean ensureLeasedTokens() {
        // 아직 유효한 임대 토큰이 남아 있으면 중앙 버킷을 조회하지 않고 로컬 토큰을 사용한다.
        if (leasedTokens > 0) {
            return true;
        }

        // 로컬 토큰이 없을 때만 중앙 버킷에서 leaseAmount만큼 토큰을 빌린다.
        if (leaseService.tryLeaseToken(leaseAmount)) {
            addLeasedTokens(leaseAmount, Instant.now());
            return true;
        }

        return false;
    }

    private void addLeasedTokens(int amount, Instant acquiredAt) {
        leasedTokens = amount;
        // bucket4j의 intervally aligned refill 정책과 동일한 초 경계를 로컬 토큰의 만료 시각으로 사용한다.
        leaseExpiresAtMillis = calculateNextRefillBoundary(acquiredAt).toEpochMilli();
    }

    private Instant calculateNextRefillBoundary(Instant now) {
        return Instant.ofEpochSecond(now.getEpochSecond() + 1);
    }

    private boolean consume() {
        if (leasedTokens < 1) {
            return false;
        }

        leasedTokens -= 1;
        if (leasedTokens == 0) {
            leaseExpiresAtMillis = 0;
        }

        return true;
    }
}

private void doSend(T message) {
    try {
        limiter.consumeBlocking(WAIT_TIMEOUT);
        logSending(message);
        MimeMessage emptyMimeMessage = javaMailSender.createMimeMessage();
        MimeMessage targetMimeMessage = mimeMessageCustomizer.customize(emptyMimeMessage, message);
        javaMailSender.send(targetMimeMessage);
        handleSuccess(message);
    } catch (RateLimitExceededException e) {
        logMailSendingFailed(e);
        throw new RetryableMailException(e);
    } catch (MailException e) {
        logMailSendingFailed(e);
        throwWhenCanRetry(e);
        handleAmbiguous(message);
    } catch (Exception e) {
        log.error("예기치 않은 오류 발생: {}", e.getMessage(), e);
        handleFailure(message);
    }
}