diff --git a/remoting/src/main/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandler.java b/remoting/src/main/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandler.java index c6a97fe441b..a9f2908504d 100644 --- a/remoting/src/main/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandler.java +++ b/remoting/src/main/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandler.java @@ -33,18 +33,44 @@ public class RemotingCodeDistributionHandler extends ChannelDuplexHandler { private final ConcurrentMap inboundDistribution; private final ConcurrentMap outboundDistribution; + // Packed volatile holder for atomic publication of (code, adder) pair. + // Avoids the race where another EventLoop sees updated code but stale adder. + private static final class CodeAdderPair { + final int code; + final LongAdder adder; + CodeAdderPair(int code, LongAdder adder) { + this.code = code; + this.adder = adder; + } + } + + private volatile CodeAdderPair lastIn = new CodeAdderPair(Integer.MIN_VALUE, null); + private volatile CodeAdderPair lastOut = new CodeAdderPair(Integer.MIN_VALUE, null); + public RemotingCodeDistributionHandler() { inboundDistribution = new ConcurrentHashMap<>(); outboundDistribution = new ConcurrentHashMap<>(); } private void countInbound(int requestCode) { + CodeAdderPair pair = lastIn; + if (requestCode == pair.code && pair.adder != null) { + pair.adder.increment(); + return; + } LongAdder item = inboundDistribution.computeIfAbsent(requestCode, k -> new LongAdder()); + lastIn = new CodeAdderPair(requestCode, item); item.increment(); } private void countOutbound(int responseCode) { + CodeAdderPair pair = lastOut; + if (responseCode == pair.code && pair.adder != null) { + pair.adder.increment(); + return; + } LongAdder item = outboundDistribution.computeIfAbsent(responseCode, k -> new LongAdder()); + lastOut = new CodeAdderPair(responseCode, item); item.increment(); } diff --git a/remoting/src/test/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandlerTest.java b/remoting/src/test/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandlerTest.java index eb623a9de92..e6f3e836fdc 100644 --- a/remoting/src/test/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandlerTest.java +++ b/remoting/src/test/java/org/apache/rocketmq/remoting/netty/RemotingCodeDistributionHandlerTest.java @@ -69,4 +69,38 @@ public void remotingCodeCountTest() throws Exception { return f1 && f2; }); } + + @Test + public void testMultipleCodesIndependent() throws Exception { + RemotingCodeDistributionHandler handler = new RemotingCodeDistributionHandler(); + Class clazz = RemotingCodeDistributionHandler.class; + Method methodIn = clazz.getDeclaredMethod("countInbound", int.class); + methodIn.setAccessible(true); + + // Count code=10 three times, code=20 twice + methodIn.invoke(handler, 10); + methodIn.invoke(handler, 10); + methodIn.invoke(handler, 10); + methodIn.invoke(handler, 20); + methodIn.invoke(handler, 20); + + String snapshot = handler.getInBoundSnapshotString(); + Assert.assertTrue(snapshot.contains("10:3")); + Assert.assertTrue(snapshot.contains("20:2")); + } + + @Test + public void testCacheHitSameCode() throws Exception { + RemotingCodeDistributionHandler handler = new RemotingCodeDistributionHandler(); + Class clazz = RemotingCodeDistributionHandler.class; + Method methodIn = clazz.getDeclaredMethod("countInbound", int.class); + methodIn.setAccessible(true); + + // Same code repeatedly should hit the CodeAdderPair cache and still count correctly + for (int i = 0; i < 100; i++) { + methodIn.invoke(handler, 42); + } + String snapshot = handler.getInBoundSnapshotString(); + Assert.assertTrue(snapshot.contains("42:100")); + } } \ No newline at end of file