From 502efd80b2bd9e6b3606f7b5a57d9cb814ce7342 Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Fri, 27 Oct 2023 13:13:47 +0300
Subject: [PATCH] fix(server): Correctly set expiration from Lua scripts
 (#2080)

We used to set `time_now_ms_` only in the non-squashed execution path.

Fixes #2034
---
 src/server/main_service.cc |  9 +++++----
 src/server/multi_test.cc   | 13 +++++++++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index 3d5df1cbf..72f582b51 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -1685,16 +1685,17 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
     });
 
     ++ServerState::tlocal()->stats.eval_shardlocal_coordination_cnt;
-    boost::intrusive_ptr<Transaction> stub_tx = new Transaction{tx, *sid};
-    cntx->transaction = stub_tx.get();
-
     tx->PrepareMultiForScheduleSingleHop(*sid, tx->GetDbIndex(), args);
     tx->ScheduleSingleHop([&](Transaction*, EngineShard*) {
+      boost::intrusive_ptr<Transaction> stub_tx = new Transaction{tx, *sid};
+      cntx->transaction = stub_tx.get();
+
       result = interpreter->RunFunction(eval_args.sha, &error);
+
+      cntx->transaction = tx;
       return OpStatus::OK;
     });
 
-    cntx->transaction = tx;
     if (*sid != ServerState::tlocal()->thread_index()) {
       VLOG(1) << "Migrating connection " << cntx->conn() << " from " << ProactorBase::GetIndex()
               << " to " << *sid;
diff --git a/src/server/multi_test.cc b/src/server/multi_test.cc
index b0335c6b4..405d03353 100644
--- a/src/server/multi_test.cc
+++ b/src/server/multi_test.cc
@@ -876,6 +876,19 @@ TEST_F(MultiTest, TestLockedKeys) {
   EXPECT_FALSE(service_->IsLocked(0, "key2"));
 }
 
+TEST_F(MultiTest, EvalExpiration) {
+  // Make sure expiration is correctly set even from Lua scripts
+  if (auto config = absl::GetFlag(FLAGS_default_lua_flags); config != "") {
+    GTEST_SKIP() << "Skipped Eval test because default_lua_flags is set";
+    return;
+  }
+
+  absl::FlagSaver fs;
+  absl::SetFlag(&FLAGS_multi_exec_mode, Transaction::LOCK_AHEAD);
+  Run({"eval", "redis.call('set', 'x', 0, 'ex', 5, 'nx')", "1", "x"});
+  EXPECT_LE(CheckedInt({"pttl", "x"}), 5000);
+}
+
 class MultiEvalTest : public BaseFamilyTest {
  protected:
   MultiEvalTest() : BaseFamilyTest() {