diff --git a/src/common/session/SessionManager.h b/src/common/session/SessionManager.h index 00ef5fe63ce..c1d0f32ad61 100644 --- a/src/common/session/SessionManager.h +++ b/src/common/session/SessionManager.h @@ -6,6 +6,7 @@ #ifndef COMMON_SESSION_SESSIONMANAGER_H_ #define COMMON_SESSION_SESSIONMANAGER_H_ +#include #include #include "clients/meta/MetaClient.h" @@ -23,7 +24,7 @@ namespace nebula { class SessionCount { private: - std::atomic count_ = 1; + std::atomic count_{0}; public: int fetch_add(int step) { @@ -75,11 +76,48 @@ class SessionManager { protected: using SessionPtr = std::shared_ptr; using SessionCountPtr = std::shared_ptr; + + // 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()); + 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 activeSessions_; - folly::ConcurrentHashMap userIpSessionCount_; std::unique_ptr scavenger_; meta::MetaClient* metaClient_{nullptr}; HostAddr myAddr_; + + private: + folly::RWSpinLock sessCntLock_; + std::unordered_map userIpSessionCount_; }; } // namespace nebula diff --git a/src/graph/session/GraphSessionManager.cpp b/src/graph/session/GraphSessionManager.cpp index fbd7d02fd43..b8f0d3f5b96 100644 --- a/src/graph/session/GraphSessionManager.cpp +++ b/src/graph/session/GraphSessionManager.cpp @@ -114,9 +114,8 @@ folly::Future>> 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::vectorsecond.get()->fetch_add(1); - } else { - userIpSessionCount_.insert_or_assign(key, std::make_shared()); - } -} - -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 diff --git a/src/graph/session/GraphSessionManager.h b/src/graph/session/GraphSessionManager.h index cc13f1a956e..045403b877e 100644 --- a/src/graph/session/GraphSessionManager.h +++ b/src/graph/session/GraphSessionManager.h @@ -112,12 +112,6 @@ class GraphSessionManager final : public SessionManager { // 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