feat(server): adding support for EXAT PXAT option at set command (#306)

feat(server): EXAT PXAT support changes according to code review #284

Signed-off-by: Boaz Sade <boaz@dragonflydb.io>
Co-authored-by: Boaz Sade <boaz@dragonflydb.io>

src/server/string_family.cc

@@ -11,7 +11,10 @@ extern "C" {
 #include <absl/container/inlined_vector.h>
 #include <double-conversion/string-to-double.h>
 
+#include <chrono>
+
 #include "base/logging.h"
+#include "redis/util.h"
 #include "server/command_registry.h"
 #include "server/conn_context.h"
 #include "server/engine_shard_set.h"
@@ -138,7 +141,8 @@ OpResult<string> OpGetRange(const OpArgs& op_args, string_view key, int32_t star
   return string(slice.substr(start, end - start + 1));
 };
 
-size_t ExtendExisting(const OpArgs& op_args, PrimeIterator it, string_view key, string_view val, bool prepend) {
+size_t ExtendExisting(const OpArgs& op_args, PrimeIterator it, string_view key, string_view val,
+                      bool prepend) {
   string tmp, new_val;
   auto* shard = op_args.shard;
   string_view slice = GetSlice(shard, it->second, &tmp);
@@ -296,6 +300,19 @@ OpResult<int64_t> OpIncrBy(const OpArgs& op_args, std::string_view key, int64_t
   return new_val;
 }
 
+int64_t CalculateAbsTime(int64_t unix_time, bool as_milli) {
+  using std::chrono::duration_cast;
+  using std::chrono::milliseconds;
+  using std::chrono::seconds;
+  using std::chrono::system_clock;
+
+  if (as_milli) {
+    return unix_time - duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
+  } else {
+    return unix_time - duration_cast<seconds>(system_clock::now().time_since_epoch()).count();
+  }
+}
+
 // Returns true if keys were set, false otherwise.
 OpStatus OpMSet(const OpArgs& op_args, ArgSlice args) {
   DCHECK(!args.empty() && args.size() % 2 == 0);
@@ -452,8 +469,8 @@ void StringFamily::Set(CmdArgList args, ConnectionContext* cntx) {
 
     string_view cur_arg = ArgS(args, i);
 
-    if (cur_arg == "EX" || cur_arg == "PX") {
-      bool is_ms = (cur_arg == "PX");
+    if (cur_arg == "EX" || cur_arg == "PX" || cur_arg == "EXAT" || cur_arg == "PXAT") {
+      bool is_ms = (cur_arg == "PX" || cur_arg == "PXAT");
       ++i;
       if (i == args.size()) {
         builder->SendError(kSyntaxErr);
@@ -464,7 +481,21 @@ void StringFamily::Set(CmdArgList args, ConnectionContext* cntx) {
         return builder->SendError(kInvalidIntErr);
       }
 
-      if (int_arg <= 0 || (!is_ms && int_arg >= kMaxExpireDeadlineSec)) {
+      // Since PXAT/EXAT can change this, we need to check this ahead
+      if (int_arg <= 0) {
+        return builder->SendError(InvalidExpireTime("set"));
+      }
+      // for []AT we need to take expiration time as absolute from the value given
+      // check here and if the time is in the past, return OK but don't set it
+      // Note that the time pass here for PXAT is in milliseconds, we must not change it!
+      if (cur_arg == "EXAT" || cur_arg == "PXAT") {
+        int_arg = CalculateAbsTime(int_arg, is_ms);
+        if (int_arg < 0) {
+          // this happened in the past, just return, for some reason Redis reports OK in this case
+          return builder->SendStored();
+        }
+      }
+      if (!is_ms && int_arg >= kMaxExpireDeadlineSec) {
         return builder->SendError(InvalidExpireTime("set"));
       }
 
@@ -517,7 +548,7 @@ void StringFamily::SetEx(CmdArgList args, ConnectionContext* 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 
+  // 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);

src/server/string_family_test.cc

@@ -355,4 +355,35 @@ TEST_F(StringFamilyTest, SetNx) {
   EXPECT_EQ(Run({"get", "foo"}), "bar");  // the value was not changed
 }
 
+TEST_F(StringFamilyTest, SetPxAtExAt) {
+  using std::chrono::duration_cast;
+  using std::chrono::milliseconds;
+  using std::chrono::seconds;
+  using std::chrono::system_clock;
+
+  // Expiration time as set at unix time
+  auto resp = Run({"set", "foo", "bar", "EXAT", "-1"});
+  ASSERT_THAT(resp, ErrArg("invalid expire time"));
+  resp = Run({"set", "foo", "bar", "EXAT", std::to_string(time(nullptr) - 1)});
+  ASSERT_THAT(resp, "OK");  // it would return OK but will not set the value - expiration time is 0
+                            // (checked with Redis)
+  EXPECT_EQ(Run({"get", "foo"}).type, facade::RespExpr::NIL);
+
+  resp = Run({"set", "foo", "bar", "PXAT", "-1"});
+  ASSERT_THAT(resp, ErrArg("invalid expire time"));
+
+  auto now = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
+  resp = Run({"set", "foo", "bar", "PXAT", std::to_string(now - 23)});
+  ASSERT_THAT(resp, "OK");  // it would return OK but will not set the value (checked with Redis)
+  EXPECT_EQ(Run({"get", "foo"}).type, facade::RespExpr::NIL);
+
+  resp = Run({"set", "foo", "bar", "EXAT", std::to_string(time(nullptr) + 1)});
+  ASSERT_THAT(resp, "OK");  // valid expiration time
+  EXPECT_EQ(Run({"get", "foo"}), "bar");
+
+  resp = Run({"set", "foo2", "abc", "PXAT", std::to_string(now + 300)});
+  ASSERT_THAT(resp, "OK");
+  EXPECT_EQ(Run({"get", "foo2"}), "abc");
+}
+
 }  // namespace dfly