mirrors/dragonfly
1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2025-05-11 10:25:47 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

feat(server): support for setnx command #281 (#283)

Browse source 
Signed-off-by: Boaz Sade <boaz1sade@gmail.com>
This commit is contained in:
Boaz Sade 2022-09-13 17:40:22 +03:00 • committed by GitHub
parent add252c301
commit d3359f1a0a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 93 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

68
src/server/string_family.cc
View file

@ -325,16 +325,23 @@ OpStatus SetCmd::Set(const SetParams& params, string_view key, string_view value
VLOG(2) << "Set " << key << "(" << db_slice.shard_id() << ") ";
if (params.how == SET_IF_EXISTS) {
auto [it, expire_it] = db_slice.FindExt(params.db_index, key);
if (IsValid(it)) { // existing
return SetExisting(params, it, expire_it, value);
if (params.IsConditionalSet()) {
const auto [it, expire_it] = db_slice.FindExt(params.db_index, key);
// Make sure that we have this key, and only add it if it does exists
if (params.how == SET_IF_EXISTS) {
if (IsValid(it)) {
return SetExisting(params, it, expire_it, value);
} else {
return OpStatus::SKIPPED;
}
} else {
if (IsValid(it)) { // if the policy is not to overide and have the key, just return
return OpStatus::SKIPPED;
}
}
return OpStatus::SKIPPED;
}
// At this point we either need to add missing entry, or we
// will override an existing one
// Trying to add a new entry.
tuple<PrimeIterator, ExpireIterator, bool> add_res;
try {
@ -347,7 +354,6 @@ OpStatus SetCmd::Set(const SetParams& params, string_view key, string_view value
if (!get<2>(add_res)) { // Existing.
return SetExisting(params, it, get<1>(add_res), value);
}
//
// Adding new value.
PrimeValue tvalue{value};
@ -480,14 +486,7 @@ void StringFamily::Set(CmdArgList args, ConnectionContext* cntx) {
}
}
DCHECK(cntx->transaction);
auto cb = [&](Transaction* t, EngineShard* shard) {
SetCmd sg(t->GetOpArgs(shard));
return sg.Set(sparams, key, value);
};
OpResult<void> result = cntx->transaction->ScheduleSingleHop(std::move(cb));
const auto result{SetGeneric(cntx, std::move(sparams), key, value)};
if (result == OpStatus::OK) {
return builder->SendStored();
}
@ -501,10 +500,44 @@ void StringFamily::Set(CmdArgList args, ConnectionContext* cntx) {
return builder->SendSetSkipped();
}
OpResult<void> StringFamily::SetGeneric(ConnectionContext* cntx, SetCmd::SetParams sparams,
std::string_view key, std::string_view value) {
DCHECK(cntx->transaction);
auto cb = [&](Transaction* t, EngineShard* shard) {
SetCmd sg(t->GetOpArgs(shard));
return sg.Set(sparams, key, value);
};
return cntx->transaction->ScheduleSingleHop(std::move(cb));
}
void StringFamily::SetEx(CmdArgList args, ConnectionContext* cntx) {
SetExGeneric(true, std::move(args), cntx);
}
void StringFamily::SetNx(CmdArgList args, ConnectionContext* cntx) {
// This is the same as calling the "Set" function, only in this case we are
// change the value only if the key does not exist. Otherwise the function
// will not modify it. in which case it would return 0
// it would return to the caller 1 in case the key did not exists and was added
string_view key = ArgS(args, 1);
string_view value = ArgS(args, 2);
SetCmd::SetParams sparams{cntx->db_index()};
sparams.how = SetCmd::SET_IF_NOTEXIST;
sparams.memcache_flags = cntx->conn_state.memcache_flag;
const auto results{SetGeneric(cntx, std::move(sparams), key, value)};
SinkReplyBuilder* builder = cntx->reply_builder();
if (results == OpStatus::OK) {
return builder->SendLong(1); // this means that we successfully set the value
}
if (results == OpStatus::OUT_OF_MEMORY) {
return builder->SendError(kOutOfMemory);
}
CHECK_EQ(results, OpStatus::SKIPPED); // in this case it must be skipped!
return builder->SendLong(0); // value do exists, we need to report that we didn't change it
}
void StringFamily::Get(CmdArgList args, ConnectionContext* cntx) {
get_qps.Inc();
@ -962,6 +995,7 @@ void StringFamily::Shutdown() {
void StringFamily::Register(CommandRegistry* registry) {
*registry << CI{"SET", CO::WRITE | CO::DENYOOM, -3, 1, 1, 1}.HFUNC(Set)
<< CI{"SETEX", CO::WRITE | CO::DENYOOM, 4, 1, 1, 1}.HFUNC(SetEx)
<< CI{"SETNX", CO::WRITE | CO::DENYOOM, 3, 1, 1, 1}.HFUNC(SetNx)
<< CI{"PSETEX", CO::WRITE | CO::DENYOOM, 4, 1, 1, 1}.HFUNC(PSetEx)
<< CI{"APPEND", CO::WRITE | CO::FAST, 3, 1, 1, 1}.HFUNC(Append)
<< CI{"PREPEND", CO::WRITE | CO::FAST, 3, 1, 1, 1}.HFUNC(Prepend)

11
src/server/string_family.h
View file

@ -19,7 +19,8 @@ class SetCmd {
const OpArgs op_args_;
public:
explicit SetCmd(const OpArgs& op_args) : op_args_(op_args) {}
explicit SetCmd(const OpArgs& op_args) : op_args_(op_args) {
}
enum SetHow { SET_ALWAYS, SET_IF_NOTEXIST, SET_IF_EXISTS };
@ -35,6 +36,10 @@ class SetCmd {
explicit SetParams(DbIndex dib) : db_index(dib) {
}
constexpr bool IsConditionalSet() const {
return how == SET_IF_NOTEXIST || how == SET_IF_EXISTS;
}
};
OpStatus Set(const SetParams& params, std::string_view key, std::string_view value);
@ -67,14 +72,18 @@ class StringFamily {
static void Set(CmdArgList args, ConnectionContext* cntx);
static void SetEx(CmdArgList args, ConnectionContext* cntx);
static void SetNx(CmdArgList args, ConnectionContext* cntx);
static void SetRange(CmdArgList args, ConnectionContext* cntx);
static void StrLen(CmdArgList args, ConnectionContext* cntx);
static void Prepend(CmdArgList args, ConnectionContext* cntx);
static void PSetEx(CmdArgList args, ConnectionContext* cntx);
// These functions are used internally, they do not implement any specific command
static void IncrByGeneric(std::string_view key, int64_t val, ConnectionContext* cntx);
static void ExtendGeneric(CmdArgList args, bool prepend, ConnectionContext* cntx);
static void SetExGeneric(bool seconds, CmdArgList args, ConnectionContext* cntx);
static OpResult<void> SetGeneric(ConnectionContext* cntx, SetCmd::SetParams sparams,
std::string_view key, std::string_view value);
struct GetResp {
std::string value;

37
src/server/string_family_test.cc
View file

@ -40,7 +40,6 @@ vector<int64_t> ToIntArr(const RespExpr& e) {
return res;
}
TEST_F(StringFamilyTest, SetGet) {
auto resp = Run({"set", "key", "val"});
@ -122,7 +121,7 @@ TEST_F(StringFamilyTest, SetHugeKey) {
}
TEST_F(StringFamilyTest, MGetSet) {
Run({"mset", "z", "0"}); // single key
Run({"mset", "z", "0"}); // single key
auto resp = Run({"mget", "z"}); // single key
EXPECT_THAT(resp, "0");
@ -188,7 +187,6 @@ TEST_F(StringFamilyTest, MSetGet) {
get_fb.join();
}
TEST_F(StringFamilyTest, MSetDel) {
auto mset_fb = pp_->at(0)->LaunchFiber([&] {
for (size_t i = 0; i < 1000; ++i) {
@ -310,8 +308,8 @@ TEST_F(StringFamilyTest, Range) {
EXPECT_EQ(Run({"getrange", "key3", "4", "5"}), "");
Run({"SET", "num", "1234"});
EXPECT_EQ(Run({"getrange","num", "3", "5000"}), "4");
EXPECT_EQ(Run({"getrange","num", "-5000", "10000"}), "1234");
EXPECT_EQ(Run({"getrange", "num", "3", "5000"}), "4");
EXPECT_EQ(Run({"getrange", "num", "-5000", "10000"}), "1234");
}
TEST_F(StringFamilyTest, IncrByFloat) {
@ -328,4 +326,33 @@ TEST_F(StringFamilyTest, IncrByFloat) {
EXPECT_EQ(resp, "3.566");
}
TEST_F(StringFamilyTest, SetNx) {
// Make sure that we "screen out" invalid parameters for this command
// this is important as it uses similar path as the "normal" set
auto resp = Run({"setnx", "foo", "bar", "XX"});
EXPECT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"setnx", "foo", "bar", "NX"});
ASSERT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"setnx", "foo", "bar", "xx"});
ASSERT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"setnx", "foo", "bar", "ex", "abc"});
ASSERT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"setnx", "foo", "bar", "ex", "-1"});
ASSERT_THAT(resp, ErrArg("wrong number of arguments"));
resp = Run({"setnx", "foo", "bar", "ex", "1"});
ASSERT_THAT(resp, ErrArg("wrong number of arguments"));
// now let see how it goes for the valid parameters
EXPECT_EQ(1, CheckedInt({"setnx", "foo", "bar"}));
EXPECT_EQ(Run({"get", "foo"}), "bar");
// second call to the same key should return 0 as we have it
EXPECT_EQ(0, CheckedInt({"setnx", "foo", "hello"}));
EXPECT_EQ(Run({"get", "foo"}), "bar"); // the value was not changed
}
} // namespace dfly