fix(server): Small fixes

Signed-off-by: Vladislav Oleshko <[email protected]>
dragonflydb · Feb 19, 2023 · af21149 · af21149
1 parent b80c8a4
commit af21149
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 23 deletions.
diff --git a/src/server/main_service.cc b/src/server/main_service.cc
@@ -537,29 +537,32 @@ static void MultiSetError(ConnectionContext* cntx) {
   }
 }
 
-bool CheckEvalKeys(const ConnectionState::ScriptInfo& eval_info, const CommandId* cid,
-                   CmdArgList args, Transaction* trans) {
-  auto multi_mode = trans->GetMultiMode();
-
-  // We allow any keys in GLOBAL and NON_ATOMIC mode.
+// Return OK if all keys are allowed to be accessed: either declared in EVAL or
+// transaction is running in global or non-atomic mode.
+OpStatus CheckKeysDeclared(const ConnectionState::ScriptInfo& eval_info, const CommandId* cid,
+                           CmdArgList args, Transaction* trans) {
+  Transaction::MultiMode multi_mode = trans->GetMultiMode();
+
+  // We either scheduled on all shards or re-schedule for each operation,
+  // so we are not restricted to any keys.
   if (multi_mode == Transaction::GLOBAL || multi_mode == Transaction::NON_ATOMIC)
-    return true;
+    return OpStatus::OK;
 
   OpResult<KeyIndex> key_index_res = DetermineKeys(cid, args);
   if (!key_index_res)
-    return false;  // TODO: Propagate error
+    return key_index_res.status();
 
   const auto& key_index = *key_index_res;
   for (unsigned i = key_index.start; i < key_index.end; ++i) {
     if (!eval_info.keys.contains(ArgS(args, i))) {
-      return false;
+      return OpStatus::KEY_NOTFOUND;
     }
   }
 
   if (unsigned i = key_index.bonus; i && !eval_info.keys.contains(ArgS(args, i)))
-    return false;
+    return OpStatus::KEY_NOTFOUND;
 
-  return true;
+  return OpStatus::OK;
 }
 
 void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx) {
@@ -684,15 +687,20 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx)
   if (under_script) {
     DCHECK(dfly_cntx->transaction);
     if (IsTransactional(cid)) {
-      if (!CheckEvalKeys(*dfly_cntx->conn_state.script_info, cid, args, dfly_cntx->transaction))
+      OpStatus status =
+          CheckKeysDeclared(*dfly_cntx->conn_state.script_info, cid, args, dfly_cntx->transaction);
+
+      if (status == OpStatus::KEY_NOTFOUND)
         return (*cntx)->SendError("script tried accessing undeclared key");
 
+      if (status != OpStatus::OK)
+        return (*cntx)->SendError(status);
+
       dfly_cntx->transaction->MultiSwitchCmd(cid);
-      OpStatus st = dfly_cntx->transaction->InitByArgs(dfly_cntx->conn_state.db_index, args);
+      status = dfly_cntx->transaction->InitByArgs(dfly_cntx->conn_state.db_index, args);
 
-      if (st != OpStatus::OK) {
-        return (*cntx)->SendError(st);
-      }
+      if (status != OpStatus::OK)
+        return (*cntx)->SendError(status);
     }
   } else {
     DCHECK(dfly_cntx->transaction == nullptr);

diff --git a/src/server/transaction.cc b/src/server/transaction.cc
@@ -149,7 +149,7 @@ void Transaction::InitMultiData(KeyIndex key_index) {
   DCHECK(multi_);
   auto args = cmd_with_full_args_;
 
-  if (multi_->mode == Transaction::NON_ATOMIC)
+  if (multi_->mode == NON_ATOMIC)
     return;
 
   // TODO: determine correct locking mode for transactions, scripts and regular commands.
@@ -231,8 +231,7 @@ void Transaction::InitByKeys(KeyIndex key_index) {
   DCHECK_LT(key_index.start, args.size());
 
   bool needs_reverse_mapping = cid_->opt_mask() & CO::REVERSE_MAPPING;
-  bool single_key =
-      key_index.HasSingleKey() && (!multi_ || multi_->mode == Transaction::NON_ATOMIC);
+  bool single_key = key_index.HasSingleKey() && !IsRunningMulti();
 
   if (single_key) {
     DCHECK_GT(key_index.step, 0u);
@@ -270,7 +269,7 @@ void Transaction::InitByKeys(KeyIndex key_index) {
   // Compress shard data, if we occupy only one shard.
   if (unique_shard_cnt_ == 1) {
     PerShardData* sd;
-    if (multi_ && multi_->mode != NON_ATOMIC) {
+    if (IsRunningMulti()) {
       sd = &shard_data_[unique_shard_id_];
     } else {
       shard_data_.resize(1);

diff --git a/src/server/transaction.h b/src/server/transaction.h
@@ -208,10 +208,6 @@ class Transaction {
     return bool(multi_);
   }
 
-  bool IsRunningMulti() const {
-    return multi_ && multi_->mode != NON_ATOMIC;
-  }
-
   MultiMode GetMultiMode() const {
     return multi_->mode;
   }
@@ -406,6 +402,14 @@ class Transaction {
     return use_count_.load(std::memory_order_relaxed);
   }
 
+  // Whether the transaction is multi and runs in an atomic mode.
+  // This, instead of just IsMulti(), should be used to check for the possibility of
+  // different optimizations, because they can safely be applied to non-atomic multi
+  // transactions as well.
+  bool IsRunningMulti() const {
+    return multi_ && multi_->mode != NON_ATOMIC;
+  }
+
   unsigned SidToId(ShardId sid) const {
     return sid < shard_data_.size() ? sid : 0;
   }