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

Signed-off-by: Boaz Sade <boaz1sade@gmail.com>

src/server/string_family.cc

@@ -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) {
+  if (params.IsConditionalSet()) {
-    auto [it, expire_it] = db_slice.FindExt(params.db_index, key);
+    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 (IsValid(it)) {  // existing
+    if (params.how == SET_IF_EXISTS) {
-      return SetExisting(params, it, expire_it, value);
+      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);
+  const auto result{SetGeneric(cntx, std::move(sparams), key, value)};
-
-  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));
-
   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)

src/server/string_family.h

@@ -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;

src/server/string_family_test.cc

@@ -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", "3", "5000"}), "4");
-  EXPECT_EQ(Run({"getrange","num", "-5000", "10000"}), "1234");
+  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