fix: add limit chain

core/src/main/java/org/apache/seata/core/protocol/MessageType.java

@@ -22,6 +22,11 @@
  */
 public interface MessageType {
 
+    /**
+     * The constant TYPE_NOT_EXIST.
+     */
+    short TYPE_NOT_EXIST = 0;
+
     /**
      * The constant TYPE_GLOBAL_BEGIN.
      */

server/src/main/java/org/apache/seata/server/coordinator/DefaultCoordinator.java

@@ -26,8 +26,8 @@
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
-
 import io.netty.channel.Channel;
+
 import org.apache.commons.lang.time.DateFormatUtils;
 import org.apache.seata.common.DefaultValues;
 import org.apache.seata.common.thread.NamedThreadFactory;
@@ -65,8 +65,8 @@
 import org.apache.seata.core.rpc.netty.ChannelManager;
 import org.apache.seata.core.rpc.netty.NettyRemotingServer;
 import org.apache.seata.server.AbstractTCInboundHandler;
+import org.apache.seata.server.limit.LimitRequestDecorator;
 import org.apache.seata.server.metrics.MetricsPublisher;
-import org.apache.seata.server.ratelimit.RateLimiterHandler;
 import org.apache.seata.server.session.BranchSession;
 import org.apache.seata.server.session.GlobalSession;
 import org.apache.seata.server.session.SessionCondition;
@@ -197,8 +197,6 @@ public class DefaultCoordinator extends AbstractTCInboundHandler implements Tran
 
     private final ThreadPoolExecutor branchRemoveExecutor;
 
-    private RateLimiterHandler rateLimiterHandler;
-
     private RemotingServer remotingServer;
 
     private final DefaultCore core;
@@ -229,8 +227,6 @@ protected DefaultCoordinator(RemotingServer remotingServer) {
         } else {
             branchRemoveExecutor = null;
         }
-        // create server rate limter
-        rateLimiterHandler = RateLimiterHandler.getInstance();
     }
 
     public static DefaultCoordinator getInstance(RemotingServer remotingServer) {
@@ -646,11 +642,9 @@ public AbstractResultMessage onRequest(AbstractMessage request, RpcContext conte
         }
         AbstractTransactionRequestToTC transactionRequest = (AbstractTransactionRequestToTC) request;
         transactionRequest.setTCInboundHandler(this);
-        AbstractResultMessage resultMessage = processRateLimit(request, context);
-        if (resultMessage != null) {
-            return resultMessage;
-        }
-        return transactionRequest.handle(context);
+
+        LimitRequestDecorator limitRequestDecorator = new LimitRequestDecorator(transactionRequest);
+        return limitRequestDecorator.handle(context);
     }
 
     @Override
@@ -770,8 +764,4 @@ private void doRemove(BranchSession bt) {
             }
         }
     }
-
-    private AbstractResultMessage processRateLimit(AbstractMessage request, RpcContext context) {
-        return rateLimiterHandler.handle(request, context);
-    }
 }

server/src/main/java/org/apache/seata/server/limit/LimitRequestDecorator.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.seata.server.limit;
+
+import org.apache.seata.core.protocol.transaction.AbstractTransactionRequestToTC;
+import org.apache.seata.core.protocol.transaction.AbstractTransactionResponse;
+import org.apache.seata.core.rpc.RpcContext;
+import org.apache.seata.server.limit.ratelimit.RateLimiterHandler;
+
+/**
+ * LimitRequestDecorator decorate AbstractTransactionRequestToTC to use limiter
+ */
+public class LimitRequestDecorator extends AbstractTransactionRequestToTC {
+
+    private AbstractTransactionRequestToTC originalRequest;
+
+    private TransactionRequestLimitHandler requestLimitHandler;
+
+    public LimitRequestDecorator(AbstractTransactionRequestToTC originalRequest) {
+        this.originalRequest = originalRequest;
+
+        // create server rate limter
+        RateLimiterHandler rateLimiterHandler = RateLimiterHandler.getInstance();
+        rateLimiterHandler.setTransactionRequestLimitHandler(null);
+        requestLimitHandler = rateLimiterHandler;
+    }
+
+
+    @Override
+    public AbstractTransactionResponse handle(RpcContext rpcContext) {
+        return requestLimitHandler.handle(originalRequest, rpcContext);
+    }
+
+    @Override
+    public short getTypeCode() {
+        return originalRequest.getTypeCode();
+    }
+}

