diff --git a/src/main/java/com/eactive/eai/common/inflow/CustomGroupBucket.java b/src/main/java/com/eactive/eai/common/inflow/CustomGroupBucket.java new file mode 100644 index 0000000..296fb15 --- /dev/null +++ b/src/main/java/com/eactive/eai/common/inflow/CustomGroupBucket.java @@ -0,0 +1,77 @@ +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; + } +} diff --git a/src/main/java/com/eactive/eai/common/inflow/InflowControlManager.java b/src/main/java/com/eactive/eai/common/inflow/InflowControlManager.java index 5b2d6d3..29933bf 100644 --- a/src/main/java/com/eactive/eai/common/inflow/InflowControlManager.java +++ b/src/main/java/com/eactive/eai/common/inflow/InflowControlManager.java @@ -27,6 +27,7 @@ import com.ext.eai.common.stdmessage.STDMessageKeys; import io.github.bucket4j.Bandwidth; import io.github.bucket4j.Bucket4j; +import io.github.bucket4j.local.LocalBucket; import io.github.bucket4j.local.LocalBucketBuilder; @Component @@ -40,7 +41,7 @@ public class InflowControlManager implements Lifecycle, Bucket { private Map adapterBucketList = new HashMap<>(); private Map interfaceBucketList = new HashMap<>(); - private Map groupBucketList = new HashMap<>(); + private Map groupBucketList = new HashMap<>(); private Map interfaceToGroupMap = new HashMap<>(); private Map groupVoMap = new HashMap<>(); private boolean started; @@ -134,7 +135,7 @@ public class InflowControlManager implements Lifecycle, Bucket { for (InflowGroupVO groupVo : inflowGroupVoList) { if (isInflowGroupTarget(groupVo)) { - CustomBucket bucket = makeBucketUsingGroupVO(groupVo); + CustomGroupBucket bucket = makeGroupBucket(groupVo); if (bucket != null) { groupBucketList.put(groupVo.getGroupId(), bucket); groupVoMap.put(groupVo.getGroupId(), groupVo); @@ -262,7 +263,7 @@ public class InflowControlManager implements Lifecycle, Bucket { InflowGroupVO groupVo = map.get(key); if (isInflowGroupTarget(groupVo)) { - CustomBucket bucket = makeBucketUsingGroupVO(groupVo); + CustomGroupBucket bucket = makeGroupBucket(groupVo); if (bucket != null) { groupBucketList.put(groupVo.getGroupId(), bucket); groupVoMap.put(groupVo.getGroupId(), groupVo); @@ -390,30 +391,33 @@ public class InflowControlManager implements Lifecycle, Bucket { return null; } - private CustomBucket makeBucketUsingGroupVO(InflowGroupVO groupVo) { - if (isInflowGroupTarget(groupVo)) { - LocalBucketBuilder builder = Bucket4j.builder(); - if (groupVo.getThresholdPerSecond() > 0) { - builder.addLimit(Bandwidth.simple(groupVo.getThresholdPerSecond(), Duration.ofSeconds(1))); - } - - if (groupVo.getThreshold() > 0 && StringUtils.isNotBlank(groupVo.getThresholdTimeUnit())) { - builder.addLimit(Bandwidth.simple(groupVo.getThreshold(), - InflowControlManager.getPeriod(groupVo.getThresholdTimeUnit()))); - } - - // InflowTargetVO로 변환하여 CustomBucket 생성 - InflowTargetVO targetVo = new InflowTargetVO(); - targetVo.setName(groupVo.getGroupId()); - targetVo.setThreshold(groupVo.getThreshold()); - targetVo.setThresholdPerSecond(groupVo.getThresholdPerSecond()); - targetVo.setThresholdTimeUnit(groupVo.getThresholdTimeUnit()); - targetVo.setActivate(groupVo.isActivate()); - - return new CustomBucket(builder.build(), targetVo); + /** + * 그룹용 버킷 생성 - 초당/추가 임계치를 별도 버킷으로 분리 + */ + private CustomGroupBucket makeGroupBucket(InflowGroupVO groupVo) { + if (!isInflowGroupTarget(groupVo)) { + return null; } - return null; + LocalBucket perSecondBucket = null; + LocalBucket thresholdBucket = null; + + // 초당 임계치 버킷 (별도) + if (groupVo.getThresholdPerSecond() > 0) { + perSecondBucket = Bucket4j.builder() + .addLimit(Bandwidth.simple(groupVo.getThresholdPerSecond(), Duration.ofSeconds(1))) + .build(); + } + + // 추가 임계치 버킷 (별도) + if (groupVo.getThreshold() > 0 && StringUtils.isNotBlank(groupVo.getThresholdTimeUnit())) { + thresholdBucket = Bucket4j.builder() + .addLimit(Bandwidth.simple(groupVo.getThreshold(), + InflowControlManager.getPeriod(groupVo.getThresholdTimeUnit()))) + .build(); + } + + return new CustomGroupBucket(perSecondBucket, thresholdBucket, groupVo); } @Override @@ -467,12 +471,12 @@ public class InflowControlManager implements Lifecycle, Bucket { @Override public boolean isGroupPass(String groupId) { - CustomBucket b = groupBucketList.get(groupId); + CustomGroupBucket b = groupBucketList.get(groupId); if (b == null) return true; - return b.getLocalBucket().tryConsume(1); + return b.tryConsume(); } @Override