sserver: Limit the total number of pending user register requests.

btccom · Feb 21, 2019 · 6aaf0b7 · 6aaf0b7
1 parent 3cb62e2
commit 6aaf0b7
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/src/UserInfo.cc b/src/UserInfo.cc
@@ -581,6 +581,11 @@ void UserInfo::handleAutoRegEvent(
   string nodePath(path);
   string userName = nodePath.substr(userInfo->zkAutoRegWatchDir_.size());
 
+  {
+    ScopeLock lock(userInfo->autoRegPendingUsersLock_);
+    userInfo->autoRegPendingUsers_.erase(userName);
+  }
+
   size_t sessions = userInfo->server_->autoRegCallback(userName);
 
   LOG(INFO) << "Auto Reg: User '" << userName << "' (" << sessions
@@ -594,13 +599,17 @@ bool UserInfo::tryAutoReg(
     userName = filterWorkerName(userName);
     regularUserName(userName);
 
-    if (autoRegPendingUsers_.find(userName) != autoRegPendingUsers_.end()) {
-      return true;
-    }
-    if (autoRegPendingUsers_.size() >= autoRegMaxPendingUsers_) {
-      LOG(INFO) << "UserInfo: too many pending registing request, user: "
-                << userName;
-      return false;
+    {
+      ScopeLock lock(autoRegPendingUsersLock_);
+      if (autoRegPendingUsers_.find(userName) != autoRegPendingUsers_.end()) {
+        return true;
+      }
+      if (autoRegPendingUsers_.size() >= autoRegMaxPendingUsers_) {
+        LOG(INFO) << "UserInfo: too many pending registing request, user: "
+                  << userName;
+        return false;
+      }
+      autoRegPendingUsers_.insert(userName);
     }
 
     string userInfo = Strings::Format(
@@ -620,5 +629,8 @@ bool UserInfo::tryAutoReg(
   } catch (...) {
     LOG(ERROR) << "UserInfo::tryAutoReg(): unknown exception";
   }
+
+  ScopeLock lock(autoRegPendingUsersLock_);
+  autoRegPendingUsers_.erase(userName);
   return false;
 }
diff --git a/src/UserInfo.h b/src/UserInfo.h
@@ -86,6 +86,7 @@ class UserInfo {
   string zkAutoRegWatchDir_;
   uint autoRegMaxPendingUsers_;
   std::set<string> autoRegPendingUsers_;
+  std::mutex autoRegPendingUsersLock_;
 
   pthread_rwlock_t nameChainlock_;
   // username -> chainId