fix: Handle bonus key + init error in squashing (#5303)

* fix: Handle bonus key + init error in squashing

Author: dranikpg

src/server/multi_command_squasher.cc

@@ -153,10 +153,14 @@ bool MultiCommandSquasher::ExecuteStandalone(RedisReplyBuilder* rb, const Stored
   auto* tx = cntx_->transaction;
   if (cmd->Cid()->IsTransactional()) {
     tx->MultiSwitchCmd(cmd->Cid());
-    tx->InitByArgs(cntx_->ns, cntx_->conn_state.db_index, args);
+    auto status = tx->InitByArgs(cntx_->ns, cntx_->conn_state.db_index, args);
+    if (status != OpStatus::OK) {
+      rb->SendError(status);
+      rb->ConsumeLastError();
+      return !opts_.error_abort;
+    }
   }
   service_->InvokeCmd(cmd->Cid(), args, CommandContext{tx, rb, cntx_});
-
   return true;
 }
 
@@ -195,10 +199,13 @@ OpStatus MultiCommandSquasher::SquashedHopCb(EngineShard* es, RespVersion resp_v
     crb.SetReplyMode(dispatched.cmd->ReplyMode());
 
     local_tx->MultiSwitchCmd(dispatched.cmd->Cid());
-    local_tx->InitByArgs(cntx_->ns, local_cntx.conn_state.db_index, args);
-    service_->InvokeCmd(dispatched.cmd->Cid(), args,
-                        CommandContext{local_cntx.transaction, &crb, &local_cntx});
-
+    auto status = local_tx->InitByArgs(cntx_->ns, local_cntx.conn_state.db_index, args);
+    if (status != OpStatus::OK) {
+      crb.SendError(status);
+    } else {
+      service_->InvokeCmd(dispatched.cmd->Cid(), args,
+                          CommandContext{local_cntx.transaction, &crb, &local_cntx});
+    }
     move_reply(crb.Take(), &dispatched.reply);
 
     // Assert commands made no persistent state changes to stub context state

src/server/multi_test.cc

@@ -330,6 +330,29 @@ TEST_F(MultiTest, MultiWithoutTx) {
   EXPECT_EQ(resp.GetVec()[4], "OK3");
 }
 
+TEST_F(MultiTest, MultiCommandsWithBonusKeys) {
+  absl::FlagSaver fs;
+  absl::SetFlag(&FLAGS_multi_exec_squash, true);
+
+  EXPECT_EQ(Shard("za", shard_set->size()), Shard("zb", shard_set->size()));
+  EXPECT_EQ(Shard("zb", shard_set->size()), Shard("ze", shard_set->size()));
+
+  // Check bonus keys are correctly processed with squashing
+  Run({"multi"});
+  Run({"zadd", "za", "1", "a", "2", "b"});
+  Run({"zadd", "zb", "2", "b", "3", "c"});
+  Run({"zinterstore", "ze", "2", "za", "zb"});
+  auto resp = Run({"exec"});
+  EXPECT_THAT(resp.GetVec()[2], IntArg(1));
+  EXPECT_THAT(Run({"zcard", "ze"}), IntArg(1));
+
+  // Check squashing correctly pre-validates commands
+  Run({"multi"});
+  Run({"zinterstore", "ze", "2", "za", "zb", "z one extra"});
+  resp = Run({"exec"});
+  EXPECT_THAT(resp, ErrArg("syntax error"));
+}
+
 TEST_F(MultiTest, MultiHop) {
   Run({"set", kKey1, "1"});
 

src/server/transaction.cc

@@ -307,12 +307,14 @@ void Transaction::PrepareMultiFps(CmdArgList keys) {
 }
 
 void Transaction::StoreKeysInArgs(const KeyIndex& key_index) {
-  DCHECK(!key_index.bonus);
   DCHECK(kv_fp_.empty());
   DCHECK(args_slices_.empty());
 
   // even for a single key we may have multiple arguments per key (MSET).
+  if (key_index.bonus)
+    args_slices_.emplace_back(*key_index.bonus, *key_index.bonus + 1);
   args_slices_.emplace_back(key_index.start, key_index.end);
+
   for (string_view key : key_index.Range(full_args_))
     kv_fp_.push_back(LockTag(key).Fingerprint());
 }