Add the read/write lock for the namespace

src/commands/cmd_server.cc

@@ -1275,7 +1275,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loadin
                         MakeCmdAttr<CommandInfo>("info", -1, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandRole>("role", 1, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandConfig>("config", -2, "read-only", 0, 0, 0, GenerateConfigFlag),
-                        MakeCmdAttr<CommandNamespace>("namespace", -3, "read-only exclusive", 0, 0, 0),
+                        MakeCmdAttr<CommandNamespace>("namespace", -3, "write", 0, 0, 0),
                         MakeCmdAttr<CommandKeys>("keys", 2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFlushDB>("flushdb", 1, "write no-dbsize-check", 0, 0, 0),
                         MakeCmdAttr<CommandFlushAll>("flushall", 1, "write no-dbsize-check", 0, 0, 0),

src/server/namespace.cc

@@ -79,6 +79,7 @@ Status Namespace::LoadAndRewrite() {
     return {Status::NotOK, "cannot switch off repl_namespace_enabled when namespaces exist in db"};
   }
 
+  std::unique_lock<std::shared_mutex> lock(tokens_mu_);
   // Load from the configuration file first
   tokens_ = config->load_tokens;
   // Merge the tokens from the database if the token is not in the configuration file
@@ -92,10 +93,11 @@ Status Namespace::LoadAndRewrite() {
   // is enabled. So we don't need to do that if no tokens are loaded or the namespace replication is disabled.
   if (config->load_tokens.empty() || !config->repl_namespace_enabled) return Status::OK();
 
-  return Rewrite();
+  return Rewrite(tokens_);
 }
 
-StatusOr<std::string> Namespace::Get(const std::string& ns) const {
+StatusOr<std::string> Namespace::Get(const std::string& ns) {
+  std::shared_lock lock(tokens_mu_);
   for (const auto& iter : tokens_) {
     if (iter.second == ns) {
       return iter.first;
@@ -104,7 +106,8 @@ StatusOr<std::string> Namespace::Get(const std::string& ns) const {
   return {Status::NotFound};
 }
 
-StatusOr<std::string> Namespace::GetByToken(const std::string& token) const {
+StatusOr<std::string> Namespace::GetByToken(const std::string& token) {
+  std::shared_lock lock(tokens_mu_);
   auto iter = tokens_.find(token);
   if (iter == tokens_.end()) {
     return {Status::NotFound};
@@ -132,6 +135,7 @@ Status Namespace::Set(const std::string& ns, const std::string& token) {
     return {Status::NotOK, kErrInvalidToken};
   }
 
+  std::unique_lock lock(tokens_mu_);
   for (const auto& iter : tokens_) {
     if (iter.second == ns) {  // need to delete the old token first
       tokens_.erase(iter.first);
@@ -140,7 +144,7 @@ Status Namespace::Set(const std::string& ns, const std::string& token) {
   }
   tokens_[token] = ns;
 
-  s = Rewrite();
+  s = Rewrite(tokens_);
   if (!s.IsOK()) {
     tokens_.erase(token);
     return s;
@@ -149,17 +153,22 @@ Status Namespace::Set(const std::string& ns, const std::string& token) {
 }
 
 Status Namespace::Add(const std::string& ns, const std::string& token) {
-  // duplicate namespace
-  for (const auto& iter : tokens_) {
-    if (iter.second == ns) {
-      if (iter.first == token) return Status::OK();
-      return {Status::NotOK, kErrNamespaceExists};
+  {
+    std::shared_lock lock(tokens_mu_);
+    // duplicate namespace
+    for (const auto& iter : tokens_) {
+      if (iter.second == ns) {
+        if (iter.first == token) return Status::OK();
+        return {Status::NotOK, kErrNamespaceExists};
+      }
+    }
+    // duplicate token
+    if (tokens_.find(token) != tokens_.end()) {
+      return {Status::NotOK, kErrTokenExists};
     }
   }
-  // duplicate token
-  if (tokens_.find(token) != tokens_.end()) {
-    return {Status::NotOK, kErrTokenExists};
-  }
+
+  // we don't need to lock the mutex here because the Set method will lock it
   return Set(ns, token);
 }
 
@@ -171,10 +180,11 @@ Status Namespace::Del(const std::string& ns) {
     return {Status::NotOK, kErrCantModifyNamespace};
   }
 
+  std::unique_lock lock(tokens_mu_);
   for (const auto& iter : tokens_) {
     if (iter.second == ns) {
       tokens_.erase(iter.first);
-      auto s = Rewrite();
+      auto s = Rewrite(tokens_);
       if (!s.IsOK()) {
         tokens_[iter.first] = iter.second;
         return s;
@@ -185,11 +195,11 @@ Status Namespace::Del(const std::string& ns) {
   return {Status::NotOK, kErrNamespaceNotFound};
 }
 
-Status Namespace::Rewrite() {
+Status Namespace::Rewrite(const std::map<std::string, std::string>& tokens) const {
   auto config = storage_->GetConfig();
   // Rewrite the configuration file only if it's running with the configuration file
   if (config->HasConfigFile()) {
-    auto s = config->Rewrite(tokens_);
+    auto s = config->Rewrite(tokens);
     if (!s.IsOK()) {
       return s;
     }
@@ -206,7 +216,7 @@ Status Namespace::Rewrite() {
     return Status::OK();
   }
   jsoncons::json json;
-  for (const auto& iter : tokens_) {
+  for (const auto& iter : tokens) {
     json[iter.first] = iter.second;
   }
   return storage_->WriteToPropagateCF(kNamespaceDBKey, json.to_string());

src/server/namespace.h

@@ -34,18 +34,20 @@ class Namespace {
   Namespace &operator=(const Namespace &) = delete;
 
   Status LoadAndRewrite();
-  StatusOr<std::string> Get(const std::string &ns) const;
-  StatusOr<std::string> GetByToken(const std::string &token) const;
+  StatusOr<std::string> Get(const std::string &ns);
+  StatusOr<std::string> GetByToken(const std::string &token);
   Status Set(const std::string &ns, const std::string &token);
   Status Add(const std::string &ns, const std::string &token);
   Status Del(const std::string &ns);
   const std::map<std::string, std::string> &List() const { return tokens_; }
-  Status Rewrite();
+  Status Rewrite(const std::map<std::string, std::string> &tokens) const;
   bool IsAllowModify() const;
 
  private:
   engine::Storage *storage_;
   rocksdb::ColumnFamilyHandle *cf_ = nullptr;
+
+  std::shared_mutex tokens_mu_;
   std::map<std::string, std::string> tokens_;
 
   Status loadFromDB(std::map<std::string, std::string> *db_tokens) const;