Skip to content

Commit

Permalink
Fix concurrent bug about session count (vesoft-inc#2636)
Browse files Browse the repository at this point in the history
Co-authored-by: Yee <2520865+yixinglu@users.noreply.github.com>
  • Loading branch information
nebula-bots and yixinglu authored Apr 17, 2023
1 parent 7cb8287 commit 2a19754
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 31 deletions.
42 changes: 40 additions & 2 deletions src/common/session/SessionManager.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {
Expand Down Expand Up @@ -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
Expand Down
25 changes: 2 additions & 23 deletions src/graph/session/GraphSessionManager.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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
6 changes: 0 additions & 6 deletions src/graph/session/GraphSessionManager.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 2a19754

Please sign in to comment.