From ffce1cb79646af08a256814da5d1a26ec371ce84 Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Wed, 29 Jan 2025 16:35:39 +0200
Subject: [PATCH] chore: Add Lua force atomicity flag (#4523)

* chore: Add Lua force atomicity flag

We accidentally instructed Sidekiq to add `disable-atomicity` to their
script, despite them needing to run atomically.

This hack-ish PR adds a `--lua_force_atomicity_shas` flag to allow
specifying which SHAs are forced to run in an atomic fashion, even if
they are marked as non-atomic.

Fixes #4522

* fix build on clang
---
 src/server/cluster/cluster_family.cc |  2 +-
 src/server/multi_test.cc             | 25 +++++++++++++++++++++++++
 src/server/script_mgr.cc             | 14 ++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/server/cluster/cluster_family.cc b/src/server/cluster/cluster_family.cc
index 5d70cf9b4..241b63cde 100644
--- a/src/server/cluster/cluster_family.cc
+++ b/src/server/cluster/cluster_family.cc
@@ -900,7 +900,7 @@ void ClusterFamily::InitMigration(CmdArgList args, SinkReplyBuilder* builder) {
 
   const auto& incoming_migrations = cluster_config()->GetIncomingMigrations();
   bool found = any_of(incoming_migrations.begin(), incoming_migrations.end(),
-                      [&source_id, &slot_ranges](const MigrationInfo& info) {
+                      [source_id = source_id, &slot_ranges](const MigrationInfo& info) {
                         return info.node_info.id == source_id && info.slot_ranges == slot_ranges;
                       });
   if (!found) {
diff --git a/src/server/multi_test.cc b/src/server/multi_test.cc
index 24be97d9e..d61fc003d 100644
--- a/src/server/multi_test.cc
+++ b/src/server/multi_test.cc
@@ -20,6 +20,7 @@ ABSL_DECLARE_FLAG(bool, multi_exec_squash);
 ABSL_DECLARE_FLAG(bool, lua_auto_async);
 ABSL_DECLARE_FLAG(bool, lua_allow_undeclared_auto_correct);
 ABSL_DECLARE_FLAG(std::string, default_lua_flags);
+ABSL_DECLARE_FLAG(std::vector<std::string>, lua_force_atomicity_shas);
 
 namespace dfly {
 
@@ -27,6 +28,7 @@ using namespace std;
 using namespace util;
 using absl::StrCat;
 using ::io::Result;
+using testing::_;
 using testing::ElementsAre;
 using testing::HasSubstr;
 
@@ -1148,4 +1150,27 @@ TEST_F(MultiTest, EvalShaRo) {
   EXPECT_THAT(resp, ErrArg("Write commands are not allowed from read-only scripts"));
 }
 
+TEST_F(MultiTest, ForceAtomicityFlag) {
+  absl::FlagSaver fs;
+
+  const string kHash = "bb855c2ecfa3114d222cb11e0682af6360e9712f";
+  const string_view kScript = R"(
+    --!df flags=disable-atomicity
+    redis.call('get', 'x');
+    return "OK";
+  )";
+
+  // EVAL the script works due to disable-atomicity flag
+  EXPECT_EQ(Run({"eval", kScript, "0"}), "OK");
+
+  EXPECT_THAT(Run({"script", "list"}), RespElementsAre(kHash, kScript));
+
+  // Flush scripts to force re-evaluating of flags
+  EXPECT_EQ(Run({"script", "flush"}), "OK");
+
+  // Now it doesn't work, because we force atomicity
+  absl::SetFlag(&FLAGS_lua_force_atomicity_shas, {kHash});
+  EXPECT_THAT(Run({"eval", kScript, "0"}), ErrArg("undeclared"));
+}
+
 }  // namespace dfly
diff --git a/src/server/script_mgr.cc b/src/server/script_mgr.cc
index ee29617a7..d261c740e 100644
--- a/src/server/script_mgr.cc
+++ b/src/server/script_mgr.cc
@@ -48,6 +48,13 @@ ABSL_FLAG(
     "Comma-separated list of Lua script SHAs which are allowed to access undeclared keys. SHAs are "
     "only looked at when loading the script, and new values do not affect already-loaded script.");
 
+ABSL_FLAG(std::vector<std::string>, lua_force_atomicity_shas,
+          std::vector<std::string>({
+              "f8133be7f04abd9dfefa83c3b29a9d837cfbda86",  // Sidekiq, see #4522
+          }),
+          "Comma-separated list of Lua script SHAs which are forced to run in atomic mode, even if "
+          "the script specifies disable-atomicity.");
+
 namespace dfly {
 using namespace std;
 using namespace facade;
@@ -265,6 +272,13 @@ io::Result<string, GenericError> ScriptMgr::Insert(string_view body, Interpreter
     return params_opt.get_unexpected();
   auto params = params_opt->value_or(default_params_);
 
+  if (!params.atomic) {
+    auto force_atomic_shas = absl::GetFlag(FLAGS_lua_force_atomicity_shas);
+    if (find(force_atomic_shas.begin(), force_atomic_shas.end(), sha) != force_atomic_shas.end()) {
+      params.atomic = true;
+    }
+  }
+
   auto undeclared_shas = absl::GetFlag(FLAGS_lua_undeclared_keys_shas);
   if (find(undeclared_shas.begin(), undeclared_shas.end(), sha) != undeclared_shas.end()) {
     params.undeclared_keys = true;