diff --git a/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java b/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java index bc8400c19a2..1b1d2ded9eb 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java +++ b/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 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 prev = this.groupChannelTable.putIfAbsent(group, channelTable); - channelTable = prev != null ? prev : channelTable; + ConcurrentMap newChannelTable = new ConcurrentHashMap<>(); + newChannelTable.put(clientChannelInfo.getChannel(), clientChannelInfo); + ConcurrentMap 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()); } diff --git a/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java b/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java index 451b0e044c7..02cd46bc4ec 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java +++ b/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 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(); } -} \ No newline at end of file +}