Fix concurrent bug about session count (#5496)

src/common/session/SessionManager.h

@@ -6,6 +6,7 @@
 #ifndef COMMON_SESSION_SESSIONMANAGER_H_
 #define COMMON_SESSION_SESSIONMANAGER_H_
 
+#include <folly/RWSpinLock.h>
 #include <folly/concurrency/ConcurrentHashMap.h>
 
 #include "clients/meta/MetaClient.h"
@@ -23,7 +24,7 @@ namespace nebula {
 
 class SessionCount {
  private:
-  std::atomic<int32_t> count_ = 1;
+  std::atomic<int32_t> count_{0};
 
  public:
   int fetch_add(int step) {
@@ -75,11 +76,48 @@ class SessionManager {
  protected:
   using SessionPtr = std::shared_ptr<SessionType>;
   using SessionCountPtr = std::shared_ptr<SessionCount>;
+
+  // Get session count pointer according to key
+  SessionCountPtr sessionCnt(const std::string& key) {
+    folly::RWSpinLock::ReadHolder rh(&sessCntLock_);
+    auto iter = userIpSessionCount_.find(key);
+    if (iter != userIpSessionCount_.end()) {
+      return iter->second;
+    }
+    return nullptr;
+  }
+
+  // add sessionCount
+  void addSessionCount(std::string& key) {
+    auto sessCntPtr = sessionCnt(key);
+    if (!sessCntPtr) {
+      folly::RWSpinLock::WriteHolder wh(&sessCntLock_);
+      auto iter = userIpSessionCount_.emplace(key, std::make_shared<SessionCount>());
+      sessCntPtr = iter.first->second;
+    }
+    sessCntPtr->fetch_add(1);
+  }
+
+  // sub sessionCount
+  void subSessionCount(std::string& key) {
+    auto countFindPtr = sessionCnt(key);
+    if (countFindPtr) {
+      auto count = countFindPtr->fetch_sub(1);
+      if (count == 1) {
+        folly::RWSpinLock::WriteHolder wh(&sessCntLock_);
+        userIpSessionCount_.erase(key);
+      }
+    }
+  }
+
   folly::ConcurrentHashMap<SessionID, SessionPtr> activeSessions_;
-  folly::ConcurrentHashMap<std::string, SessionCountPtr> userIpSessionCount_;
   std::unique_ptr<thread::GenericWorker> scavenger_;
   meta::MetaClient* metaClient_{nullptr};
   HostAddr myAddr_;
+
+ private:
+  folly::RWSpinLock sessCntLock_;
+  std::unordered_map<std::string, SessionCountPtr> userIpSessionCount_;
 };
 
 }  // namespace nebula

src/graph/session/GraphSessionManager.cpp

@@ -114,9 +114,8 @@ folly::Future<StatusOr<std::shared_ptr<ClientSession>>> GraphSessionManager::cre
   // check the number of sessions per user per ip
   std::string key = userName + clientIp;
   auto maxSessions = FLAGS_max_sessions_per_ip_per_user;
-  auto uiscFindPtr = userIpSessionCount_.find(key);
-  if (uiscFindPtr != userIpSessionCount_.end() && maxSessions > 0 &&
-      uiscFindPtr->second.get()->get() > maxSessions - 1) {
+  auto uiscFindPtr = sessionCnt(key);
+  if (maxSessions > 0 && uiscFindPtr && uiscFindPtr->get() > maxSessions - 1) {
     return Status::Error(
         "Create Session failed: Too many sessions created from %s by user %s. "
         "the threshold is %d. You can change it by modifying '%s' in nebula-graphd.conf",
@@ -362,25 +361,5 @@ void GraphSessionManager::removeSessionFromLocalCache(const std::vector<SessionI
   }
 }
 
-void GraphSessionManager::addSessionCount(std::string& key) {
-  auto countFindPtr = userIpSessionCount_.find(key);
-  if (countFindPtr != userIpSessionCount_.end()) {
-    countFindPtr->second.get()->fetch_add(1);
-  } else {
-    userIpSessionCount_.insert_or_assign(key, std::make_shared<SessionCount>());
-  }
-}
-
-void GraphSessionManager::subSessionCount(std::string& key) {
-  auto countFindPtr = userIpSessionCount_.find(key);
-  if (countFindPtr == userIpSessionCount_.end()) {
-    VLOG(1) << "Session count not found for key: " << key;
-    return;
-  }
-  auto count = countFindPtr->second.get()->fetch_sub(1);
-  if (count <= 0) {
-    userIpSessionCount_.erase(countFindPtr);
-  }
-}
 }  // namespace graph
 }  // namespace nebula

src/graph/session/GraphSessionManager.h

@@ -112,12 +112,6 @@ class GraphSessionManager final : public SessionManager<ClientSession> {
   // Updates session info locally.
   // session: ClientSession which will be updated.
   void updateSessionInfo(ClientSession* session);
-
-  // add sessionCount
-  void addSessionCount(std::string& key);
-
-  // sub sessionCount
-  void subSessionCount(std::string& key);
 };
 
 }  // namespace graph