[ISSUE apache#7945] Make HAProxyMessageForwarder Scalable (apache#7946)

Jianing-Xu · Mar 19, 2024 · 1b63e8e · 1b63e8e
1 parent 9b4f16a
commit 1b63e8e
Showing 1 changed file with 25 additions and 8 deletions.
diff --git a/.../java/org/apache/rocketmq/proxy/remoting/protocol/http2proxy/HAProxyMessageForwarder.java b/.../java/org/apache/rocketmq/proxy/remoting/protocol/http2proxy/HAProxyMessageForwarder.java
@@ -89,8 +89,6 @@ private void forwardHAProxyMessage(Channel inboundChannel, Channel outboundChann
     protected HAProxyMessage buildHAProxyMessage(Channel inboundChannel) throws IllegalAccessException, DecoderException {
         String sourceAddress = null, destinationAddress = null;
         int sourcePort = 0, destinationPort = 0;
-        List<HAProxyTLV> haProxyTLVs = new ArrayList<>();
-
         if (inboundChannel.hasAttr(AttributeKeys.PROXY_PROTOCOL_ADDR)) {
             Attribute<?>[] attributes = (Attribute<?>[]) FieldUtils.readField(FIELD_ATTRIBUTE, inboundChannel);
             if (ArrayUtils.isEmpty(attributes)) {
@@ -117,12 +115,6 @@ protected HAProxyMessage buildHAProxyMessage(Channel inboundChannel) throws Ille
                 if (attribute.key() == AttributeKeys.PROXY_PROTOCOL_SERVER_PORT) {
                     destinationPort = Integer.parseInt(attributeValue);
                 }
-                if (StringUtils.startsWith(attributeKey, HAProxyConstants.PROXY_PROTOCOL_TLV_PREFIX)) {
-                    HAProxyTLV haProxyTLV = buildHAProxyTLV(attributeKey, attributeValue);
-                    if (haProxyTLV != null) {
-                        haProxyTLVs.add(haProxyTLV);
-                    }
-                }
             }
         } else {
             String remoteAddr = RemotingHelper.parseChannelRemoteAddr(inboundChannel);
@@ -137,10 +129,35 @@ protected HAProxyMessage buildHAProxyMessage(Channel inboundChannel) throws Ille
         HAProxyProxiedProtocol proxiedProtocol = AclUtils.isColon(sourceAddress) ? HAProxyProxiedProtocol.TCP6 :
             HAProxyProxiedProtocol.TCP4;
 
+        List<HAProxyTLV> haProxyTLVs = buildHAProxyTLV(inboundChannel);
+
         return new HAProxyMessage(HAProxyProtocolVersion.V2, HAProxyCommand.PROXY,
             proxiedProtocol, sourceAddress, destinationAddress, sourcePort, destinationPort, haProxyTLVs);
     }
 
+    protected List<HAProxyTLV> buildHAProxyTLV(Channel inboundChannel) throws IllegalAccessException, DecoderException {
+        List<HAProxyTLV> result = new ArrayList<>();
+        if (!inboundChannel.hasAttr(AttributeKeys.PROXY_PROTOCOL_ADDR)) {
+            return result;
+        }
+        Attribute<?>[] attributes = (Attribute<?>[]) FieldUtils.readField(FIELD_ATTRIBUTE, inboundChannel);
+        if (ArrayUtils.isEmpty(attributes)) {
+            return result;
+        }
+        for (Attribute<?> attribute : attributes) {
+            String attributeKey = attribute.key().name();
+            if (!StringUtils.startsWith(attributeKey, HAProxyConstants.PROXY_PROTOCOL_TLV_PREFIX)) {
+                continue;
+            }
+            String attributeValue = (String) attribute.get();
+            HAProxyTLV haProxyTLV = buildHAProxyTLV(attributeKey, attributeValue);
+            if (haProxyTLV != null) {
+                result.add(haProxyTLV);
+            }
+        }
+        return result;
+    }
+
     protected HAProxyTLV buildHAProxyTLV(String attributeKey, String attributeValue) throws DecoderException {
         String typeString = StringUtils.substringAfter(attributeKey, HAProxyConstants.PROXY_PROTOCOL_TLV_PREFIX);
         ByteBuf byteBuf = Unpooled.buffer();