From 734401098cf58a3e4946b84c58b42660ed935f9c Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Wed, 9 Aug 2023 14:18:34 +0300
Subject: [PATCH] opt(server): Execute lua on target shard, if it's 1 (#1639)

* opt(server): Execute lua on target shard, if it's 1

This will save hops by short-circuiting execution of commands.

* Reuse unique shard id from tx
Only switch threads for LOCK_AHEAD

* Signedness
---
 src/server/main_service.cc | 13 ++++++++++---
 src/server/transaction.cc  |  5 +++++
 src/server/transaction.h   |  3 +++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index 23e14cb0d..d3ce75d85 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -1438,15 +1438,22 @@ void Service::EvalInternal(const EvalArgs& eval_args, Interpreter* interpreter,
     sinfo->keys.insert(KeyLockArgs::GetLockKey(ArgS(eval_args.keys, i)));
   }
   sinfo->async_cmds_heap_limit = absl::GetFlag(FLAGS_multi_eval_squash_buffer);
-  DCHECK(cntx->transaction);
+  Transaction* tx = cntx->transaction;
+  CHECK(tx != nullptr);
 
-  bool scheduled = StartMultiEval(cntx->db_index(), eval_args.keys, *params, cntx->transaction);
+  bool scheduled = StartMultiEval(cntx->db_index(), eval_args.keys, *params, tx);
 
   interpreter->SetGlobalArray("KEYS", eval_args.keys);
   interpreter->SetGlobalArray("ARGV", eval_args.args);
   interpreter->SetRedisFunc([cntx, this](auto args) { CallFromScript(cntx, args); });
 
-  Interpreter::RunResult result = interpreter->RunFunction(eval_args.sha, &error);
+  Interpreter::RunResult result;
+  if (tx->GetMultiMode() == Transaction::LOCK_AHEAD && tx->GetUniqueShardCnt() == 1) {
+    ShardId sid = tx->GetUniqueShard();
+    pp_.at(sid)->Await([&]() { result = interpreter->RunFunction(eval_args.sha, &error); });
+  } else {
+    result = interpreter->RunFunction(eval_args.sha, &error);
+  }
   absl::Cleanup clean = [interpreter]() { interpreter->ResetStack(); };
 
   if (auto err = FlushEvalAsyncCmds(cntx, true); err) {
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index 2234ab8e1..0a940728d 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -946,6 +946,11 @@ string_view Transaction::Name() const {
   return cid_->name();
 }
 
+ShardId Transaction::GetUniqueShard() const {
+  DCHECK_EQ(GetUniqueShardCnt(), 1U);
+  return unique_shard_id_;
+}
+
 KeyLockArgs Transaction::GetLockArgs(ShardId sid) const {
   KeyLockArgs res;
   res.db_index = db_index_;
diff --git a/src/server/transaction.h b/src/server/transaction.h
index 54d1c59d9..ac6df7409 100644
--- a/src/server/transaction.h
+++ b/src/server/transaction.h
@@ -263,6 +263,9 @@ class Transaction {
     return unique_shard_cnt_;
   }
 
+  // This method is meaningless if GetUniqueShardCnt() != 1.
+  ShardId GetUniqueShard() const;
+
   bool IsMulti() const {
     return bool(multi_);
   }