com.github.javaclub.sword.spring.MethodInvokeLimitedInterceptor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of sword Show documentation
Javaclub sword
The newest version!
/*
 * @(#)MethodInvokeLimitedInterceptor.java	2017年5月27日
 *
 * Copyright (c) 2017. All Rights Reserved.
 *
 */

package com.github.javaclub.sword.spring;

import static com.github.javaclub.sword.core.Strings.concat;
import static com.github.javaclub.sword.core.Strings.isEmpty;

import java.util.Objects;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.javaclub.sword.InvokeLimitedException;
import com.github.javaclub.sword.annotation.InvokeLimited;
import com.github.javaclub.sword.annotation.InvokeLimited.InvokeLimitedType;
import com.github.javaclub.sword.cache.CacheManager;
import com.github.javaclub.sword.component.SpringContextUtil;
import com.github.javaclub.sword.core.Strings;
import com.github.javaclub.sword.domain.ResultDO;
import com.github.javaclub.sword.domain.enumtype.BasicMessage;
import com.github.javaclub.sword.util.AnnotationUtil;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;


/**
 * MethodInvokeLimitedInterceptor 方法调用QPS超限拦截
 *
 * @author Gerald Chen
 * @version $Id: MethodInvokeLimitedInterceptor.java 2017年5月27日 18:02:50 Exp $
 */
public class MethodInvokeLimitedInterceptor implements MethodInterceptor {
	
	static final Logger log = LoggerFactory.getLogger(MethodInvokeLimitedInterceptor.class);
	
	static volatile CacheManager cacheManager;
	
	static volatile LoadingCache counter =
	        CacheBuilder.newBuilder()
	                .expireAfterWrite(2, TimeUnit.SECONDS)
	                .build(new CacheLoader() {
	                    @Override
	                    public AtomicLong load(String methodAndSeconds) throws Exception {
	                        return new AtomicLong(0);
	                    }
	                });
	
	@SuppressWarnings("unchecked")
	@Override
	public Object invoke(MethodInvocation mi) throws Throwable {
		
        // 判断该方法是否加了 @InvokeLimited 注解
        if (!mi.getMethod().isAnnotationPresent(InvokeLimited.class)) {
        		return mi.proceed();
        }
        InvokeLimited limited = AnnotationUtil.getAnnotation(mi.getMethod(), InvokeLimited.class);
        InvokeLimitedType type  = limited.type();
        int max = limited.limit(); // 最大调用限制数/s
        
        String methodKey = this.getClassAndMethodName(mi);
        boolean isLimited = false;
        if (InvokeLimitedType.LOCAL.equals(type)) { // 单机
	        	// 得到当前秒
	    	    long currentSeconds = System.currentTimeMillis() / 1000;
	    	    String key = Strings.concat(methodKey, "#", currentSeconds);
	    	    try {
	    			if (counter.get(key).incrementAndGet() > max) {
	    				isLimited = true;
	    			}
	    		} catch (ExecutionException e) {
	    			// 如果真的抛出了这个异常，让方法正常执行(例外)
	    			log.error("", e);
	    			return mi.proceed();
	    		}
        } else if (InvokeLimitedType.DISTRIBUTED.equals(type)) { // 分布式
	        	try {
				if(null == cacheManager) {
					cacheManager = (CacheManager) SpringContextUtil.getBean(CacheManager.class);
				}
			} catch (Exception e) {
				log.error("", e);
			}
	        	if (null == cacheManager) {
	        		return mi.proceed(); // 正常执行
	        	}
	        	long currentSeconds = System.currentTimeMillis() / 1000;
	    	    String key = Strings.concat(methodKey, "#", currentSeconds);
	    	    // long incredVal = cacheManager.incr(key);
	    	    long incredVal = 1;
	    	    if (incredVal == 1) {
	    	    		cacheManager.expire(key, 2); // 2秒后过期
	    	    }
	        	if (log.isInfoEnabled()) {
	        		log.info("key={}\t incredVal={}", key, incredVal);
	        	}
	        	if (incredVal > max) {
	        		isLimited = true;
	        	}
        }
        
        Class returnType = mi.getMethod().getReturnType();
        if (isLimited) {
	        log.warn("method={}\t message=Method Call is Limited by max {}/s", methodKey, max);
	        if(Objects.equals(returnType, ResultDO.class)) {
	            ResultDO ret = ResultDO.failure(Strings.format("[{}] Method Call is Limited by max QPS={}", methodKey, max));
	            ret.setCode(BasicMessage.METHOD_CALL_IS_LIMITED.getCode());
	            return ret;
	        }
	        throw new InvokeLimitedException(Strings.format("[{}] Method Call is Limited by max QPS={}", methodKey, max));
        }
        
        Object result = null;
		try {
			result = mi.proceed();
		} catch (Exception e) {
			log.error("", e);
			throw e;
		}
		
        return result;
	}
	
	String getClassAndMethodName(MethodInvocation mi) { 
		String className = "";
		try {
			className = mi.getThis().getClass().getName();
		} catch (Exception e) {
		}
		return isEmpty(className) ? mi.getMethod().getName() : concat(className, ".", mi.getMethod().getName());
	}
	
	public static void main(String[] args) {
		ResultDO success = ResultDO.success();
		ResultDO failure = ResultDO.failure();
		
		System.out.println(Objects.equals(success.getClass(), failure.getClass()));
	}

}