001/**
002 * Copyright (c) 2015-2022, Michael Yang 杨福海 (fuhai999@gmail.com).
003 * <p>
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 * <p>
008 * http://www.apache.org/licenses/LICENSE-2.0
009 * <p>
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package io.jboot.components.limiter.interceptor;
017
018
019import com.google.common.util.concurrent.RateLimiter;
020import com.jfinal.aop.Invocation;
021import io.jboot.components.limiter.LimiterManager;
022import io.jboot.components.limiter.redis.RedisRateLimitUtil;
023import io.jboot.utils.ClassUtil;
024import io.jboot.utils.StrUtil;
025
026import javax.servlet.http.HttpServletRequest;
027import java.util.concurrent.Semaphore;
028
029public abstract class BaseLimiterInterceptor {
030
031
032    protected void doInterceptForConcurrency(int rate, String resource, String fallback, Invocation inv) {
033        Semaphore semaphore = LimiterManager.me().getOrCreateSemaphore(resource, rate);
034        boolean acquire = false;
035        try {
036            acquire = semaphore.tryAcquire();
037            if (acquire) {
038                inv.invoke();
039            }
040            //不允许通行
041            else {
042                doExecFallback(resource, fallback, inv);
043            }
044        } finally {
045            if (acquire) {
046                semaphore.release();
047            }
048        }
049    }
050
051
052    protected void doInterceptForTokenBucket(int rate, String resource, String fallback, Invocation inv) {
053        RateLimiter limiter = LimiterManager.me().getOrCreateRateLimiter(resource, rate);
054        //允许通行
055        if (limiter.tryAcquire()) {
056            inv.invoke();
057        }
058        //不允许通行
059        else {
060            doExecFallback(resource, fallback, inv);
061        }
062    }
063
064    protected void doInterceptForTokenBucketWithCluster(int rate, String resource, String fallback, Invocation inv) {
065        //允许通行
066        if (RedisRateLimitUtil.tryAcquire(resource, rate)) {
067            inv.invoke();
068        }
069        //不允许通行
070        else {
071            doExecFallback(resource, fallback, inv);
072        }
073    }
074
075    protected void doExecFallback(String resource, String fallback, Invocation inv) {
076        LimiterManager.me().processFallback(resource, fallback, inv);
077    }
078
079
080    protected String getPackageOrTarget(Invocation inv) {
081        return inv.isActionInvocation() ? buildUrl(inv) : ClassUtil.buildMethodString(inv.getMethod());
082    }
083
084    protected String buildUrl(Invocation inv) {
085        HttpServletRequest request = inv.getController().getRequest();
086        String uri = request.getRequestURI();
087        String query = request.getQueryString();
088        return StrUtil.isBlank(query) ? uri : uri + "?" + query;
089    }
090
091}