perf[session]: use high-performance segment locks to store session maps

net/src/main/java/com/zfoo/net/NetContext.java

@@ -134,8 +134,8 @@ public class NetContext implements ApplicationListener<ApplicationContextEvent>,
         configManager.getRegistry().shutdown();
 
         // 先关闭所有session
-        IOUtils.closeIO(ArrayUtils.listToArray(new ArrayList<>(sessionManager.getClientSessionMap().values()), Session.class));
-        IOUtils.closeIO(ArrayUtils.listToArray(new ArrayList<>(sessionManager.getServerSessionMap().values()), Session.class));
+        sessionManager.forEachClientSession(it -> IOUtils.closeIO(it));
+        sessionManager.forEachServerSession(it -> IOUtils.closeIO(it));
 
         // 关闭客户端和服务器
         AbstractClient.shutdown();

net/src/main/java/com/zfoo/net/consumer/balancer/AbstractConsumerLoadBalancer.java

@@ -24,6 +24,7 @@ import com.zfoo.protocol.util.StringUtils;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 /**
@@ -52,37 +53,22 @@ public abstract class AbstractConsumerLoadBalancer implements IConsumerLoadBalan
     }
 
     public List<Session> getSessionsByModule(ProtocolModule module) {
-        return NetContext.getSessionManager().getClientSessionMap()
-                .values()
-                .stream()
-                .filter(it -> it.getConsumerAttribute() != null && it.getConsumerAttribute().getProviderConfig() != null)
-                .filter(it -> it.getConsumerAttribute().getProviderConfig().getProviders().stream().anyMatch(provider -> provider.getProtocolModule().equals(module)))
-                .collect(Collectors.toList());
+        var list = new ArrayList<Session>();
+        NetContext.getSessionManager().forEachClientSession(new Consumer<Session>() {
+            @Override
+            public void accept(Session session) {
+                if (session.getConsumerAttribute() == null || session.getConsumerAttribute().getProviderConfig() == null) {
+                    return;
+                }
+                var providerConfig = session.getConsumerAttribute().getProviderConfig();
+                if (providerConfig.getProviders().stream().anyMatch(it -> it.getProtocolModule().equals(module))) {
+                    list.add(session);
+                }
+            }
+        });
+        return list;
     }
 
-    public List<Session> sessionsByModule(ProtocolModule module) {
-        var clientSessionMap = NetContext.getSessionManager().getClientSessionMap();
-        var sessions = new ArrayList<Session>();
-        for (var clientSession : clientSessionMap.values()) {
-            var consumerAttribute = clientSession.getConsumerAttribute();
-            if (consumerAttribute == null) {
-                continue;
-            }
-
-            var registerVO = (RegisterVO) consumerAttribute;
-            var providerConfig = registerVO.getProviderConfig();
-            if (providerConfig == null) {
-                continue;
-            }
-
-            if (providerConfig.getProviders().stream().anyMatch(it -> it.getProtocolModule().getId() == module.getId())) {
-                sessions.add(clientSession);
-            }
-        }
-        return sessions;
-    }
-
-
     public boolean sessionHasModule(Session session, IPacket packet) {
         var consumerAttribute = session.getConsumerAttribute();
         if (Objects.isNull(consumerAttribute)) {

net/src/main/java/com/zfoo/net/consumer/registry/ZookeeperRegistry.java

@@ -18,6 +18,7 @@ import com.zfoo.net.NetContext;
 import com.zfoo.net.consumer.event.ConsumerStartEvent;
 import com.zfoo.net.core.tcp.TcpClient;
 import com.zfoo.net.core.tcp.TcpServer;
+import com.zfoo.net.session.Session;
 import com.zfoo.net.util.SessionUtils;
 import com.zfoo.protocol.collection.ArrayUtils;
 import com.zfoo.protocol.collection.concurrent.ConcurrentArrayList;
@@ -49,10 +50,7 @@ import org.apache.zookeeper.server.auth.DigestAuthenticationProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collections;
-import java.util.List;
-import java.util.Objects;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
@@ -452,17 +450,22 @@ public class ZookeeperRegistry implements IRegistry {
             return;
         }
 
-        logger.info("start using providerHashConsumerSet:{} to check [consumer:{}]", providerHashConsumerSet, NetContext.getSessionManager().getClientSessionMap().size());
+        logger.info("start using providerHashConsumerSet:{} to check [consumer:{}]", providerHashConsumerSet, NetContext.getSessionManager().clientSessionSize());
 
         var recheckFlag = false;
 
         for (var providerCache : providerHashConsumerSet) {
             // 先排除已经启动的consumer
             // getClientSessionMap
-            var consumerClientList = NetContext.getSessionManager().getClientSessionMap().values()
-                    .stream()
-                    .filter(it -> it.getConsumerAttribute() != null && it.getConsumerAttribute().equals(providerCache))
-                    .collect(Collectors.toList());
+            var consumerClientList = new ArrayList<Session>();
+            NetContext.getSessionManager().forEachClientSession(new Consumer<Session>() {
+                        @Override
+                        public void accept(Session session) {
+                            if (session.getConsumerAttribute() != null && session.getConsumerAttribute().equals(providerCache)) {
+                                consumerClientList.add(session);
+                            }
+                        }
+                    });
 
             if (consumerClientList.size() == 1) {
                 var consumer = consumerClientList.get(0);

net/src/main/java/com/zfoo/net/session/ISessionManager.java

@@ -12,7 +12,7 @@
 
 package com.zfoo.net.session;
 
-import java.util.Map;
+import java.util.function.Consumer;
 
 /**
  * @author godotg
@@ -26,8 +26,9 @@ public interface ISessionManager {
 
     Session getServerSession(long id);
 
-    Map<Long, Session> getServerSessionMap();
+    int serverSessionSize();
 
+    void forEachServerSession(Consumer<Session> consumer);
 
     void addClientSession(Session session);
 
@@ -35,7 +36,9 @@ public interface ISessionManager {
 
     Session getClientSession(long id);
 
-    Map<Long, Session> getClientSessionMap();
+    void forEachClientSession(Consumer<Session> consumer);
+
+    int clientSessionSize();
 
     int getClientSessionChangeId();
 

net/src/main/java/com/zfoo/net/session/SessionManager.java

@@ -13,13 +13,12 @@
 package com.zfoo.net.session;
 
 import com.zfoo.net.util.SessionUtils;
+import com.zfoo.protocol.collection.concurrent.ConcurrentHashMapLongObject;
 import com.zfoo.util.security.IdUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
 
 /**
  * @author godotg
@@ -34,14 +33,14 @@ public class SessionManager implements ISessionManager {
      * 如：自己作为网关，那肯定有一大堆客户端连接，他们连接上来后，就会保存下来这些信息。
      * 因此：要全局消息广播，其实要用这个Map
      */
