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

package com.eactive.eai.common.inflow;

import java.time.Duration;
import java.util.Arrays;
import java.util.Calendar;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;

import com.eactive.eai.common.exception.ExceptionUtil;
import com.eactive.eai.common.lifecycle.Lifecycle;
import com.eactive.eai.common.lifecycle.LifecycleException;
import com.eactive.eai.common.lifecycle.LifecycleListener;
import com.eactive.eai.common.lifecycle.LifecycleSupport;
import com.eactive.eai.common.util.Logger;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import io.github.bucket4j.local.LocalBucket;
import io.github.bucket4j.local.LocalBucketBuilder;

public abstract class AbstractInflowControlManager implements Lifecycle, Bucket {
	static Logger logger = Logger.getLogger(Logger.LOGGER_DEFAULT);
	public static final String QUOTA_TIMEUNIT_SECOND = "SEC";
	public static final String QUOTA_TIMEUNIT_MINUTE = "MIN";
	public static final String QUOTA_TIMEUNIT_HOUR = "HOU";
	public static final String QUOTA_TIMEUNIT_DAY = "DAY";
	public static final String QUOTA_TIMEUNIT_MONTH = "MON";

	private volatile Map<String, CustomBucket> adapterBucketList = new ConcurrentHashMap<>();
	private volatile Map<String, CustomBucket> interfaceBucketList = new ConcurrentHashMap<>();
	private volatile Map<String, CustomGroupBucket> groupBucketList = new ConcurrentHashMap<>();
	private volatile Map<String, String> interfaceToGroupMap = new ConcurrentHashMap<>();
	private volatile Map<String, InflowGroupVO> groupVoMap = new ConcurrentHashMap<>();
	private boolean started;
	private LifecycleSupport lifecycle = new LifecycleSupport(this);

	@Autowired
	protected InflowControlDAO inflowControlDAO;

	public void start() throws LifecycleException {
		if (started)
			throw new LifecycleException("RECEAICMM201");

		lifecycle.fireLifecycleEvent(STARTING_EVENT, this);
		try {
			init();
		} catch (Exception e) {
			throw new LifecycleException(ExceptionUtil.getErrorCode(e, "RECEAICMM202"));
		}
		started = true;
		lifecycle.fireLifecycleEvent(STARTED_EVENT, this);
	}

	private void init() throws Exception {
		initAdapter();
		initInterface();
		initGroup();
	}

	private void initAdapter() throws Exception {
		Map<String, CustomBucket> newMap = new HashMap<>();
		List<InflowTargetVO> inflowAdapterVoList = inflowControlDAO.getInflowAdapterList();

		for (InflowTargetVO inflowVo : inflowAdapterVoList) {
			if (AbstractInflowControlManager.isInflowTarget(inflowVo)) {
				CustomBucket bucket = makeBucketUsingInflowVO(inflowVo);
				if (bucket != null) {
					newMap.put(inflowVo.getName(), bucket);
				}
			}
		}
		this.adapterBucketList = newMap;
	}

	private void initInterface() throws Exception {
		Map<String, CustomBucket> newMap = new HashMap<>();
		List<InflowTargetVO> inflowInterfaceVoList = inflowControlDAO.getInflowInterfaceList();

		for (InflowTargetVO inflowVo : inflowInterfaceVoList) {
			if (AbstractInflowControlManager.isInflowTarget(inflowVo)) {
				CustomBucket bucket = makeBucketUsingInflowVO(inflowVo);
				if (bucket != null) {
					newMap.put(inflowVo.getName(), bucket);
				}
			}
		}
		this.interfaceBucketList = newMap;
	}

	private void initGroup() throws Exception {
		Map<String, CustomGroupBucket> newGroupBucketList = new HashMap<>();
		Map<String, String> newInterfaceToGroupMap = new HashMap<>();
		Map<String, InflowGroupVO> newGroupVoMap = new HashMap<>();

		List<InflowGroupVO> inflowGroupVoList = inflowControlDAO.getInflowGroupList();

		for (InflowGroupVO groupVo : inflowGroupVoList) {
			if (isInflowGroupTarget(groupVo)) {
				CustomGroupBucket bucket = makeGroupBucket(groupVo);
				if (bucket != null) {
					newGroupBucketList.put(groupVo.getGroupId(), bucket);
					newGroupVoMap.put(groupVo.getGroupId(), groupVo);
					for (String interfaceId : groupVo.getInterfaceList()) {
						newInterfaceToGroupMap.put(interfaceId, groupVo.getGroupId());
					}
				}
			}
		}

		if (logger.isInfo()) {
			logger.info("InflowControlManager] initGroup completed. groups=" + newGroupBucketList.size()
					+ ", mappings=" + newInterfaceToGroupMap.size());
		}

		this.groupBucketList = newGroupBucketList;
		this.interfaceToGroupMap = newInterfaceToGroupMap;
		this.groupVoMap = newGroupVoMap;
	}