server/src/main/java/org/apache/seata/server/limit/TransactionRequestLimitHandler.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.seata.server.limit;
+
+import org.apache.seata.core.protocol.transaction.AbstractTransactionRequestToTC;
+import org.apache.seata.core.protocol.transaction.AbstractTransactionResponse;
+import org.apache.seata.core.rpc.RpcContext;
+
+/**
+ * TransactionRequestLimitHandler
+ */
+public abstract class TransactionRequestLimitHandler {
+
+    /**
+     * limit handler
+     */
+    protected TransactionRequestLimitHandler transactionRequestLimitHandler;
+
+    public TransactionRequestLimitHandler() {
+    }
+
+    /**
+     * next handler handle
+     * @param context
+     * @return
+     */
+    protected AbstractTransactionResponse next(AbstractTransactionRequestToTC originRequest, RpcContext context) {
+        if (transactionRequestLimitHandler != null) {
+            return transactionRequestLimitHandler.next(originRequest, context);
+        }
+        return originRequest.handle(context);
+    }
+
+    public abstract AbstractTransactionResponse handle(AbstractTransactionRequestToTC originRequest, RpcContext context);
+
+    public void setTransactionRequestLimitHandler(TransactionRequestLimitHandler transactionRequestLimitHandler) {
+        this.transactionRequestLimitHandler = transactionRequestLimitHandler;
+    }
+}

server/src/main/java/org/apache/seata/server/ratelimit/RateLimitInfo.java → server/src/main/java/org/apache/seata/server/limit/ratelimit/RateLimitInfo.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.server.ratelimit;
+package org.apache.seata.server.limit.ratelimit;
 
 import org.apache.seata.common.util.UUIDGenerator;
 

server/src/main/java/org/apache/seata/server/ratelimit/RateLimiter.java → server/src/main/java/org/apache/seata/server/limit/ratelimit/RateLimiter.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.server.ratelimit;
+package org.apache.seata.server.limit.ratelimit;
 
 /**
  * RateLimiter

server/src/main/java/org/apache/seata/server/ratelimit/RateLimiterHandler.java → server/src/main/java/org/apache/seata/server/limit/ratelimit/RateLimiterHandler.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.server.ratelimit;
+package org.apache.seata.server.limit.ratelimit;
 
 import org.apache.seata.common.XID;
 import org.apache.seata.common.loader.EnhancedServiceLoader;
@@ -25,18 +25,19 @@
 import org.apache.seata.common.ConfigurationKeys;
 import org.apache.seata.config.ConfigurationFactory;
 import org.apache.seata.core.exception.TransactionExceptionCode;
-import org.apache.seata.core.protocol.AbstractMessage;
-import org.apache.seata.core.protocol.AbstractResultMessage;
+import org.apache.seata.core.protocol.MessageType;
 import org.apache.seata.core.protocol.ResultCode;
-import org.apache.seata.core.protocol.transaction.GlobalBeginRequest;
+import org.apache.seata.core.protocol.transaction.AbstractTransactionRequestToTC;
+import org.apache.seata.core.protocol.transaction.AbstractTransactionResponse;
 import org.apache.seata.core.protocol.transaction.GlobalBeginResponse;
 import org.apache.seata.core.rpc.RpcContext;
+import org.apache.seata.server.limit.TransactionRequestLimitHandler;
 import org.apache.seata.server.metrics.MetricsPublisher;
 
 /**
  * RateLimiterHandler
  */
