Fix connection closed but not clean session issue

client/src/main/java/org/apache/rocketmq/client/consumer/store/RemoteBrokerOffsetStore.java

@@ -25,7 +25,6 @@ import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.rocketmq.client.exception.MQBrokerException;
 import org.apache.rocketmq.client.exception.MQClientException;
-import org.apache.rocketmq.client.impl.FindBrokerResult;
 import org.apache.rocketmq.client.impl.factory.MQClientInstance;
 import org.apache.rocketmq.client.log.ClientLogger;
 import org.apache.rocketmq.common.MixAll;
@@ -193,6 +192,7 @@ public class RemoteBrokerOffsetStore implements OffsetStore {
         MQBrokerException, InterruptedException, MQClientException {
         updateConsumeOffsetToBroker(mq, offset, true);
     }
+
     private void updateConsumeOffsetToSnode(MessageQueue mq, long offset) throws RemotingException,
         MQBrokerException, InterruptedException, MQClientException {
         updateConsumeOffsetToBroker(mq, offset, true);
@@ -206,9 +206,9 @@ public class RemoteBrokerOffsetStore implements OffsetStore {
         MQBrokerException, InterruptedException, MQClientException {
 
         String snodeAddr = this.mQClientFactory.findSnodeAddressInPublish();
-        if (null == snodeAddr){
+        if (null == snodeAddr) {
             this.mQClientFactory.updateSnodeInfoFromNameServer();
-            snodeAddr= this.mQClientFactory.findSnodeAddressInPublish();
+            snodeAddr = this.mQClientFactory.findSnodeAddressInPublish();
         }
 
         if (snodeAddr != null) {
@@ -233,9 +233,9 @@ public class RemoteBrokerOffsetStore implements OffsetStore {
     private long fetchConsumeOffsetFromBroker(MessageQueue mq) throws RemotingException, MQBrokerException,
         InterruptedException, MQClientException {
         String snodeAddr = this.mQClientFactory.findSnodeAddressInPublish();
-        if (null == snodeAddr){
+        if (null == snodeAddr) {
             this.mQClientFactory.updateSnodeInfoFromNameServer();
-            snodeAddr= this.mQClientFactory.findSnodeAddressInPublish();
+            snodeAddr = this.mQClientFactory.findSnodeAddressInPublish();
         }
 
         if (snodeAddr != null) {

client/src/main/java/org/apache/rocketmq/client/exception/MQSnodeException.java

@@ -16,12 +16,9 @@
  */
 package org.apache.rocketmq.client.exception;
 
-import org.apache.rocketmq.common.UtilAll;
-import org.apache.rocketmq.common.help.FAQUrl;
-
 public class MQSnodeException extends MQBrokerException {
 
     public MQSnodeException(int responseCode, String errorMessage) {
-        super(responseCode,errorMessage);
+        super(responseCode, errorMessage);
     }
 }

client/src/main/java/org/apache/rocketmq/client/impl/consumer/DefaultMQPushConsumerImpl.java

@@ -56,7 +56,6 @@ import org.apache.rocketmq.common.UtilAll;
 import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
 import org.apache.rocketmq.common.filter.FilterAPI;
 import org.apache.rocketmq.common.help.FAQUrl;
-import org.apache.rocketmq.logging.InternalLogger;
 import org.apache.rocketmq.common.message.Message;
 import org.apache.rocketmq.common.message.MessageAccessor;
 import org.apache.rocketmq.common.message.MessageConst;
@@ -72,8 +71,8 @@ import org.apache.rocketmq.common.protocol.heartbeat.SubscriptionData;
 import org.apache.rocketmq.common.protocol.route.BrokerData;
 import org.apache.rocketmq.common.protocol.route.TopicRouteData;
 import org.apache.rocketmq.common.sysflag.PullSysFlag;
+import org.apache.rocketmq.logging.InternalLogger;
 import org.apache.rocketmq.remoting.RPCHook;
-import org.apache.rocketmq.remoting.common.RemotingHelper;
 import org.apache.rocketmq.remoting.exception.RemotingException;
 
 public class DefaultMQPushConsumerImpl implements MQConsumerInner {
@@ -1138,6 +1137,7 @@ public class DefaultMQPushConsumerImpl implements MQConsumerInner {
         this.consumeMessageService = consumeMessageService;
 
     }
+
     private void tryToFindSnodePublishInfo() {
         this.mQClientFactory.updateSnodeInfoFromNameServer();
     }

client/src/main/java/org/apache/rocketmq/client/impl/factory/MQClientInstance.java

@@ -38,7 +38,6 @@ import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import org.apache.rocketmq.client.admin.MQAdminExtInner;
 import org.apache.rocketmq.client.common.ThreadLocalIndex;
-import org.apache.rocketmq.client.exception.MQBrokerException;
 import org.apache.rocketmq.client.exception.MQClientException;
 import org.apache.rocketmq.client.impl.ClientRemotingProcessor;
 import org.apache.rocketmq.client.impl.FindBrokerResult;
@@ -80,7 +79,6 @@ import org.apache.rocketmq.common.protocol.route.TopicRouteData;
 import org.apache.rocketmq.logging.InternalLogger;
 import org.apache.rocketmq.remoting.ClientConfig;
 import org.apache.rocketmq.remoting.common.RemotingHelper;
-import org.apache.rocketmq.remoting.exception.RemotingException;
 import org.apache.rocketmq.remoting.interceptor.InterceptorGroup;
 import org.apache.rocketmq.remoting.protocol.RemotingCommand;
 

snode/src/main/java/org/apache/rocketmq/snode/SnodeController.java

@@ -352,6 +352,8 @@ public class SnodeController {
                 this.consumerManageExecutor);
         this.mqttRemotingServer.registerProcessor(RequestCode.MQTT_MESSAGE,
             defaultMqttMessageProcessor, handleMqttMessageExecutor);
+
+
         defaultMqttMessageProcessor.registerMessageHanlder(MqttMessageType.CONNECT,
             new MqttConnectMessageHandler(this));
         defaultMqttMessageProcessor.registerMessageHanlder(MqttMessageType.DISCONNECT,
@@ -372,6 +374,9 @@ public class SnodeController {
             new MqttSubscribeMessageHandler(this));
         defaultMqttMessageProcessor.registerMessageHanlder(MqttMessageType.UNSUBSCRIBE,
             new MqttUnsubscribeMessagHandler(this));
+
+
+
     }
 
     public void start() {

snode/src/main/java/org/apache/rocketmq/snode/client/ClientHousekeepingService.java

@@ -25,7 +25,6 @@ import org.apache.rocketmq.remoting.ChannelEventListener;
 import org.apache.rocketmq.remoting.RemotingChannel;
 import org.apache.rocketmq.remoting.common.RemotingHelper;
 import org.apache.rocketmq.remoting.netty.NettyChannelImpl;
-import org.apache.rocketmq.snode.client.impl.ClientRole;
 import org.apache.rocketmq.snode.constant.SnodeConstant;
 
 public class ClientHousekeepingService implements ChannelEventListener {
@@ -53,15 +52,13 @@ public class ClientHousekeepingService implements ChannelEventListener {
         this.iotClientManager.shutdown();
     }
 
-    private ClientRole clientRole(RemotingChannel remotingChannel) {
+    private Client getClient(RemotingChannel remotingChannel) {
         if (remotingChannel instanceof NettyChannelImpl) {
             Channel channel = ((NettyChannelImpl) remotingChannel).getChannel();
             Attribute<Client> clientAttribute = channel.attr(SnodeConstant.NETTY_CLIENT_ATTRIBUTE_KEY);
             if (clientAttribute != null) {
                 Client client = clientAttribute.get();
-                if (client != null) {
-                    return client.getClientRole();
-                }
+                return client;
             }
         }
         log.warn("RemotingChannel type error: {}", remotingChannel.getClass());
@@ -69,17 +66,17 @@ public class ClientHousekeepingService implements ChannelEventListener {
     }
 
     private void closeChannel(String remoteAddress, RemotingChannel remotingChannel) {
-        ClientRole clientRole = clientRole(remotingChannel);
-        if (clientRole != null) {
-            switch (clientRole) {
+        Client client = getClient(remotingChannel);
+        if (client != null) {
+            switch (client.getClientRole()) {
                 case Consumer:
-                    this.consumerManager.onClose(remoteAddress, remotingChannel);
+                    this.consumerManager.onClose(client.getGroups(), remotingChannel);
                     return;
                 case Producer:
-                    this.producerManager.onClose(remoteAddress, remotingChannel);
+                    this.producerManager.onClose(client.getGroups(), remotingChannel);
                     return;
                 case IOTCLIENT:
-                    this.iotClientManager.onClose(remoteAddress, remotingChannel);
+                    this.iotClientManager.onClose(client.getGroups(), remotingChannel);
                     return;
                 default:
             }

snode/src/main/java/org/apache/rocketmq/snode/client/ClientManager.java

@@ -17,6 +17,7 @@
 package org.apache.rocketmq.snode.client;
 
 import java.util.List;
+import java.util.Set;
 import org.apache.rocketmq.remoting.RemotingChannel;
 
 public interface ClientManager {
@@ -24,7 +25,7 @@ public interface ClientManager {
 
     void unRegister(String groupId, RemotingChannel remotingChannel);
 
-    void onClose(String groupId, RemotingChannel remotingChannel);
+    void onClose(Set<String> groupId, RemotingChannel remotingChannel);
 
     List<RemotingChannel> getChannels(String groupId);
 

snode/src/main/java/org/apache/rocketmq/snode/client/impl/ClientManagerImpl.java

@@ -20,6 +20,7 @@ import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -30,6 +31,8 @@ import org.apache.rocketmq.logging.InternalLogger;
 import org.apache.rocketmq.logging.InternalLoggerFactory;
 import org.apache.rocketmq.remoting.RemotingChannel;
 import org.apache.rocketmq.remoting.common.RemotingHelper;
+import org.apache.rocketmq.remoting.netty.NettyChannelHandlerContextImpl;
+import org.apache.rocketmq.remoting.netty.NettyChannelImpl;
 import org.apache.rocketmq.snode.client.Client;
 import org.apache.rocketmq.snode.client.ClientManager;
 
@@ -104,16 +107,20 @@ public abstract class ClientManagerImpl implements ClientManager {
     public boolean register(String groupId, Client client) {
         boolean updated = false;
         if (client != null) {
-            ConcurrentHashMap<RemotingChannel, Client> channelTable = groupClientTable.get(groupId);
+            ConcurrentHashMap<RemotingChannel, Client> channelTable = this.groupClientTable.get(groupId);
             if (channelTable == null) {
                 channelTable = new ConcurrentHashMap();
                 ConcurrentHashMap prev = groupClientTable.putIfAbsent(groupId, channelTable);
                 channelTable = prev != null ? prev : channelTable;
             }
-
-            Client oldClient = channelTable.get(client.getRemotingChannel());
+            log.info("*********");
+            RemotingChannel remotingChannel = client.getRemotingChannel();
+            if (remotingChannel instanceof NettyChannelHandlerContextImpl) {
+                remotingChannel = new NettyChannelImpl(((NettyChannelHandlerContextImpl) remotingChannel).getChannelHandlerContext().channel());
+            }
+            Client oldClient = channelTable.get(remotingChannel);
             if (oldClient == null) {
-                Client prev = channelTable.put(client.getRemotingChannel(), client);
+                Client prev = channelTable.put(remotingChannel, client);
                 if (prev != null) {
                     log.info("New client connected, group: {} {} {} channel: {}", groupId,
                         client.toString());
@@ -127,14 +134,14 @@ public abstract class ClientManagerImpl implements ClientManager {
                         groupId,
                         oldClient.toString(),
                         channelTable.toString());
-                    channelTable.put(client.getRemotingChannel(), client);
+                    channelTable.put(remotingChannel, client);
                 }
             }
             oldClient.setLastUpdateTimestamp(System.currentTimeMillis());
+            onRegister(groupId, remotingChannel);
         }
-        log.debug("Register client role: {}, group: {}, last: {}", client.getClientRole(), groupId,
+        log.info("Register client role: {}, group: {}, last: {}", client.getClientRole(), groupId,
             client.getLastUpdateTimestamp());
-        onRegister(groupId, client.getRemotingChannel());
         return updated;
     }
 
@@ -157,13 +164,16 @@ public abstract class ClientManagerImpl implements ClientManager {
     public void unRegister(String groupId, RemotingChannel remotingChannel) {
         removeClient(groupId, remotingChannel);
         onUnregister(groupId, remotingChannel);
+
     }
 
     @Override
-    public void onClose(String groupId, RemotingChannel remotingChannel) {
+    public void onClose(Set<String> groups, RemotingChannel remotingChannel) {
+        for (String groupId : groups) {
             removeClient(groupId, remotingChannel);
             onClosed(groupId, remotingChannel);
         }
+    }
 
     public List<RemotingChannel> getChannels(String groupId) {
         if (groupId != null) {

snode/src/main/java/org/apache/rocketmq/snode/processor/ConsumerManageProcessor.java

@@ -130,9 +130,9 @@ public class ConsumerManageProcessor implements RequestProcessor {
         if (!clientIds.isEmpty()) {
             GetConsumerListByGroupResponseBody body = new GetConsumerListByGroupResponseBody();
             body.setConsumerIdList(clientIds);
-            response.setBody(body.encode());
             response.setCode(ResponseCode.SUCCESS);
             response.setRemark(null);
+            response.setBody(body.encode());
             return response;
         } else {
             log.warn("GetAllClientId failed, {} {}", requestHeader.getConsumerGroup(),
@@ -163,11 +163,11 @@ public class ConsumerManageProcessor implements RequestProcessor {
         RemotingSendRequestException, RemotingConnectException, RemotingCommandException {
         final RemotingCommand response =
             RemotingCommand.createResponseCommand(QueryConsumerOffsetResponseHeader.class);
-        final QueryConsumerOffsetResponseHeader responseHeader =
-            (QueryConsumerOffsetResponseHeader) response.readCustomHeader();
         final QueryConsumerOffsetRequestHeader requestHeader =
             (QueryConsumerOffsetRequestHeader) request
                 .decodeCommandCustomHeader(QueryConsumerOffsetRequestHeader.class);
+        final QueryConsumerOffsetResponseHeader responseHeader =
+            (QueryConsumerOffsetResponseHeader) response.readCustomHeader();
 
         long offset =
             this.snodeController.getConsumerOffsetManager().queryOffset(requestHeader.getEnodeName(),

snode/src/main/java/org/apache/rocketmq/snode/service/impl/NnodeServiceImpl.java

@@ -58,7 +58,7 @@ public class NnodeServiceImpl implements NnodeService {
     }
 
     @Override
-    public void registerSnode(SnodeConfig snodeConfig) throws Exception{
+    public void registerSnode(SnodeConfig snodeConfig) throws Exception {
         List<String> nnodeAddressList = this.snodeController.getRemotingClient().getNameServerAddressList();
         RemotingCommand remotingCommand = new RemotingCommand();
         RegisterSnodeRequestHeader requestHeader = new RegisterSnodeRequestHeader();

snode/src/test/java/org/apache/rocketmq/snode/service/NnodeServiceImplTest.java

@@ -60,7 +60,7 @@ public class NnodeServiceImplTest extends SnodeTestBase {
     }
 
     @Test
-    public void registerSnodeSuccessTest() throws InterruptedException, RemotingConnectException,
+    public void registerSnodeTest() throws InterruptedException, RemotingConnectException,
         RemotingSendRequestException, RemotingTimeoutException {
         when(snodeController.getRemotingClient().getNameServerAddressList()).thenReturn(createNnodeList());
         when(snodeController.getRemotingClient().invokeSync(anyString(), any(RemotingCommand.class), anyLong())).thenReturn(createSuccessResponse());