chore: adjust transaction code to keystep/3 commands (#2941)

This commit is contained in:
Roman Gershman 2024-04-21 17:20:57 +03:00 committed by GitHub
parent d2a08e2c3c
commit 3d609504d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 52 additions and 8 deletions

View file

@ -40,6 +40,7 @@ namespace dfly {
using namespace std;
using namespace jsoncons;
using facade::kSyntaxErrType;
using facade::WrongNumArgsError;
using JsonExpression = jsonpath::jsonpath_expression<JsonType>;
using OptBool = optional<bool>;
@ -1355,6 +1356,24 @@ void JsonFamily::Set(CmdArgList args, ConnectionContext* cntx) {
}
}
// JSON.MSET key path value [key path value ...]
void JsonFamily::MSet(CmdArgList args, ConnectionContext* cntx) {
DCHECK_GE(args.size(), 3u);
if (args.size() % 3 != 0) {
return cntx->SendError(facade::WrongNumArgsError("json.mset"));
}
auto cb = [&](Transaction* t, EngineShard* shard) {
ArgSlice args = t->GetShardArgs(shard->shard_id());
LOG(INFO) << shard->shard_id() << " " << args;
return OpStatus::OK;
};
Transaction* trans = cntx->transaction;
trans->ScheduleSingleHop(cb);
cntx->SendOk();
}
void JsonFamily::Resp(CmdArgList args, ConnectionContext* cntx) {
string_view key = ArgS(args, 0);
string_view path = DefaultJsonPath;
@ -2026,9 +2045,12 @@ void JsonFamily::Register(CommandRegistry* registry) {
ArrAppend);
*registry << CI{"JSON.ARRINDEX", CO::READONLY | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(ArrIndex);
// TODO: Support negative first_key index to revive the debug sub-command
*registry << CI{"JSON.DEBUG", CO::READONLY | CO::FAST, -3, 2, 2, acl::JSON}.HFUNC(Debug);
*registry << CI{"JSON.RESP", CO::READONLY | CO::FAST, -2, 1, 1, acl::JSON}.HFUNC(Resp);
*registry << CI{"JSON.SET", CO::WRITE | CO::DENYOOM | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(Set);
*registry << CI{"JSON.DEBUG", CO::READONLY | CO::FAST, -3, 2, 2, acl::JSON}.HFUNC(Debug)
<< CI{"JSON.RESP", CO::READONLY | CO::FAST, -2, 1, 1, acl::JSON}.HFUNC(Resp)
<< CI{"JSON.SET", CO::WRITE | CO::DENYOOM | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(Set)
<< CI{"JSON.MSET", CO::WRITE | CO::DENYOOM | CO::FAST | CO::INTERLEAVED_KEYS, -4, 1, -1,
acl::JSON}
.HFUNC(MSet);
}
} // namespace dfly

View file

@ -41,6 +41,7 @@ class JsonFamily {
static void Debug(CmdArgList args, ConnectionContext* cntx);
static void Resp(CmdArgList args, ConnectionContext* cntx);
static void Set(CmdArgList args, ConnectionContext* cntx);
static void MSet(CmdArgList args, ConnectionContext* cntx);
};
} // namespace dfly

View file

@ -1081,4 +1081,18 @@ TEST_F(JsonFamilyTest, Set) {
EXPECT_EQ(resp, R"([{"a":2,"b":8,"c":[1,2,3]}])");
}
TEST_F(JsonFamilyTest, MSet) {
string json = R"(
{"a":{"a":1, "b":2, "c":3}}
)";
auto resp = Run({"JSON.MSET", "j1", "$"});
EXPECT_THAT(resp, ErrArg("wrong number"));
resp = Run({"JSON.MSET", "j1", "$", json, "j3", "$"});
EXPECT_THAT(resp, ErrArg("wrong number"));
resp = Run({"JSON.MSET", "j1", "$", json, "j3", "$", json});
EXPECT_EQ(resp, "OK");
}
} // namespace dfly

View file

@ -208,8 +208,8 @@ void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vector<PerShar
shard_index[sid].key_step = key_index.step;
add(sid, i);
DCHECK_LE(key_index.step, 2u);
if (key_index.step == 2) { // Handle value associated with preceding key.
// Handle values associated with preceding key.
for (unsigned j = 1; j < key_index.step; ++j) {
add(sid, ++i);
}
}
@ -333,8 +333,7 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
}
shard_data_.resize(shard_set->size()); // shard_data isn't sparse, so we must allocate for all :(
DCHECK(key_index.step == 1 || key_index.step == 2);
DCHECK(key_index.step != 2 || (full_args_.size() % 2) == 0);
DCHECK_EQ(full_args_.size() % key_index.step, 0u);
// Safe, because flow below is not preemptive.
auto& shard_index = tmp_space.GetShardIndex(shard_data_.size());
@ -1578,7 +1577,15 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
} else {
key_index.end = last > 0 ? last : (int(args.size()) + last + 1);
}
key_index.step = cid->opt_mask() & CO::INTERLEAVED_KEYS ? 2 : 1;
if (cid->opt_mask() & CO::INTERLEAVED_KEYS) {
if (cid->name() == "JSON.MSET") {
key_index.step = 3;
} else {
key_index.step = 2;
}
} else {
key_index.step = 1;
}
if (cid->opt_mask() & CO::STORE_LAST_KEY) {
string_view name{cid->name()};