-public class RateLimiterHandler implements CachedConfigurationChangeListener {
+public class RateLimiterHandler extends TransactionRequestLimitHandler implements CachedConfigurationChangeListener {
     /**
      * The instance of RateLimiterHandler
      */
@@ -68,6 +69,27 @@ private RateLimiterHandler() {
         config.addConfigListener(ConfigurationKeys.RATE_LIMIT_BUCKET_TOKEN_INITIAL_NUM, this);
     }
 
+    @Override
+    public AbstractTransactionResponse handle(AbstractTransactionRequestToTC originRequest, RpcContext context) {
+        if (!rateLimiter.isEnable()) {
+            return next(originRequest, context);
+        }
+
+        if (MessageType.TYPE_GLOBAL_BEGIN == originRequest.getTypeCode()) {
+            if (!rateLimiter.canPass()) {
+                GlobalBeginResponse response = new GlobalBeginResponse();
+                response.setTransactionExceptionCode(TransactionExceptionCode.BeginFailed);
+                response.setResultCode(ResultCode.Failed);
+                RateLimitInfo rateLimitInfo = RateLimitInfo.generateRateLimitInfo(context.getApplicationId(),
+                        RateLimitInfo.GLOBAL_BEGIN_FAILED, context.getClientId(), XID.getIpAddressAndPort());
+                MetricsPublisher.postRateLimitEvent(rateLimitInfo);
+                response.setMsg(String.format("TransactionException[rate limit exception, rate limit info: %s]", rateLimitInfo));
+                return response;
+            }
+        }
+        return next(originRequest, context);
+    }
+
     public static RateLimiterHandler getInstance() {
         if (instance == null) {
             synchronized (RateLimiterHandler.class) {
@@ -97,24 +119,4 @@ public void onChangeEvent(ConfigurationChangeEvent event) {
         }
         rateLimiter.reInit(config);
     }
-
-    public AbstractResultMessage handle(AbstractMessage request, RpcContext rpcContext) {
-        if (!rateLimiter.isEnable()) {
-            return null;
-        }
-
-        if (request instanceof GlobalBeginRequest) {
-            if (!rateLimiter.canPass()) {
-                GlobalBeginResponse response = new GlobalBeginResponse();
-                response.setTransactionExceptionCode(TransactionExceptionCode.BeginFailed);
-                response.setResultCode(ResultCode.Failed);
-                RateLimitInfo rateLimitInfo = RateLimitInfo.generateRateLimitInfo(rpcContext.getApplicationId(),
-                        RateLimitInfo.GLOBAL_BEGIN_FAILED, rpcContext.getClientId(), XID.getIpAddressAndPort());
-                MetricsPublisher.postRateLimitEvent(rateLimitInfo);
-                response.setMsg(String.format("TransactionException[rate limit exception, rate limit info: %s]", rateLimitInfo));
-                return response;
-            }
-        }
-        return null;
-    }
 }

server/src/main/java/org/apache/seata/server/ratelimit/RateLimiterHandlerConfig.java → server/src/main/java/org/apache/seata/server/limit/ratelimit/RateLimiterHandlerConfig.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.server.ratelimit;
+package org.apache.seata.server.limit.ratelimit;
 
 /**
  * RateLimiterHandlerConfig

server/src/main/java/org/apache/seata/server/ratelimit/TokenBucketLimiter.java → server/src/main/java/org/apache/seata/server/limit/ratelimit/TokenBucketLimiter.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.seata.server.ratelimit;
+package org.apache.seata.server.limit.ratelimit;
 
 import org.apache.seata.common.ConfigurationKeys;
 import org.apache.seata.common.executor.Initialize;

server/src/main/java/org/apache/seata/server/metrics/MetricsPublisher.java

@@ -21,7 +21,7 @@
 import org.apache.seata.core.event.RateLimitEvent;
 import org.apache.seata.core.model.GlobalStatus;
 import org.apache.seata.server.event.EventBusManager;
-import org.apache.seata.server.ratelimit.RateLimitInfo;
+import org.apache.seata.server.limit.ratelimit.RateLimitInfo;
 import org.apache.seata.server.session.GlobalSession;
 
 /**

server/src/main/resources/META-INF/services/org.apache.seata.server.ratelimit.RateLimiter → server/src/main/resources/META-INF/services/org.apache.seata.server.limit.ratelimit.RateLimiter

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-org.apache.seata.server.ratelimit.TokenBucketLimiter
+org.apache.seata.server.limit.ratelimit.TokenBucketLimiter