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

Author: chakaz

src/server/multi_test.cc

@@ -21,13 +21,15 @@ 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 {
 
 using namespace std;
 using namespace util;
 using absl::StrCat;
 using ::io::Result;
+using testing::_;
 using testing::ElementsAre;
 using testing::HasSubstr;
 
@@ -1219,4 +1221,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

src/server/script_mgr.cc

@@ -47,6 +47,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;
@@ -264,6 +271,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;