1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2024-12-14 11:58:02 +00:00

feat: Allow pre-declaring Lua SHAs to run with undeclared keys (#3465)

* feat: Allow pre-declaring Lua SHAs to run with undeclared keys

By using `--lua_undeclared_keys_shas=SHA,SHA,SHA` users can now specify
which scripts should run globally (undeclared keys) without explicit
support from the scripts themselves.

Fixes #2442
This commit is contained in:
Shahar Mike 2024-08-08 10:27:27 +03:00 committed by GitHub
parent 1b1a83dc58
commit aa424c81af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 47 additions and 0 deletions

View file

@ -380,6 +380,7 @@ TEST_F(MultiTest, Eval) {
GTEST_SKIP() << "Skipped Eval test because default_lua_flags is set";
return;
}
absl::FlagSaver saver;
absl::SetFlag(&FLAGS_lua_allow_undeclared_auto_correct, true);
RespExpr resp;
@ -762,6 +763,37 @@ TEST_F(MultiTest, ScriptFlagsEmbedded) {
EXPECT_THAT(Run({"eval", s2, "0"}), ErrArg("Invalid flag: this-is-an-error"));
}
TEST_F(MultiTest, UndeclaredKeyFlag) {
if (auto mode = absl::GetFlag(FLAGS_multi_exec_mode); mode != Transaction::LOCK_AHEAD) {
GTEST_SKIP() << "Skipped test because multi_exec_mode is not default";
return;
}
absl::FlagSaver fs; // lua_undeclared_keys_shas changed via CONFIG cmd below
const char* script = "return redis.call('GET', 'random-key');";
Run({"set", "random-key", "works"});
// Get SHA for script in a persistent way
string sha = Run({"script", "load", script}).GetString();
// Make sure we can't run the script before setting the flag
EXPECT_THAT(Run({"evalsha", sha, "0"}), ErrArg("undeclared"));
EXPECT_THAT(Run({"eval", script, "0"}), ErrArg("undeclared"));
// Clear all Lua scripts so we can configure the cache
EXPECT_THAT(Run({"script", "flush"}), "OK");
EXPECT_THAT(Run({"script", "exists", sha}), IntArg(0));
EXPECT_THAT(
Run({"config", "set", "lua_undeclared_keys_shas", absl::StrCat(sha, ",NON-EXISTING-HASH")}),
"OK");
// Check eval finds script flags.
EXPECT_EQ(Run({"eval", script, "0"}), "works");
EXPECT_EQ(Run({"evalsha", sha, "0"}), "works");
}
// todo: ASAN fails heres on arm
#ifndef SANITIZERS
TEST_F(MultiTest, ScriptBadCommand) {

View file

@ -38,6 +38,15 @@ ABSL_FLAG(bool, lua_allow_undeclared_auto_correct, false,
"access undeclared keys, automaticaly set the script flag to be able to run with "
"undeclared key.");
ABSL_FLAG(
std::vector<std::string>, lua_undeclared_keys_shas,
std::vector<std::string>({
"351130589c64523cb98978dc32c64173a31244f3", // Sidekiq, see #2442
"6ae15ef4678593dc61f991c9953722d67d822776", // Sidekiq, see #2442
}),
"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.");
namespace dfly {
using namespace std;
using namespace facade;
@ -262,6 +271,11 @@ io::Result<string, GenericError> ScriptMgr::Insert(string_view body, Interpreter
return params_opt.get_unexpected();
auto params = params_opt->value_or(default_params_);
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;
}
// If the script is atomic, check for possible squashing optimizations.
// For non atomic modes, squashing increases the time locks are held, which
// can decrease throughput with frequently accessed keys.

View file

@ -824,6 +824,7 @@ void ServerFamily::Init(util::AcceptServer* acceptor, std::vector<facade::Listen
config_registry.RegisterMutable("tls_ca_cert_file");
config_registry.RegisterMutable("tls_ca_cert_dir");
config_registry.RegisterMutable("replica_priority");
config_registry.RegisterMutable("lua_undeclared_keys_shas");
pb_task_ = shard_set->pool()->GetNextProactor();
if (pb_task_->GetKind() == ProactorBase::EPOLL) {