diff --git a/src/base_cmd.h b/src/base_cmd.h index f387112e1..f988f88f3 100644 --- a/src/base_cmd.h +++ b/src/base_cmd.h @@ -24,6 +24,7 @@ namespace pikiwidb { // string cmd const std::string kCmdNameSet = "set"; const std::string kCmdNameGet = "get"; +const std::string kCmdNameIncrby = "incrby"; // multi const std::string kCmdNameMulti = "multi"; diff --git a/src/cmd_kv.cc b/src/cmd_kv.cc index 35fc0f532..b9f7aa970 100644 --- a/src/cmd_kv.cc +++ b/src/cmd_kv.cc @@ -163,7 +163,7 @@ BitCountCmd::BitCountCmd(const std::string& name, int16_t arity) bool BitCountCmd::DoInitial(PClient* client) { size_t paramSize = client->argv_.size(); - if(paramSize != 2 && paramSize != 4) { + if (paramSize != 2 && paramSize != 4) { client->SetRes(CmdRes::kSyntaxErr, kCmdNameBitCount); return false; } @@ -188,7 +188,7 @@ void BitCountCmd::DoCmd(PClient* client) { if (pstd::String2int(client->argv_[2], &start_offset_) == 0 || pstd::String2int(client->argv_[3], &end_offset_) == 0) { client->SetRes(CmdRes::kInvalidInt); - return ; + return; } auto str = GetDecodedString(value); @@ -215,4 +215,39 @@ void BitCountCmd::DoCmd(PClient* client) { client->AppendInteger(static_cast(count)); } +IncrbyCmd::IncrbyCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, CmdFlagsWrite, AclCategoryWrite | AclCategoryString) {} + +bool IncrbyCmd::DoInitial(PClient* client) { + int64_t by_ = 0; + if (!(pstd::String2int(client->argv_[2].data(), client->argv_[2].size(), &by_))) { + client->SetRes(CmdRes::kInvalidInt); + return false; + } + client->SetKey(client->argv_[1]); + return true; +} + +void IncrbyCmd::DoCmd(PClient* client) { + int64_t new_value_ = 0; + int64_t by_ = 0; + pstd::String2int(client->argv_[2].data(), client->argv_[2].size(), &by_); + PError err = PSTORE.Incrby(client->Key(), by_, &new_value_); + switch (err) { + case PError_type: + client->SetRes(CmdRes::kInvalidInt); + break; + case PError_notExist: // key not exist, set a new value + PSTORE.ClearExpire(client->Key()); // clear key's old ttl + PSTORE.SetValue(client->Key(), PObject::CreateString(by_)); + client->AppendInteger(by_); + break; + case PError_ok: + client->AppendInteger(new_value_); + break; + default: + client->SetRes(CmdRes::kErrOther, "incrby cmd error"); + break; + } +} } // namespace pikiwidb diff --git a/src/cmd_kv.h b/src/cmd_kv.h index 68432e7ae..e3ef2bab7 100644 --- a/src/cmd_kv.h +++ b/src/cmd_kv.h @@ -88,4 +88,15 @@ class BitCountCmd : public BaseCmd { void DoCmd(PClient *client) override; }; +class IncrbyCmd : public BaseCmd { + public: + IncrbyCmd(const std::string &name, int16_t arity); + + protected: + bool DoInitial(PClient *client) override; + + private: + void DoCmd(PClient *client) override; +}; + } // namespace pikiwidb diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc index 0a03d54d6..92215f1be 100644 --- a/src/cmd_table_manager.cc +++ b/src/cmd_table_manager.cc @@ -42,6 +42,9 @@ void CmdTableManager::InitCmdTable() { cmds_->insert(std::make_pair(kCmdNameMset, std::move(msetPtr))); std::unique_ptr bitcountPtr = std::make_unique(kCmdNameBitCount, -2); cmds_->insert(std::make_pair(kCmdNameBitCount, std::move(bitcountPtr))); + + std::unique_ptr incrbyPtr = std::make_unique(kCmdNameIncrby, 3); + cmds_->insert(std::make_pair(kCmdNameIncrby, std::move(incrbyPtr))); } std::pair CmdTableManager::GetCommand(const std::string& cmdName, PClient* client) { diff --git a/src/store.cc b/src/store.cc index 43c5a1c59..7379aa902 100644 --- a/src/store.cc +++ b/src/store.cc @@ -441,9 +441,9 @@ bool PStore::DeleteKey(const PString& key) { size_t ret = 0; // erase() from folly ConcurrentHashmap will throw an exception if hash function crashes try { - ret = db->erase(key); + ret = db->erase(key); } catch (const std::exception& e) { - return false; + return false; } return ret != 0; @@ -556,6 +556,39 @@ PObject* PStore::SetValue(const PString& key, PObject&& value) { return const_cast(&obj); } +PError PStore::Incrby(const PString& key, int64_t value, int64_t* ret) { + PObject* old_value = nullptr; + auto db = &dbs_[dbno_]; + + // shared when reading + std::unique_lock lock(mutex_); + PError err = getValueByType(key, old_value, PType_string); + if (err != PError_ok) { + return err; + } + char* end = nullptr; + auto str = pikiwidb::GetDecodedString(old_value); + int64_t ival = strtoll(str->c_str(), &end, 10); + if (*end != 0) { + // value is not a integer + return PError_type; + } + + PObject new_value; + *ret = ival + value; + new_value = PObject::CreateString((long)(*ret)); + new_value.lru = PObject::lruclock; + auto [realObj, status] = db->insert_or_assign(key, std::move(new_value)); + const PObject& obj = realObj->second; + + // put this key to sync list + if (!waitSyncKeys_.empty()) { + waitSyncKeys_[dbno_].insert_or_assign(key, &obj); + } + + return PError_ok; +} + void PStore::SetExpire(const PString& key, uint64_t when) const { expiredDBs_[dbno_].SetExpire(key, when); } int64_t PStore::TTL(const PString& key, uint64_t now) { return expiredDBs_[dbno_].TTL(key, now); } diff --git a/src/store.h b/src/store.h index 081d2cdcc..b16cb4cc3 100644 --- a/src/store.h +++ b/src/store.h @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include namespace pikiwidb { @@ -117,6 +119,8 @@ class PStore { PError GetValueByTypeNoTouch(const PString& key, PObject*& value, PType type = PType_invalid); PObject* SetValue(const PString& key, PObject&& value); + // incr + PError Incrby(const PString& key, int64_t value, int64_t* ret); // for expire key enum ExpireResult : std::int8_t { @@ -157,6 +161,8 @@ class PStore { private: PStore() : dbno_(0) {} + // mutex + mutable std::shared_mutex mutex_; PError getValueByType(const PString& key, PObject*& value, PType type = PType_invalid, bool touch = true); @@ -187,7 +193,7 @@ class PStore { size_t Size() const { return blockedClients_.size(); } private: - using Clients = std::list, uint64_t, ListPosition> >; + using Clients = std::list, uint64_t, ListPosition>>; using WaitingList = folly::ConcurrentHashMap>; WaitingList blockedClients_; @@ -199,9 +205,9 @@ class PStore { mutable std::vector dbs_; mutable std::vector expiredDBs_; std::vector blockedClients_; - std::vector > backends_; + std::vector> backends_; - using ToSyncDB = folly::ConcurrentHashMap >; + using ToSyncDB = folly::ConcurrentHashMap>; std::vector waitSyncKeys_; int dbno_ = -1; };