chore: move single shard lua to new interface

Signed-off-by: Vladislav Oleshko <[email protected]>

Author: dranikpg

src/server/command_registry.h

@@ -93,6 +93,10 @@ class CommandId : public facade::CommandId {
 
   bool IsTransactional() const;
 
+  bool IsMultiTransactional() const {
+    return CO::IsTransKind(name()) || CO::IsEvalKind(name());
+  }
+
   bool IsReadOnly() const {
     return opt_mask_ & CO::READONLY;
   }

src/server/main_service.cc

@@ -1721,7 +1721,7 @@ optional<bool> StartMultiEval(DbIndex dbid, CmdArgList keys, ScriptMgr::ScriptPa
       trans->StartMultiGlobal(dbid);
       return true;
     case Transaction::LOCK_AHEAD:
-      trans->StartMultiLockedAhead(dbid, keys);
+      trans->StartMultiLockedAhead(dbid, keys, true);
       return true;
     case Transaction::NON_ATOMIC:
       trans->StartMultiNonAtomic();
@@ -1762,26 +1762,6 @@ std::pair<const CommandId*, CmdArgList> Service::FindCmd(CmdArgList args) const
   return {res, args.subspan(1)};
 }
 
-static bool CanRunSingleShardMulti(optional<ShardId> sid, const ScriptMgr::ScriptParams& params,
-                                   const Transaction& tx) {
-  if (!sid.has_value()) {
-    return false;
-  }
-
-  if (DetermineMultiMode(params) != Transaction::LOCK_AHEAD) {
-    return false;
-  }
-
-  if (tx.GetMultiMode() != Transaction::NOT_DETERMINED) {
-    // We may be running EVAL under MULTI. Currently RunSingleShardMulti() will attempt to lock
-    // keys, in which case will be already locked by MULTI. We could optimize this path as well
-    // though.
-    return false;
-  }
-
-  return true;
-}
-
 void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpreter* interpreter,
                            ConnectionContext* cntx) {
   DCHECK(!eval_args.sha.empty());
@@ -1806,20 +1786,6 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
   sinfo = make_unique<ConnectionState::ScriptInfo>();
   sinfo->keys.reserve(eval_args.keys.size());
 
-  optional<ShardId> sid;
-  for (size_t i = 0; i < eval_args.keys.size(); ++i) {
-    string_view key = ArgS(eval_args.keys, i);
-    sinfo->keys.insert(KeyLockArgs::GetLockKey(key));
-
-    ShardId cur_sid = Shard(key, shard_count());
-    if (i == 0) {
-      sid = cur_sid;
-    }
-    if (sid.has_value() && *sid != cur_sid) {
-      sid = nullopt;
-    }
-  }
-
   sinfo->async_cmds_heap_limit = absl::GetFlag(FLAGS_multi_eval_squash_buffer);
   Transaction* tx = cntx->transaction;
   CHECK(tx != nullptr);
@@ -1834,19 +1800,29 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
 
   Interpreter::RunResult result;
 
-  if (CanRunSingleShardMulti(sid, *params, *tx)) {
-    // If script runs on a single shard, we run it remotely to save hops.
+  optional<bool> scheduled = StartMultiEval(cntx->db_index(), eval_args.keys, *params, cntx);
+  if (!scheduled) {
+    return;
+  }
+
+  // If script runs on a single shard, we run it remotely to save hops
+  if (!tx->IsScheduled() && tx->GetMultiMode() == Transaction::LOCK_AHEAD &&
+      tx->GetUniqueShardCnt() == 1) {
+    DCHECK(*scheduled);         // because tx multi mode is lock ahead
+    CHECK(!tx->IsScheduled());  // skip_scheduling = true in StartMultiEval
+
     interpreter->SetRedisFunc([cntx, tx, this](Interpreter::CallArgs args) {
       // Disable squashing, as we're using the squashing mechanism to run remotely.
       args.async = false;
       CallFromScript(cntx, args);
     });
 
     ++ServerState::tlocal()->stats.eval_shardlocal_coordination_cnt;
-    // TODO: remove doubule key iteration
-    tx->PrepareMultiForScheduleSingleHop(*sid, tx->GetDbIndex(), args);
+
+    auto sid = tx->GetUniqueShard();
+    tx->MultiBecomeSquasher();
     tx->ScheduleSingleHop([&](Transaction*, EngineShard*) {
-      boost::intrusive_ptr<Transaction> stub_tx = new Transaction{tx, *sid};
+      boost::intrusive_ptr<Transaction> stub_tx = new Transaction{tx, sid};
       stub_tx->MultiUpdateWithParent(tx);
       cntx->transaction = stub_tx.get();
 
@@ -1856,16 +1832,14 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
       return OpStatus::OK;
     });
 
-    if (*sid != ServerState::tlocal()->thread_index()) {
+    if (sid != ServerState::tlocal()->thread_index()) {
       VLOG(1) << "Migrating connection " << cntx->conn() << " from "
-              << ProactorBase::me()->GetPoolIndex() << " to " << *sid;
-      cntx->conn()->RequestAsyncMigration(shard_set->pool()->at(*sid));
+              << ProactorBase::me()->GetPoolIndex() << " to " << sid;
+      cntx->conn()->RequestAsyncMigration(shard_set->pool()->at(sid));
     }
   } else {
-    optional<bool> scheduled = StartMultiEval(cntx->db_index(), eval_args.keys, *params, cntx);
-    if (!scheduled) {
-      return;
-    }
+    if (*scheduled && !tx->IsScheduled())
+      tx->Schedule();
 
     ++ServerState::tlocal()->stats.eval_io_coordination_cnt;
     interpreter->SetRedisFunc(
@@ -2080,6 +2054,7 @@ void Service::Exec(CmdArgList args, ConnectionContext* cntx) {
     if (scheduled && allow_squashing) {
       MultiCommandSquasher::Execute(absl::MakeSpan(exec_info.body), cntx, this);
     } else {
+      DCHECK(!scheduled || !delay_scheduling);
       for (auto& scmd : exec_info.body) {
         VLOG(2) << "TX CMD " << scmd.Cid()->name() << " " << scmd.NumArgs();
 

src/server/multi_command_squasher.cc

@@ -86,7 +86,7 @@ void MultiCommandSquasher::PerformHop() {
   // Doesn't work with replication, so disallow if inline scheduling is not allowed.
   bool singlehop_possible = IsAtomic() && !tx->IsScheduled() && cmds_.empty();
   if (singlehop_possible) {
-    tx->MultiSwitchToNonAtomic();
+    tx->MultiBecomeSquasher();
     DCHECK_GT(tx->GetUniqueShardCnt(), 0u);  // it was initialized and determined active shards
     tx->ScheduleSingleHop(run_cb);
     return;

src/server/transaction.cc

@@ -347,11 +347,10 @@ void Transaction::PrepareSquashedMultiHop(const CommandId* cid,
   }
 }
 
-void Transaction::MultiSwitchToNonAtomic() {
+void Transaction::MultiBecomeSquasher() {
   DCHECK_EQ(multi_->mode, LOCK_AHEAD);
   CHECK_EQ(coordinator_state_, 0);  // not scheduled and certainly not executing
   multi_->role = SQUASHER;
-  multi_->locks.clear();
 }
 
 void Transaction::StartMultiGlobal(DbIndex dbid) {
@@ -440,15 +439,6 @@ string Transaction::DebugId() const {
   return res;
 }
 
-void Transaction::PrepareMultiForScheduleSingleHop(ShardId sid, DbIndex db, CmdArgList args) {
-  multi_.reset();
-  InitBase(db, args);
-  EnableShard(sid);
-  OpResult<KeyIndex> key_index = DetermineKeys(cid_, args);
-  CHECK(key_index);
-  StoreKeysInArgs(*key_index, false);
-}
-
 // Runs in the dbslice thread. Returns true if the transaction continues running in the thread.
 bool Transaction::RunInShard(EngineShard* shard, bool txq_ooo) {
   DCHECK_GT(run_count_.load(memory_order_relaxed), 0u);
@@ -588,6 +578,7 @@ void Transaction::ScheduleInternal() {
   DCHECK(!shard_data_.empty());
   DCHECK_EQ(0u, txid_);
   DCHECK_EQ(0, coordinator_state_ & (COORD_SCHED | COORD_OOO));
+  DCHECK(!multi_ || cid_->IsMultiTransactional());
 
   bool span_all = IsGlobal();
 
@@ -706,6 +697,7 @@ OpStatus Transaction::ScheduleSingleHop(RunnableType cb) {
   bool scheduled = (coordinator_state_ & COORD_SCHED) > 0;
   DCHECK(!scheduled || IsAtomicMulti());  // if we don't need to schedule, we're an atomic multi
   DCHECK(scheduled || !IsAtomicMulti() || multi_->role == SQUASHER);  // only useful for squashing
+  DCHECK(scheduled || !multi_ || cid_->IsMultiTransactional());  // don't schedule on normal cmd
 
   if (!scheduled)  // Conclude only if we schedule, otherwise we're part of multi-hop multi
     coordinator_state_ |= COORD_EXEC_CONCLUDING;

src/server/transaction.h

@@ -320,8 +320,8 @@ class Transaction {
   // to it must not block.
   void PrepareMultiForScheduleSingleHop(ShardId sid, DbIndex db, CmdArgList args);
 
-  // Switch lock-ahead transaction to non-atomic. Can only be called *before* it was scheduled.
-  void MultiSwitchToNonAtomic();
+  // Mark as squasher.
+  void MultiBecomeSquasher();
 
   // Write a journal entry to a shard journal with the given payload. When logging a non-automatic
   // journal command, multiple journal entries may be necessary. In this case, call with set