elink-online-common/src/main/java/com/eactive/eai/common/inflow/CustomGroupBucket.java

package com.eactive.eai.common.inflow;

import io.github.bucket4j.local.LocalBucket;

/**
 * 그룹 유량제어 전용 버킷
 * 초당 임계치와 추가 임계치를 별도 버킷으로 관리하여 개별 토큰 조회 가능
 */
public class CustomGroupBucket {
    private LocalBucket perSecondBucket;   // 초당 임계치 버킷 (null 가능)
    private LocalBucket thresholdBucket;   // 추가 임계치 버킷 (null 가능)
    private InflowGroupVO groupVo;

    public CustomGroupBucket(LocalBucket perSecondBucket, LocalBucket thresholdBucket, InflowGroupVO groupVo) {
        this.perSecondBucket = perSecondBucket;
        this.thresholdBucket = thresholdBucket;
        this.groupVo = groupVo;
    }

    /**
     * 두 버킷 모두에서 토큰 소비 시도
     * @return 모든 버킷에서 소비 성공 시 true
     */
    public boolean tryConsume() {
        boolean perSecondPass = (perSecondBucket == null) || perSecondBucket.tryConsume(1);
        if (!perSecondPass) {
            return false;
        }

        boolean thresholdPass = (thresholdBucket == null) || thresholdBucket.tryConsume(1);
        if (!thresholdPass) {
            // 초당 버킷에서 이미 소비했으므로 롤백은 불가 (Token Bucket 특성상)
            // 하지만 초당 버킷은 빠르게 리필되므로 실질적 영향 미미
            return false;
        }

        return true;
    }

    /**
     * 초당 임계치 버킷의 사용 가능한 토큰 수
     */
    public long getPerSecondAvailableTokens() {
        return (perSecondBucket != null) ? perSecondBucket.getAvailableTokens() : -1;
    }

    /**
     * 추가 임계치 버킷의 사용 가능한 토큰 수
     */
    public long getThresholdAvailableTokens() {
        return (thresholdBucket != null) ? thresholdBucket.getAvailableTokens() : -1;
    }

    public LocalBucket getPerSecondBucket() {
        return perSecondBucket;
    }

    public void setPerSecondBucket(LocalBucket perSecondBucket) {
        this.perSecondBucket = perSecondBucket;
    }

    public LocalBucket getThresholdBucket() {
        return thresholdBucket;
    }

    public void setThresholdBucket(LocalBucket thresholdBucket) {
        this.thresholdBucket = thresholdBucket;
    }

    public InflowGroupVO getGroupVo() {
        return groupVo;
    }

    public void setGroupVo(InflowGroupVO groupVo) {
        this.groupVo = groupVo;
    }
}