|
3 | 3 | import com.google.common.base.Strings; |
4 | 4 | import io.prometheus.client.Histogram; |
5 | 5 | import java.io.IOException; |
6 | | -import java.lang.reflect.Constructor; |
| 6 | +import java.util.Arrays; |
| 7 | +import java.util.Collections; |
| 8 | +import java.util.HashMap; |
| 9 | +import java.util.List; |
| 10 | +import java.util.Map; |
7 | 11 | import javax.annotation.PostConstruct; |
8 | 12 | import javax.servlet.ServletException; |
9 | 13 | import javax.servlet.http.HttpServlet; |
|
31 | 35 | @Slf4j |
32 | 36 | public abstract class RateLimiterServlet extends HttpServlet { |
33 | 37 | private static final String KEY_PREFIX_HTTP = "http_"; |
34 | | - private static final String ADAPTER_PREFIX = "org.tron.core.services.ratelimiter.adapter."; |
| 38 | + |
| 39 | + static final Map<String, Class<? extends IRateLimiter>> ALLOWED_ADAPTERS; |
| 40 | + static final String DEFAULT_ADAPTER_NAME = DefaultBaseQqsAdapter.class.getSimpleName(); |
| 41 | + |
| 42 | + static { |
| 43 | + List<Class<? extends IRateLimiter>> adapters = Arrays.asList( |
| 44 | + GlobalPreemptibleAdapter.class, |
| 45 | + QpsRateLimiterAdapter.class, |
| 46 | + IPQPSRateLimiterAdapter.class, |
| 47 | + DefaultBaseQqsAdapter.class); |
| 48 | + Map<String, Class<? extends IRateLimiter>> m = new HashMap<>(); |
| 49 | + for (Class<? extends IRateLimiter> c : adapters) { |
| 50 | + m.put(c.getSimpleName(), c); |
| 51 | + } |
| 52 | + ALLOWED_ADAPTERS = Collections.unmodifiableMap(m); |
| 53 | + } |
35 | 54 |
|
36 | 55 | @Autowired |
37 | 56 | private RateLimiterContainer container; |
38 | 57 |
|
39 | 58 | @PostConstruct |
40 | 59 | private void addRateContainer() { |
41 | | - RateLimiterInitialization.HttpRateLimiterItem item = Args.getInstance() |
42 | | - .getRateLimiterInitialization().getHttpMap().get(getClass().getSimpleName()); |
43 | | - boolean success = false; |
44 | 60 | final String name = getClass().getSimpleName(); |
45 | | - if (item != null) { |
46 | | - String cName = ""; |
47 | | - String params = ""; |
48 | | - Object obj; |
49 | | - try { |
50 | | - cName = item.getStrategy(); |
51 | | - params = item.getParams(); |
52 | | - // add the specific rate limiter strategy of servlet. |
53 | | - Class<?> c = Class.forName(ADAPTER_PREFIX + cName); |
54 | | - Constructor constructor; |
55 | | - if (c == GlobalPreemptibleAdapter.class || c == QpsRateLimiterAdapter.class |
56 | | - || c == IPQPSRateLimiterAdapter.class) { |
57 | | - constructor = c.getConstructor(String.class); |
58 | | - obj = constructor.newInstance(params); |
59 | | - container.add(KEY_PREFIX_HTTP, name, (IRateLimiter) obj); |
60 | | - } else { |
61 | | - constructor = c.getConstructor(); |
62 | | - obj = constructor.newInstance(QpsStrategy.DEFAULT_QPS_PARAM); |
63 | | - container.add(KEY_PREFIX_HTTP, name, (IRateLimiter) obj); |
64 | | - } |
65 | | - success = true; |
66 | | - } catch (Exception e) { |
67 | | - this.throwTronError(cName, params, name, e); |
68 | | - } |
| 61 | + RateLimiterInitialization.HttpRateLimiterItem item = Args.getInstance() |
| 62 | + .getRateLimiterInitialization().getHttpMap().get(name); |
| 63 | + |
| 64 | + String cName; |
| 65 | + String params; |
| 66 | + if (item == null) { |
| 67 | + cName = DEFAULT_ADAPTER_NAME; |
| 68 | + params = QpsStrategy.DEFAULT_QPS_PARAM; |
| 69 | + } else { |
| 70 | + cName = item.getStrategy(); |
| 71 | + params = item.getParams(); |
69 | 72 | } |
70 | | - if (!success) { |
71 | | - // if the specific rate limiter strategy of servlet is not defined or fail to add, |
72 | | - // then add a default Strategy. |
73 | | - try { |
74 | | - IRateLimiter rateLimiter = new DefaultBaseQqsAdapter(QpsStrategy.DEFAULT_QPS_PARAM); |
75 | | - container.add(KEY_PREFIX_HTTP, name, rateLimiter); |
76 | | - } catch (Exception e) { |
77 | | - this.throwTronError("DefaultBaseQqsAdapter", QpsStrategy.DEFAULT_QPS_PARAM, name, e); |
78 | | - } |
| 73 | + |
| 74 | + try { |
| 75 | + container.add(KEY_PREFIX_HTTP, name, buildAdapter(cName, params, name)); |
| 76 | + } catch (Exception e) { |
| 77 | + throw rateLimiterInitError(cName, params, name, e); |
| 78 | + } |
| 79 | + } |
| 80 | + |
| 81 | + static IRateLimiter buildAdapter(String cName, String params, String name) { |
| 82 | + Class<? extends IRateLimiter> c = ALLOWED_ADAPTERS.get(cName); |
| 83 | + if (c == null) { |
| 84 | + throw rateLimiterInitError(cName, params, name, |
| 85 | + new IllegalArgumentException("unknown rate limiter adapter; allowed=" |
| 86 | + + ALLOWED_ADAPTERS.keySet())); |
| 87 | + } |
| 88 | + try { |
| 89 | + return c.getConstructor(String.class).newInstance(params); |
| 90 | + } catch (Exception e) { |
| 91 | + throw rateLimiterInitError(cName, params, name, e); |
79 | 92 | } |
80 | 93 | } |
81 | 94 |
|
82 | | - private void throwTronError(String strategy, String params, String servlet, Exception e) { |
83 | | - throw new TronError("failure to add the rate limiter strategy. servlet = " + servlet |
| 95 | + private static TronError rateLimiterInitError(String strategy, String params, String servlet, |
| 96 | + Exception e) { |
| 97 | + return new TronError("failure to add the rate limiter strategy. servlet = " + servlet |
84 | 98 | + ", strategy name = " + strategy + ", params = \"" + params + "\".", |
85 | 99 | e, TronError.ErrCode.RATE_LIMITER_INIT); |
86 | 100 | } |
|
0 commit comments