[ISSUE #10698] Restore atomic producer group registration in ProducerManager

broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java

@@ -206,7 +206,8 @@ public boolean doChannelCloseEvent(final String remoteAddr, final Channel channe
     public void registerProducer(final String group, final ClientChannelInfo clientChannelInfo) {
 
         long start = System.currentTimeMillis();
-        ClientChannelInfo clientChannelInfoFound;
+        ClientChannelInfo clientChannelInfoFound = null;
+        boolean newChannel = false;
 
         ConcurrentMap<Channel, ClientChannelInfo> channelTable = this.groupChannelTable.get(group);
         // note that we must take care of the exist groups and channels,
@@ -232,24 +233,30 @@ public void registerProducer(final String group, final ClientChannelInfo clientC
         }
 
         if (null == channelTable) {
-            channelTable = new ConcurrentHashMap<>();
-            ConcurrentMap<Channel, ClientChannelInfo> prev = this.groupChannelTable.putIfAbsent(group, channelTable);
-            channelTable = prev != null ? prev : channelTable;
+            ConcurrentMap<Channel, ClientChannelInfo> newChannelTable = new ConcurrentHashMap<>();
+            newChannelTable.put(clientChannelInfo.getChannel(), clientChannelInfo);
+            ConcurrentMap<Channel, ClientChannelInfo> prev =
+                this.groupChannelTable.putIfAbsent(group, newChannelTable);
+            if (prev == null) {
+                channelTable = newChannelTable;
+                newChannel = true;
+            } else {
+                channelTable = prev;
+            }
+        }
+
+        if (!newChannel) {
+            clientChannelInfoFound = channelTable.putIfAbsent(clientChannelInfo.getChannel(), clientChannelInfo);
+            newChannel = clientChannelInfoFound == null;
         }
 
-        clientChannelInfoFound = channelTable.get(clientChannelInfo.getChannel());
-        // Add client-channel info to existing producer group
-        if (null == clientChannelInfoFound) {
-            channelTable.put(clientChannelInfo.getChannel(), clientChannelInfo);
+        if (newChannel) {
             clientChannelTable.put(clientChannelInfo.getClientId(), clientChannelInfo.getChannel());
             log.info("new producer connected, group: {} channel: {}", group, clientChannelInfo.toString());
             if (this.brokerConfig != null && this.brokerConfig.isEnableFastChannelEventProcess()) {
                 ClientChannelAttributeHelper.addProducerGroup(clientChannelInfo.getChannel(), group);
             }
-        }
-
-        // Refresh existing client-channel-info update-timestamp
-        if (clientChannelInfoFound != null) {
+        } else {
             clientChannelInfoFound.setLastUpdateTimestamp(System.currentTimeMillis());
         }
 

broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java

@@ -18,11 +18,14 @@
 
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
+import io.netty.channel.embedded.EmbeddedChannel;
 import java.lang.reflect.Field;
+import java.lang.reflect.Proxy;
 import java.util.Map;
-
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
-
 import org.apache.rocketmq.common.BrokerConfig;
 import org.apache.rocketmq.remoting.protocol.LanguageCode;
 import org.junit.Before;
@@ -165,6 +168,82 @@ public void testRegisterProducer() {
         assertThat(channel1).isEqualTo(channel);
     }
 
+    @Test
+    public void registerProducerShouldNotPublishEmptyGroupDuringConcurrentScan() throws Exception {
+        CountDownLatch hashEntered = new CountDownLatch(1);
+        CountDownLatch releaseHash = new CountDownLatch(1);
+        AtomicInteger hashCalls = new AtomicInteger();
+        Channel blockingChannel = (Channel) Proxy.newProxyInstance(
+            Channel.class.getClassLoader(), new Class<?>[] {Channel.class}, (proxy, method, args) -> {
+                switch (method.getName()) {
+                    case "hashCode":
+                        if (hashCalls.incrementAndGet() == 1) {
+                            hashEntered.countDown();
+                            if (!releaseHash.await(10, TimeUnit.SECONDS)) {
+                                throw new AssertionError("Timed out waiting to release Channel.hashCode()");
+                            }
+                        }
+                        return System.identityHashCode(proxy);
+                    case "equals":
+                        return proxy == args[0];
+                    case "toString":
+                        return "blockingChannel";
+                    default:
+                        throw new UnsupportedOperationException(method.getName());
+                }
+            });
+        String clientId = "concurrent-client";
+        ClientChannelInfo concurrentClientInfo =
+            new ClientChannelInfo(blockingChannel, clientId, LanguageCode.JAVA, 0);
+        AtomicInteger groupUnregisterCount = new AtomicInteger();
+        producerManager.appendProducerChangeListener((event, changedGroup, info) -> {
+            if (event == ProducerGroupEvent.GROUP_UNREGISTER && group.equals(changedGroup)) {
+                groupUnregisterCount.incrementAndGet();
+            }
+        });
+        AtomicReference<Throwable> registrationFailure = new AtomicReference<>();
+        Thread registerThread = new Thread(() -> {
+            try {
+                producerManager.registerProducer(group, concurrentClientInfo);
+            } catch (Throwable t) {
+                registrationFailure.set(t);
+            }
+        });
+
+        registerThread.start();
+        try {
+            assertThat(hashEntered.await(10, TimeUnit.SECONDS)).isTrue();
+            producerManager.scanNotActiveChannel();
+        } finally {
+            releaseHash.countDown();
+            registerThread.join(10_000);
+        }
+
+        assertThat(registerThread.isAlive()).isFalse();
+        assertThat(registrationFailure.get()).isNull();
+        assertThat(producerManager.getGroupChannelTable().get(group))
+            .containsEntry(blockingChannel, concurrentClientInfo);
+        assertThat(producerManager.findChannel(clientId)).isSameAs(blockingChannel);
+        assertThat(groupUnregisterCount.get()).isZero();
+    }
+
+    @Test
+    public void registerProducerShouldAddFastChannelAttributeOnlyOnce() {
+        brokerConfig.setEnableFastChannelEventProcess(true);
+        EmbeddedChannel fastChannel = new EmbeddedChannel();
+        try {
+            ClientChannelInfo fastClientInfo =
+                new ClientChannelInfo(fastChannel, "fast-client", LanguageCode.JAVA, 0);
+
+            producerManager.registerProducer(group, fastClientInfo);
+            producerManager.registerProducer(group, fastClientInfo);
+
+            assertThat(ClientChannelAttributeHelper.getProducerGroups(fastChannel)).containsExactly(group);
+        } finally {
+            fastChannel.finishAndReleaseAll();
+        }
+    }
+
     @Test
     public void unregisterProducer() throws Exception {
         producerManager.registerProducer(group, clientInfo);
@@ -225,4 +304,4 @@ public void testGetAvailableChannel() {
         assertThat(c).isNull();
     }
 
-}
+}