Merge pull request #5093 from wubin01/rate-limit

feat(net): add global rate limiter function

Author: xxo1shine

common/src/main/java/org/tron/common/parameter/CommonParameter.java

@@ -398,6 +398,12 @@ public class CommonParameter {
   @Setter
   public RateLimiterInitialization rateLimiterInitialization;
   @Getter
+  @Setter
+  public int rateLimiterGlobalQps;
+  @Getter
+  @Setter
+  public int rateLimiterGlobalIpQps;
+  @Getter
   public DbBackupConfig dbBackupConfig;
   @Getter
   public RocksDbSettings rocksDBCustomSettings;

common/src/main/java/org/tron/core/Constant.java

@@ -235,6 +235,10 @@ public class Constant {
 
   public static final String RATE_LIMITER = "rate.limiter";
 
+  public static final String RATE_LIMITER_GLOBAL_QPS = "rate.limiter.global.qps";
+
+  public static final String RATE_LIMITER_GLOBAL_IP_QPS = "rate.limiter.global.ip.qps";
+
   public static final String COMMITTEE_CHANGED_DELEGATION = "committee.changedDelegation";
 
   public static final String CRYPTO_ENGINE = "crypto.engine";

framework/src/main/java/org/tron/core/config/args/Args.java

@@ -221,6 +221,8 @@ public static void clearParam() {
     PARAMETER.allowNewRewardAlgorithm = 0;
     PARAMETER.allowNewReward = 0;
     PARAMETER.memoFee = 0;
+    PARAMETER.rateLimiterGlobalQps = 50000;
+    PARAMETER.rateLimiterGlobalIpQps = 10000;
     PARAMETER.p2pDisable = false;
   }
 
@@ -928,6 +930,14 @@ public static void setParam(final String[] args, final String confFileName) {
       PARAMETER.fullNodeAllowShieldedTransactionArgs = true;
     }
 
+    PARAMETER.rateLimiterGlobalQps =
+        config.hasPath(Constant.RATE_LIMITER_GLOBAL_QPS) ? config
+            .getInt(Constant.RATE_LIMITER_GLOBAL_QPS) : 50000;
+
+    PARAMETER.rateLimiterGlobalIpQps =
+        config.hasPath(Constant.RATE_LIMITER_GLOBAL_IP_QPS) ? config
+            .getInt(Constant.RATE_LIMITER_GLOBAL_IP_QPS) : 10000;
+
     PARAMETER.rateLimiterInitialization =
         config.hasPath(Constant.RATE_LIMITER) ? getRateLimiterFromConfig(config)
             : new RateLimiterInitialization();
@@ -1480,6 +1490,8 @@ public static void logConfig() {
     logger.info("Max connection with same IP: {}", parameter.getMaxConnectionsWithSameIp());
     logger.info("Solidity threads: {}", parameter.getSolidityThreads());
     logger.info("Trx reference block: {}", parameter.getTrxReferenceBlock());
+    logger.info("Rate limiter global qps: {}", parameter.getRateLimiterGlobalQps());
+    logger.info("Rate limiter global ip qps: {}", parameter.getRateLimiterGlobalIpQps());
     logger.info("************************ Backup config ************************");
     logger.info("Backup priority: {}", parameter.getBackupPriority());
     logger.info("Backup listen port: {}", parameter.getBackupPort());

framework/src/main/java/org/tron/core/services/http/RateLimiterServlet.java

@@ -16,6 +16,7 @@
 import org.tron.common.prometheus.MetricLabels;
 import org.tron.common.prometheus.Metrics;
 import org.tron.core.config.args.Args;
+import org.tron.core.services.ratelimiter.GlobalRateLimiter;
 import org.tron.core.services.ratelimiter.RateLimiterContainer;
 import org.tron.core.services.ratelimiter.RuntimeData;
 import org.tron.core.services.ratelimiter.adapter.DefaultBaseQqsAdapter;
@@ -91,12 +92,16 @@ private void addRateContainer() {
   @Override
   protected void service(HttpServletRequest req, HttpServletResponse resp)
       throws ServletException, IOException {
+    
+    RuntimeData runtimeData = new RuntimeData(req);
+    GlobalRateLimiter.acquire(runtimeData);
+
     IRateLimiter rateLimiter = container.get(KEY_PREFIX_HTTP, getClass().getSimpleName());
 
     boolean acquireResource = true;
 
     if (rateLimiter != null) {
-      acquireResource = rateLimiter.acquire(new RuntimeData(req));
+      acquireResource = rateLimiter.acquire(runtimeData);
     }
     String url = Strings.isNullOrEmpty(req.getRequestURI())
         ? MetricLabels.UNDEFINED : req.getRequestURI();

framework/src/main/java/org/tron/core/services/ratelimiter/GlobalRateLimiter.java

@@ -0,0 +1,35 @@
+package org.tron.core.services.ratelimiter;
+
+import com.google.common.base.Strings;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.RateLimiter;
+import java.util.concurrent.TimeUnit;
+import org.tron.core.config.args.Args;
+
+public class GlobalRateLimiter {
+
+  private static double QPS = Args.getInstance().getRateLimiterGlobalQps();
+
+  private static double IP_QPS = Args.getInstance().getRateLimiterGlobalIpQps();
+
+  private static Cache<String, RateLimiter> cache = CacheBuilder.newBuilder()
+      .maximumSize(10000).expireAfterWrite(1, TimeUnit.HOURS).build();
+
+  private static RateLimiter rateLimiter = RateLimiter.create(QPS);
+
+  public static void acquire(RuntimeData runtimeData) {
+    rateLimiter.acquire();
+    String ip = runtimeData.getRemoteAddr();
+    if (Strings.isNullOrEmpty(ip)) {
+      return;
+    }
+    RateLimiter r = cache.getIfPresent(ip);
+    if (r == null) {
+      r = RateLimiter.create(IP_QPS);
+      cache.put(ip, r);
+    }
+    r.acquire();
+  }
+
+}

framework/src/main/java/org/tron/core/services/ratelimiter/RateLimiterInterceptor.java

@@ -107,10 +107,13 @@ public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call,
     IRateLimiter rateLimiter = container
         .get(KEY_PREFIX_RPC, call.getMethodDescriptor().getFullMethodName());
 
+    RuntimeData runtimeData = new RuntimeData(call);
+    GlobalRateLimiter.acquire(runtimeData);
+
     boolean acquireResource = true;
 
     if (rateLimiter != null) {
-      acquireResource = rateLimiter.acquire(new RuntimeData(call));
+      acquireResource = rateLimiter.acquire(runtimeData);
     }
 
     Listener<ReqT> listener = new ServerCall.Listener<ReqT>() {

framework/src/test/java/org/tron/common/config/args/ArgsTest.java

@@ -30,6 +30,8 @@ public void testConfig() {
     Assert.assertEquals(Args.getInstance().getNodeDiscoveryPingTimeout(), 15_000);
     Assert.assertEquals(Args.getInstance().getMaxFastForwardNum(), 3);
     Assert.assertEquals(Args.getInstance().getBlockCacheTimeout(), 60);
+    Assert.assertEquals(Args.getInstance().getRateLimiterGlobalQps(), 50000);
+    Assert.assertEquals(Args.getInstance().getRateLimiterGlobalIpQps(), 10000);
     Assert.assertEquals(Args.getInstance().p2pDisable, true);
   }
 }

framework/src/test/java/org/tron/core/services/ratelimiter/GlobalRateLimiterTest.java

@@ -0,0 +1,22 @@
+package org.tron.core.services.ratelimiter;
+
+import java.lang.reflect.Field;
+import org.junit.Assert;
+import org.junit.Test;
+import org.tron.core.Constant;
+import org.tron.core.config.args.Args;
+
+public class GlobalRateLimiterTest {
+
+  @Test
+  public void testAcquire() throws Exception {
+    String[] a = new String[0];
+    Args.setParam(a, Constant.TESTNET_CONF);
+    RuntimeData runtimeData = new RuntimeData(null);
+    Field field =  runtimeData.getClass().getDeclaredField("address");
+    field.setAccessible(true);
+    field.set(runtimeData, "127.0.0.1");
+    Assert.assertEquals(runtimeData.getRemoteAddr(), "127.0.0.1");
+    GlobalRateLimiter.acquire(runtimeData);
+  }
+}