-    private final Map<Long, Session> serverSessionMap = new ConcurrentHashMap<>();
+    private final ConcurrentHashMapLongObject<Session> serverSessionMap = new ConcurrentHashMapLongObject<>(128);
 
 
     /**
      * 作为客户端，连接别的服务器上后，保存下来的Session
      * 如：自己配置了Consumer，说明自己作为消费者将要消费远程接口，就会创建一个TcpClient去连接Provider，那么连接上后，就会保存下来到这个Map中
      */
-    private final Map<Long, Session> clientSessionMap = new ConcurrentHashMap<>();
+    private final ConcurrentHashMapLongObject<Session> clientSessionMap = new ConcurrentHashMapLongObject<>();
 
     private volatile int clientSessionChangeId = IdUtils.getLocalIntId();
 
@@ -61,8 +60,9 @@ public class SessionManager implements ISessionManager {
             logger.error("[session:{}] does not exist", SessionUtils.sessionInfo(session));
             return;
         }
-        serverSessionMap.remove(session.getSid());
-        session.close();
+        try (session) {
+            serverSessionMap.remove(session.getSid());
+        }
     }
 
     @Override
@@ -71,8 +71,13 @@ public class SessionManager implements ISessionManager {
     }
 
     @Override
-    public Map<Long, Session> getServerSessionMap() {
-        return Collections.unmodifiableMap(serverSessionMap);
+    public int serverSessionSize() {
+        return serverSessionMap.size();
+    }
+
+    @Override
+    public void forEachServerSession(Consumer<Session> consumer) {
+        serverSessionMap.forEachPrimitive(it -> consumer.accept(it.value()));
     }
 
     @Override
@@ -91,8 +96,9 @@ public class SessionManager implements ISessionManager {
             logger.error("[session:{}] does not exist", SessionUtils.sessionInfo(session));
             return;
         }
-        clientSessionMap.remove(session.getSid());
-        session.close();
+        try (session) {
+            clientSessionMap.remove(session.getSid());
+        }
         clientSessionChangeId = IdUtils.getLocalIntId();
     }
 
@@ -102,8 +108,13 @@ public class SessionManager implements ISessionManager {
     }
 
     @Override
-    public Map<Long, Session> getClientSessionMap() {
-        return Collections.unmodifiableMap(clientSessionMap);
+    public void forEachClientSession(Consumer<Session> consumer) {
+        clientSessionMap.forEachPrimitive(it -> consumer.accept(it.value()));
+    }
+
+    @Override
+    public int clientSessionSize() {
+        return clientSessionMap.size();
     }
 
     @Override

net/src/test/java/com/zfoo/net/session/SessionUtils.java

@@ -18,6 +18,8 @@ import com.zfoo.protocol.util.FileUtils;
 import com.zfoo.protocol.util.StringUtils;
 import com.zfoo.util.ThreadUtils;
 
+import java.util.function.Consumer;
+
 /**
  * @author godotg
  * @version 3.0
@@ -29,19 +31,24 @@ public abstract class SessionUtils {
             while (true) {
                 ThreadUtils.sleep(10_000);
                 var builder = new StringBuilder();
-                builder.append(StringUtils.format("clientSession count：[{}]", NetContext.getSessionManager().getClientSessionMap().size()));
                 builder.append(FileUtils.LS);
-                for (Session session : NetContext.getSessionManager().getClientSessionMap().values()) {
-                    builder.append(StringUtils.format("[session:{}]", session.getChannel().remoteAddress()));
-                    builder.append(FileUtils.LS);
-                }
+                NetContext.getSessionManager().forEachClientSession(new Consumer<Session>() {
+                    @Override
+                    public void accept(Session session) {
+                        builder.append(StringUtils.format("[session:{}]", session.getChannel().remoteAddress()));
+                        builder.append(FileUtils.LS);
+                    }
+                });
 
-                builder.append(StringUtils.format("serverSession count：[{}]", NetContext.getSessionManager().getServerSessionMap().size()));
+                builder.append(StringUtils.format("serverSession count：[{}]", NetContext.getSessionManager().serverSessionSize()));
                 builder.append(FileUtils.LS);
-                for (Session session : NetContext.getSessionManager().getServerSessionMap().values()) {
-                    builder.append(StringUtils.format("[session:{}]", session.getChannel().remoteAddress()));
-                    builder.append(FileUtils.LS);
-                }
+                NetContext.getSessionManager().forEachServerSession(new Consumer<Session>() {
+                    @Override
+                    public void accept(Session session) {
+                        builder.append(StringUtils.format("[session:{}]", session.getChannel().remoteAddress()));
+                        builder.append(FileUtils.LS);
+                    }
+                });
 
                 System.out.println(builder.toString());