	public synchronized void reloadAdapter() throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload all Started ...");
		}
		initAdapter();
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload all finished ...");
		}
	}

	public synchronized void reloadInterface() throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload all Started ...");
		}
		initInterface();
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload all finished ...");
		}
	}

	public synchronized void reloadAdapter(String adapter) throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload Started ...");
		}
		Map<String, InflowTargetVO> map = inflowControlDAO.getInflowTargetByAdater(adapter);
		if (map != null) {
			Iterator<String> it = map.keySet().iterator();
			String key = "";
			while (it.hasNext()) {
				key = it.next();

				InflowTargetVO inflowTarget = map.get(key);
				if (AbstractInflowControlManager.isInflowTarget(inflowTarget)) {
					CustomBucket bucket = makeBucketUsingInflowVO(inflowTarget);
					if (bucket != null) {
						adapterBucketList.put(inflowTarget.getName(), bucket);
					} else {
						adapterBucketList.remove(inflowTarget.getName());
					}
				} else {
					adapterBucketList.remove(inflowTarget.getName());
				}
			}
		}
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload finished ...");
		}
	}

	public synchronized void reloadInterface(String inter) throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload Started ...");
		}
		Map<String, InflowTargetVO> map = inflowControlDAO.getInflowTargetByInterface(inter);
		if (map != null) {
			Iterator<String> it = map.keySet().iterator();
			String key = "";
			while (it.hasNext()) {
				key = it.next();

				InflowTargetVO inflowTarget = map.get(key);
				if (AbstractInflowControlManager.isInflowTarget(inflowTarget)) {
					CustomBucket bucket = makeBucketUsingInflowVO(inflowTarget);
					if (bucket != null) {
						interfaceBucketList.put(inflowTarget.getName(), bucket);
					} else {
						interfaceBucketList.remove(inflowTarget.getName());
					}
				} else {
					interfaceBucketList.remove(inflowTarget.getName());
				}
			}
		}

		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reload finished ...");
		}
	}

	public synchronized void reloadGroup() throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reloadGroup all Started ...");
		}
		initGroup();
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reloadGroup all finished ...");
		}
	}

	public synchronized void reloadGroup(String groupId) throws Exception {
		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reloadGroup Started ... groupId=" + groupId);
		}
		Map<String, InflowGroupVO> map = inflowControlDAO.getInflowTargetByGroup(groupId);
		if (map != null) {
			Iterator<Map.Entry<String, String>> mappingIt = interfaceToGroupMap.entrySet().iterator();
			while (mappingIt.hasNext()) {
				Map.Entry<String, String> entry = mappingIt.next();
				if (groupId.equals(entry.getValue())) {
					mappingIt.remove();
				}
			}

			Iterator<String> it = map.keySet().iterator();
			while (it.hasNext()) {
				String key = it.next();
				InflowGroupVO groupVo = map.get(key);

				if (isInflowGroupTarget(groupVo)) {
					CustomGroupBucket bucket = makeGroupBucket(groupVo);
					if (bucket != null) {
						groupBucketList.put(groupVo.getGroupId(), bucket);
						groupVoMap.put(groupVo.getGroupId(), groupVo);
						for (String interfaceId : groupVo.getInterfaceList()) {
							interfaceToGroupMap.put(interfaceId, groupVo.getGroupId());
						}
					} else {
						groupBucketList.remove(groupVo.getGroupId());
						groupVoMap.remove(groupVo.getGroupId());
					}
				} else {
					groupBucketList.remove(groupVo.getGroupId());
					groupVoMap.remove(groupVo.getGroupId());
				}
			}
		}

		if (logger.isWarn()) {
			logger.warn("InflowControlManager] reloadGroup finished ...");
		}
	}

	public void stop() throws LifecycleException {
		if (!started)
			throw new LifecycleException("RECEAICMM203");

		lifecycle.fireLifecycleEvent(STOPING_EVENT, this);

		adapterBucketList = new ConcurrentHashMap<>();
		interfaceBucketList = new ConcurrentHashMap<>();
		groupBucketList = new ConcurrentHashMap<>();
		interfaceToGroupMap = new ConcurrentHashMap<>();
		groupVoMap = new ConcurrentHashMap<>();

		started = false;
		lifecycle.fireLifecycleEvent(STOPPED_EVENT, this);
	}

	public void addLifecycleListener(LifecycleListener listener) {
		lifecycle.addLifecycleListener(listener);
	}

	public LifecycleListener[] findLifecycleListeners() {
		return lifecycle.findLifecycleListeners();
	}

	public void removeLifecycleListener(LifecycleListener listener) {
		lifecycle.removeLifecycleListener(listener);
	}

	public boolean isStarted() {
		return this.started;
	}

	public String[] getAdapterAllKeys() {
		Iterator<String> it = this.adapterBucketList.keySet().iterator();
		String[] svcCd = new String[this.adapterBucketList.size()];
		for (int i = 0; it.hasNext(); i++) {
			svcCd[i] = it.next();
		}
		Arrays.sort(svcCd);
		return svcCd;
	}

	public String[] getInterfaceAllKeys() {
		Iterator<String> it = this.interfaceBucketList.keySet().iterator();
		String[] svcCd = new String[this.interfaceBucketList.size()];
		for (int i = 0; it.hasNext(); i++) {
			svcCd[i] = it.next();
		}
		Arrays.sort(svcCd);
		return svcCd;
	}

	public synchronized void removeAdapter(String adapter) {
		adapterBucketList.remove(adapter);
	}

	public synchronized void removeInterface(String inter) {
		interfaceBucketList.remove(inter);
	}

	public synchronized void removeGroup(String groupId) {
		groupBucketList.remove(groupId);
		groupVoMap.remove(groupId);
		Iterator<Map.Entry<String, String>> it = interfaceToGroupMap.entrySet().iterator();
		while (it.hasNext()) {
			Map.Entry<String, String> entry = it.next();
			if (groupId.equals(entry.getValue())) {
				it.remove();
			}
		}
	}

	public String[] getGroupAllKeys() {
		Iterator<String> it = this.groupBucketList.keySet().iterator();
		String[] groupIds = new String[this.groupBucketList.size()];
		for (int i = 0; it.hasNext(); i++) {
			groupIds[i] = it.next();
		}
		Arrays.sort(groupIds);
		return groupIds;
	}

	public CustomGroupBucket getGroupBucket(String groupId) {
		return groupBucketList.get(groupId);
	}

	private CustomBucket makeBucketUsingInflowVO(InflowTargetVO inflowTarget) {
		if (AbstractInflowControlManager.isInflowTarget(inflowTarget)) {
			LocalBucketBuilder builder = Bucket4j.builder();
			if (inflowTarget.getThresholdPerSecond() > 0) {
				builder.addLimit(Bandwidth.classic(inflowTarget.getThresholdPerSecond(),
						Refill.greedy(inflowTarget.getThresholdPerSecond(), Duration.ofSeconds(1))));
			}

			if (inflowTarget.getThreshold() > 0 && StringUtils.isNotBlank(inflowTarget.getThresholdTimeUnit())) {
				builder.addLimit(Bandwidth.classic(inflowTarget.getThreshold(),
						Refill.intervally(inflowTarget.getThreshold(),
								AbstractInflowControlManager.getPeriod(inflowTarget.getThresholdTimeUnit()))));
			}

			return new CustomBucket(builder.build(), inflowTarget);
		}

		return null;
	}

	private CustomGroupBucket makeGroupBucket(InflowGroupVO groupVo) {
		if (!isInflowGroupTarget(groupVo)) {
			return null;
		}

		LocalBucket perSecondBucket = null;
		LocalBucket thresholdBucket = null;

		if (groupVo.getThresholdPerSecond() > 0) {
			perSecondBucket = Bucket4j.builder()
					.addLimit(Bandwidth.classic(groupVo.getThresholdPerSecond(),
							Refill.greedy(groupVo.getThresholdPerSecond(), Duration.ofSeconds(1))))
					.build();
		}

		if (groupVo.getThreshold() > 0 && StringUtils.isNotBlank(groupVo.getThresholdTimeUnit())) {
			thresholdBucket = Bucket4j.builder()
					.addLimit(Bandwidth.classic(groupVo.getThreshold(),
							Refill.intervally(groupVo.getThreshold(),
									AbstractInflowControlManager.getPeriod(groupVo.getThresholdTimeUnit()))))
					.build();
		}

		return new CustomGroupBucket(perSecondBucket, thresholdBucket, groupVo);
	}

	@Override
	public boolean isAdapterPass(String adapter) {
		CustomBucket b = adapterBucketList.get(adapter);

		if (b == null)
			return true;

		return b.getLocalBucket().tryConsume(1);
	}

	@Override
	public boolean isInterfacePass(String inter) {
		CustomBucket b = interfaceBucketList.get(inter);

		if (b == null)
			return true;

		return b.getLocalBucket().tryConsume(1);
	}

	@Override
	public InflowTargetVO getAdapterInflowThreashold(String adapter) {
		CustomBucket b = adapterBucketList.get(adapter);
		if (b == null)
			return null;
		return b.getInflowTargetVo();
	}

	@Override
	public InflowTargetVO getInterfaceInflowThreashold(String inter) {
		CustomBucket b = interfaceBucketList.get(inter);
		if (b == null)
			return null;
		return b.getInflowTargetVo();
	}

	@Override
	public InflowTargetVO getInflowThreashold(String inter) {
		return null;
	}

	public InflowTargetVO getAdapterInflow(String adapter) {
		return getAdapterInflowThreashold(adapter);
	}

	public InflowTargetVO getInterfaceInflow(String inter) {
		return getInterfaceInflowThreashold(inter);
	}

	@Override
	public String isGroupPass(String groupId) {
		CustomGroupBucket b = groupBucketList.get(groupId);

		if (b == null)
			return null;

		return b.tryConsume();
	}

	@Override
	public InflowGroupVO getGroupInflowThreshold(String groupId) {
		return groupVoMap.get(groupId);
	}

	@Override
	public String getGroupIdByInterface(String interfaceId) {
		return interfaceToGroupMap.get(interfaceId);
	}

	public InflowGroupVO getGroupInflow(String groupId) {
		return getGroupInflowThreshold(groupId);
	}

	public static boolean isInflowTarget(InflowTargetVO inflowTarget) {
		if (inflowTarget == null || !inflowTarget.isActivate()) {
			return false;
		}

		return inflowTarget.getThresholdPerSecond() > 0 || (inflowTarget.getThreshold() > 0 && StringUtils
				.equalsAnyIgnoreCase(inflowTarget.getThresholdTimeUnit(), QUOTA_TIMEUNIT_DAY, QUOTA_TIMEUNIT_SECOND,
						QUOTA_TIMEUNIT_MINUTE, QUOTA_TIMEUNIT_HOUR, QUOTA_TIMEUNIT_DAY, QUOTA_TIMEUNIT_MONTH));

	}

	public static boolean isInflowGroupTarget(InflowGroupVO groupVo) {
		if (groupVo == null || !groupVo.isActivate()) {
			return false;
		}

		return groupVo.getThresholdPerSecond() > 0 || (groupVo.getThreshold() > 0 && StringUtils
				.equalsAnyIgnoreCase(groupVo.getThresholdTimeUnit(), QUOTA_TIMEUNIT_DAY, QUOTA_TIMEUNIT_SECOND,
						QUOTA_TIMEUNIT_MINUTE, QUOTA_TIMEUNIT_HOUR, QUOTA_TIMEUNIT_DAY, QUOTA_TIMEUNIT_MONTH));
	}

	public static Duration getPeriod(String timeUnit) {
		if (StringUtils.equalsIgnoreCase(timeUnit, QUOTA_TIMEUNIT_SECOND)) {
			return Duration.ofSeconds(1);
		} else if (StringUtils.equalsIgnoreCase(timeUnit, QUOTA_TIMEUNIT_MINUTE)) {
			return Duration.ofMinutes(1);
		} else if (StringUtils.equalsIgnoreCase(timeUnit, QUOTA_TIMEUNIT_HOUR)) {
			return Duration.ofHours(1);
		} else if (StringUtils.equalsIgnoreCase(timeUnit, QUOTA_TIMEUNIT_DAY)) {
			return Duration.ofDays(1);
		} else if (StringUtils.equalsIgnoreCase(timeUnit, QUOTA_TIMEUNIT_MONTH)) {
			Calendar calendar = Calendar.getInstance();
			return Duration.ofDays(calendar.getActualMaximum(Calendar.DAY_OF_MONTH));
		} else {
			return null;
		}
	}
}