mirror of
https://github.com/dragonflydb/dragonfly.git
synced 2025-05-13 03:15:44 +02:00
Support script configuration (#889)
Store script parameters for each script that allow configuring it's transactions multi mode. They can be configured either for a specific scripts with `SCRIPT CONIFG <sha> [params...]` or changed globally as defaults with `default_lua_config`. The current supported options are `allow-undeclared-keys` and `disable-atomicity`. Based on those flags, we determine the correct multi mode. `disable-atomicity` allow running in non-atomic mode, whereas being atomic and enabling `allow-undeclared-keys` requires the global mode.
This commit is contained in:
Vladislav
GitHub
parent
eb5fd2867f
commit
edbd43a3b3
9 changed files with 348 additions and 185 deletions
|
|
@ -5,26 +5,55 @@
|
|||
#include "server/script_mgr.h"
|
||||
|
||||
#include <absl/cleanup/cleanup.h>
|
||||
#include <absl/strings/ascii.h>
|
||||
#include <absl/strings/match.h>
|
||||
#include <absl/strings/numbers.h>
|
||||
#include <absl/strings/str_cat.h>
|
||||
#include <absl/strings/str_split.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "base/flags.h"
|
||||
#include "base/logging.h"
|
||||
#include "core/interpreter.h"
|
||||
#include "facade/error.h"
|
||||
#include "server/engine_shard_set.h"
|
||||
#include "server/server_state.h"
|
||||
|
||||
ABSL_FLAG(std::string, default_lua_config, "",
|
||||
"Configure default mode for running Lua scripts: \n - Use 'allow-undeclared-keys' to "
|
||||
"allow accessing undeclared keys, \n - Use 'disable-atomicity' to allow "
|
||||
"running scripts non-atomically. \nSpecify multiple values "
|
||||
"separated by space, for example 'allow-undeclared-keys disable-atomicity' runs scripts "
|
||||
"non-atomically and allows accessing undeclared keys");
|
||||
|
||||
namespace dfly {
|
||||
|
||||
using namespace std;
|
||||
using namespace facade;
|
||||
|
||||
ScriptMgr::ScriptMgr() {
|
||||
// Build default script config
|
||||
std::string config = absl::GetFlag(FLAGS_default_lua_config);
|
||||
|
||||
static_assert(ScriptParams{}.atomic && !ScriptParams{}.undeclared_keys);
|
||||
|
||||
auto parts = absl::StrSplit(config, absl::ByAnyChar(",; "), absl::SkipEmpty());
|
||||
for (auto pragma : parts) {
|
||||
CHECK(ScriptParams::ApplyPragma(pragma, &default_params_))
|
||||
<< "Bad format of default_lua_config flag";
|
||||
}
|
||||
}
|
||||
|
||||
ScriptMgr::ScriptKey::ScriptKey(string_view sha) : array{} {
|
||||
DCHECK_EQ(sha.size(), size());
|
||||
memcpy(data(), sha.data(), size());
|
||||
}
|
||||
|
||||
void ScriptMgr::Run(CmdArgList args, ConnectionContext* cntx) {
|
||||
string_view subcmd = ArgS(args, 0);
|
||||
|
||||
if (args.size() == 1 && subcmd == "HELP") {
|
||||
if (subcmd == "HELP") {
|
||||
string_view kHelp[] = {
|
||||
"SCRIPT <subcommand> [<arg> [value] [opt] ...]",
|
||||
"Subcommands are:",
|
||||
|
|
@ -32,6 +61,11 @@ void ScriptMgr::Run(CmdArgList args, ConnectionContext* cntx) {
|
|||
" Return information about the existence of the scripts in the script cache.",
|
||||
"LOAD <script>",
|
||||
" Load a script into the scripts cache without executing it.",
|
||||
"CONFIGURE <sha> [options ...]",
|
||||
" The following options are possible: ",
|
||||
" - Use 'allow-undeclared-keys' to allow accessing undeclared keys",
|
||||
" - Use 'disable-atomicity' to allow running scripts non-atomically to improve "
|
||||
"performance",
|
||||
"LIST",
|
||||
" Lists loaded scripts.",
|
||||
"LATENCY",
|
||||
|
|
@ -41,110 +75,87 @@ void ScriptMgr::Run(CmdArgList args, ConnectionContext* cntx) {
|
|||
return (*cntx)->SendSimpleStrArr(kHelp, ABSL_ARRAYSIZE(kHelp));
|
||||
}
|
||||
|
||||
if (subcmd == "EXISTS") {
|
||||
vector<uint8_t> res(args.size() - 1, 0);
|
||||
for (size_t i = 1; i < args.size(); ++i) {
|
||||
string_view sha = ArgS(args, i);
|
||||
if (Find(sha)) {
|
||||
res[i - 1] = 1;
|
||||
}
|
||||
}
|
||||
if (subcmd == "EXISTS" && args.size() > 1)
|
||||
return ExistsCmd(args, cntx);
|
||||
|
||||
(*cntx)->StartArray(res.size());
|
||||
for (uint8_t v : res) {
|
||||
(*cntx)->SendLong(v);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (subcmd == "LIST")
|
||||
return ListCmd(cntx);
|
||||
|
||||
if (subcmd == "LIST") {
|
||||
return ListScripts(cntx);
|
||||
}
|
||||
if (subcmd == "LATENCY")
|
||||
return LatencyCmd(cntx);
|
||||
|
||||
if (subcmd == "LATENCY") {
|
||||
return PrintLatency(cntx);
|
||||
}
|
||||
if (subcmd == "LOAD" && args.size() == 2)
|
||||
return LoadCmd(args, cntx);
|
||||
|
||||
if (subcmd == "LOAD" && args.size() == 2) {
|
||||
string_view body = ArgS(args, 1);
|
||||
|
||||
if (body.empty()) {
|
||||
char sha[41];
|
||||
Interpreter::FuncSha1(body, sha);
|
||||
return (*cntx)->SendBulkString(sha);
|
||||
}
|
||||
|
||||
ServerState* ss = ServerState::tlocal();
|
||||
auto interpreter = ss->BorrowInterpreter();
|
||||
absl::Cleanup clean = [ss, interpreter]() { ss->ReturnInterpreter(interpreter); };
|
||||
|
||||
// no need to lock the interpreter since we do not mess the stack.
|
||||
string error_or_id;
|
||||
Interpreter::AddResult add_result = interpreter->AddFunction(body, &error_or_id);
|
||||
|
||||
if (add_result == Interpreter::ALREADY_EXISTS) {
|
||||
return (*cntx)->SendBulkString(error_or_id);
|
||||
}
|
||||
if (add_result == Interpreter::COMPILE_ERR) {
|
||||
return (*cntx)->SendError(error_or_id);
|
||||
}
|
||||
|
||||
InsertFunction(error_or_id, body);
|
||||
|
||||
return (*cntx)->SendBulkString(error_or_id);
|
||||
}
|
||||
if (subcmd == "CONFIG" && args.size() > 2)
|
||||
return ConfigCmd(args, cntx);
|
||||
|
||||
string err = absl::StrCat("Unknown subcommand or wrong number of arguments for '", subcmd,
|
||||
"'. Try SCRIPT HELP.");
|
||||
cntx->reply_builder()->SendError(err, kSyntaxErrType);
|
||||
}
|
||||
|
||||
bool ScriptMgr::InsertFunction(std::string_view id, std::string_view body) {
|
||||
ScriptKey key;
|
||||
CHECK_EQ(key.size(), id.size());
|
||||
memcpy(key.data(), id.data(), key.size());
|
||||
|
||||
lock_guard lk(mu_);
|
||||
auto [it, inserted] = db_.emplace(key, nullptr);
|
||||
if (inserted) {
|
||||
it->second.reset(new char[body.size() + 1]);
|
||||
memcpy(it->second.get(), body.data(), body.size());
|
||||
it->second[body.size()] = '\0';
|
||||
}
|
||||
return inserted;
|
||||
}
|
||||
|
||||
const char* ScriptMgr::Find(std::string_view sha) const {
|
||||
if (sha.size() != 40)
|
||||
return nullptr;
|
||||
|
||||
ScriptKey key;
|
||||
memcpy(key.data(), sha.data(), key.size());
|
||||
|
||||
lock_guard lk(mu_);
|
||||
auto it = db_.find(key);
|
||||
if (it == db_.end())
|
||||
return nullptr;
|
||||
|
||||
return it->second.get();
|
||||
}
|
||||
|
||||
vector<pair<string, string>> ScriptMgr::GetLuaScripts() const {
|
||||
vector<pair<string, string>> res;
|
||||
|
||||
lock_guard lk(mu_);
|
||||
res.reserve(db_.size());
|
||||
for (const auto& k_v : db_) {
|
||||
string key(k_v.first.data(), k_v.first.size());
|
||||
|
||||
res.emplace_back(move(key), k_v.second.get());
|
||||
void ScriptMgr::ExistsCmd(CmdArgList args, ConnectionContext* cntx) const {
|
||||
vector<uint8_t> res(args.size() - 1, 0);
|
||||
for (size_t i = 1; i < args.size(); ++i) {
|
||||
if (string_view sha = ArgS(args, i); Find(sha)) {
|
||||
res[i - 1] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
return res;
|
||||
(*cntx)->StartArray(res.size());
|
||||
for (uint8_t v : res) {
|
||||
(*cntx)->SendLong(v);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
void ScriptMgr::ListScripts(ConnectionContext* cntx) {
|
||||
vector<pair<string, string>> scripts = GetLuaScripts();
|
||||
void ScriptMgr::LoadCmd(CmdArgList args, ConnectionContext* cntx) {
|
||||
string_view body = ArgS(args, 1);
|
||||
|
||||
if (body.empty()) {
|
||||
char sha[41];
|
||||
Interpreter::FuncSha1(body, sha);
|
||||
return (*cntx)->SendBulkString(sha);
|
||||
}
|
||||
|
||||
ServerState* ss = ServerState::tlocal();
|
||||
auto interpreter = ss->BorrowInterpreter();
|
||||
absl::Cleanup clean = [ss, interpreter]() { ss->ReturnInterpreter(interpreter); };
|
||||
|
||||
// no need to lock the interpreter since we do not mess the stack.
|
||||
string error_or_id;
|
||||
Interpreter::AddResult add_result = interpreter->AddFunction(body, &error_or_id);
|
||||
|
||||
if (add_result == Interpreter::ALREADY_EXISTS) {
|
||||
return (*cntx)->SendBulkString(error_or_id);
|
||||
}
|
||||
if (add_result == Interpreter::COMPILE_ERR) {
|
||||
return (*cntx)->SendError(error_or_id);
|
||||
}
|
||||
|
||||
Insert(error_or_id, body);
|
||||
|
||||
return (*cntx)->SendBulkString(error_or_id);
|
||||
}
|
||||
|
||||
void ScriptMgr::ConfigCmd(CmdArgList args, ConnectionContext* cntx) {
|
||||
lock_guard lk{mu_};
|
||||
ScriptKey key{ArgS(args, 1)};
|
||||
auto& data = db_[key];
|
||||
|
||||
for (auto pragma : args.subspan(2)) {
|
||||
if (!ScriptParams::ApplyPragma(facade::ToSV(pragma), &data))
|
||||
return (*cntx)->SendError("Invalid config format");
|
||||
}
|
||||
|
||||
UpdateScriptCaches(key, data);
|
||||
|
||||
return (*cntx)->SendOk();
|
||||
}
|
||||
|
||||
void ScriptMgr::ListCmd(ConnectionContext* cntx) const {
|
||||
vector<pair<string, string>> scripts = GetAll();
|
||||
(*cntx)->StartArray(scripts.size());
|
||||
for (const auto& k_v : scripts) {
|
||||
(*cntx)->StartArray(2);
|
||||
|
|
@ -153,7 +164,7 @@ void ScriptMgr::ListScripts(ConnectionContext* cntx) {
|
|||
}
|
||||
}
|
||||
|
||||
void ScriptMgr::PrintLatency(ConnectionContext* cntx) {
|
||||
void ScriptMgr::LatencyCmd(ConnectionContext* cntx) const {
|
||||
absl::flat_hash_map<std::string, base::Histogram> result;
|
||||
boost::fibers::mutex mu;
|
||||
|
||||
|
|
@ -174,4 +185,66 @@ void ScriptMgr::PrintLatency(ConnectionContext* cntx) {
|
|||
}
|
||||
}
|
||||
|
||||
bool ScriptMgr::Insert(std::string_view id, std::string_view body) {
|
||||
bool updated = false;
|
||||
|
||||
lock_guard lk{mu_};
|
||||
auto [it, _] = db_.emplace(id, InternalScriptData{default_params_, nullptr});
|
||||
|
||||
if (auto& body_ptr = it->second.body; !body_ptr) {
|
||||
updated = true;
|
||||
|
||||
body_ptr.reset(new char[body.size() + 1]);
|
||||
memcpy(body_ptr.get(), body.data(), body.size());
|
||||
body_ptr[body.size()] = '\0';
|
||||
}
|
||||
|
||||
UpdateScriptCaches(id, it->second);
|
||||
|
||||
return updated;
|
||||
}
|
||||
|
||||
optional<ScriptMgr::ScriptData> ScriptMgr::Find(std::string_view sha) const {
|
||||
if (sha.size() != ScriptKey{}.size())
|
||||
return std::nullopt;
|
||||
|
||||
lock_guard lk{mu_};
|
||||
if (auto it = db_.find(sha); it != db_.end() && it->second.body)
|
||||
return ScriptData{it->second, it->second.body.get()};
|
||||
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
vector<pair<string, string>> ScriptMgr::GetAll() const {
|
||||
vector<pair<string, string>> res;
|
||||
|
||||
lock_guard lk{mu_};
|
||||
res.reserve(db_.size());
|
||||
for (const auto& [sha, data] : db_) {
|
||||
res.emplace_back(string{sha.data(), sha.size()}, data.body.get());
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
void ScriptMgr::UpdateScriptCaches(ScriptKey sha, ScriptParams params) const {
|
||||
shard_set->pool()->Await([&sha, ¶ms](auto index, auto* pb) {
|
||||
ServerState::tlocal()->SetScriptParams(sha, params);
|
||||
});
|
||||
}
|
||||
|
||||
bool ScriptMgr::ScriptParams::ApplyPragma(string_view pragma, ScriptParams* params) {
|
||||
if (pragma == "disable-atomicity") {
|
||||
params->atomic = false;
|
||||
return true;
|
||||
}
|
||||
|
||||
if (pragma == "allow-undeclared-keys") {
|
||||
params->undeclared_keys = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace dfly
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue