Skip to content

Commit

Permalink
chore: Add Lua force atomicity flag (#4523)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
chakaz authored and romange committed Jan 30, 2025
1 parent c056ccb commit 096fde1
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
25 changes: 25 additions & 0 deletions src/server/multi_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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
14 changes: 14 additions & 0 deletions src/server/script_mgr.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 096fde1

Please sign in to comment.