/Users/deen/code/yugabyte-db/src/yb/yql/cql/ql/exec/exec_context.cc

Source (jump to first uncovered line)
//--------------------------------------------------------------------------------------------------
// Copyright (c) YugaByte, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
// in compliance with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied.  See the License for the specific language governing permissions and limitations
// under the License.
//
//--------------------------------------------------------------------------------------------------

#include "yb/yql/cql/ql/exec/exec_context.h"

#include <boost/function.hpp>

#include "yb/client/table.h"
#include "yb/client/transaction.h"
#include "yb/client/yb_op.h"

#include "yb/common/ql_rowblock.h"
#include "yb/common/schema.h"

#include "yb/gutil/casts.h"

#include "yb/rpc/thread_pool.h"

#include "yb/util/result.h"
#include "yb/util/status_format.h"
#include "yb/util/trace.h"
#include "yb/util/tsan_util.h"

#include "yb/yql/cql/ql/exec/rescheduler.h"
#include "yb/yql/cql/ql/ptree/parse_tree.h"
#include "yb/yql/cql/ql/ptree/pt_select.h"
#include "yb/yql/cql/ql/util/statement_params.h"

DEFINE_int32(cql_prepare_child_threshold_ms, 2000 * yb::kTimeMultiplier,
             "Timeout if preparing for child transaction takes longer"
             "than the prescribed threshold.");

namespace yb {
namespace ql {

using client::CommitCallback;
using client::Restart;
using client::YBqlReadOpPtr;
using client::YBSessionPtr;
using client::YBTransactionPtr;

ExecContext::ExecContext(const ParseTree& parse_tree, const StatementParameters& params)
    : parse_tree_(parse_tree), params_(params) {
}

ExecContext::~ExecContext() {
  // Reset to abort transaction explicitly instead of letting it expire.
  // Should be ok not to take a rescheduler here since the `ExecContext` clean up should happen
  // only when we return a response to the CQL client, which is now guaranteed to happen in
  // CQL proxy's handler thread.
  Reset(client::Restart::kFalse, nullptr);
}

TnodeContext* ExecContext::AddTnode(const TreeNode *tnode) {
  restart_ = client::Restart::kFalse;
  tnode_contexts_.emplace_back(tnode);
  return &tnode_contexts_.back();
}

//--------------------------------------------------------------------------------------------------
Status ExecContext::StartTransaction(
    const IsolationLevel isolation_level, QLEnv* ql_env, Rescheduler* rescheduler) {
  TRACE("Start Transaction");
  transaction_start_time_ = MonoTime::Now();
  if (!transaction_) {
    transaction_ = VERIFY_RESULT(ql_env->NewTransaction(
        transaction_, isolation_level, rescheduler->GetDeadline()));
  } else if (1.28ktransaction_->IsRestartRequired()1.28k) {
    transaction_ = VERIFY_RESULT(transaction_->CreateRestartedTransaction());
  } else {
    // If there is no need to start or restart transaction, just return. This can happen to DMLs on
    // a table with secondary index inside a "BEGIN TRANSACTION ... END TRANSACTION" block. Each DML
    // will try to start a transaction "on-demand" and we will use the shared transaction already
    // started by "BEGIN TRANSACTION".
    return Status::OK();
  }

  if (!transactional_session_) {
    transactional_session_ = ql_env->NewSession();
    transactional_session_->SetReadPoint(client::Restart::kFalse);
  }
  transactional_session_->SetDeadline(rescheduler->GetDeadline());
  transactional_session_->SetTransaction(transaction_);

  return Status::OK();
}

Status ExecContext::PrepareChildTransaction(
    CoarseTimePoint deadline, ChildTransactionDataPB* data) {
  auto future = DCHECK_NOTNULL(transaction_.get())->PrepareChildFuture(
      client::ForceConsistentRead::kTrue, deadline);

  // Set the deadline to be the earlier of the input deadline and the current timestamp
  // plus the waiting time for the prepare child
  auto future_deadline = std::min(
      deadline,
      CoarseMonoClock::Now() + MonoDelta::FromMilliseconds(FLAGS_cql_prepare_child_threshold_ms));

  auto future_status = future.wait_until(future_deadline);

  if (future_status == std::future_status::ready) {
    *data = VERIFY_RESULT(std::move(future).get());
    return Status::OK();
  }

  auto message = Format("Timed out waiting for prepare child status, left to deadline: $0",
                        MonoDelta(deadline - CoarseMonoClock::now()));
  LOG(INFO) << message;
  return STATUS(TimedOut, message);
}

Status ExecContext::ApplyChildTransactionResult(const ChildTransactionResultPB& result) {
  return DCHECK_NOTNULL(transaction_.get())->ApplyChildResult(result);
}

void ExecContext::CommitTransaction(CoarseTimePoint deadline, CommitCallback callback) {
  if (!transaction_) {
    LOG(DFATAL) << "No transaction to commit";
    return;
  }

  // Clear the transaction from the session before committing the transaction. SetTransaction()
  // must be called before the Commit() call instead of after because when the commit callback is
  // invoked, it will finish the current transaction, return the response and make the CQLProcessor
  // available for the next statement and its operations would be aborted by SetTransaction().
  transactional_session_->SetTransaction(nullptr);
  transactional_session_ = nullptr;

  YBTransactionPtr transaction = std::move(transaction_);
  TRACE("Commit Transaction");
  transaction->Commit(deadline, std::move(callback));
}

void ExecContext::AbortTransaction() {
  if (!transaction_) {
    LOG(DFATAL) << "No transaction to abort";
    return;
  }

  // Abort the session and clear the transaction from the session before aborting the transaction.
  transactional_session_->Abort();
  transactional_session_->SetTransaction(nullptr);
  transactional_session_ = nullptr;

  YBTransactionPtr transaction = std::move(transaction_);
  TRACE("Abort Transaction");
  transaction->Abort();
}

bool ExecContext::HasPendingOperations() const {
  for (const auto& tnode_context : tnode_contexts_) {
    if (tnode_context.HasPendingOperations()) {
      return true;
    }
  }
  return false;
}

class AbortTransactionTask : public rpc::ThreadPoolTask {
 public:
  explicit AbortTransactionTask(YBTransactionPtr transaction)
      : transaction_(std::move(transaction)) {}

  void Run() override {
    transaction_->Abort();
    transaction_ = nullptr;
  }

  void Done(const Status& status) override {
    delete this;
  }

  virtual ~AbortTransactionTask() {
  }

 private:
  YBTransactionPtr transaction_;
};

//--------------------------------------------------------------------------------------------------
void ExecContext::Reset(const Restart restart, Rescheduler* rescheduler) {
  if (transactional_session_) {
    transactional_session_->Abort();
    transactional_session_->SetTransaction(nullptr);
  }
  if (transaction_ && !(42.5ktransaction_->IsRestartRequired()42.5k && restart0)) {
    YBTransactionPtr transaction = std::move(transaction_);
    TRACE("Abort Transaction");
    if (rescheduler && rescheduler->NeedReschedule()41.1k) {
      rescheduler->Reschedule(new AbortTransactionTask(std::move(transaction)));
    } else {
      transaction->Abort();
    }
  }
  restart_ = restart;
  tnode_contexts_.clear();
  if (restart) {
    num_retries_++;
  }
}

//--------------------------------------------------------------------------------------------------
TnodeContext::TnodeContext(const TreeNode* tnode) : tnode_(tnode), start_time_(MonoTime::Now()) {
}

TnodeContext::~TnodeContext() = default;

TnodeContext* TnodeContext::AddChildTnode(const TreeNode* tnode) {
  DCHECK(!child_context_);
  child_context_ = std::make_unique<TnodeContext>(tnode);
  return child_context_.get();
}

Status TnodeContext::AppendRowsResult(RowsResult::SharedPtr&& rows_result) {
  // Append data arriving from DocDB.
  // (1) SELECT without nested query.
  //  - SELECT <select_list> FROM <table or index>
  //      WHERE <filter_cond>
  //      LIMIT <limit> OFFSET <offset>
  //  - New rows are appended at the end.
  //
  // (2) SELECT with nested query.
  //  - SELECT <select_list> FROM <table>
  //      WHERE
  //        primary_key IN (SELECT primary_key FROM <index> WHERE <index_cond>)
  //        AND
  //        <filter_cond>
  //      LIMIT <limit> OFFSET <offset>
  //  - When nested INDEX query fully-covers the SELECT command, data coming from the nested node
  //    is appended without being filtered or rejected.
  //  - When nested INDEX query does NOT fully-cover the SELECT command, data is rejected if OFFSET
  //    is not yet reached. Otherwise, data is appended ONE row at a time.
  if (!rows_result) {
    return Status::OK();
  }

  int64_t number_of_new_rows =
    VERIFY_RESULT(QLRowBlock::GetRowCount(YQL_CLIENT_CQL, rows_result->rows_data()));

  if (query_state_) {
    RSTATUS_DCHECK(tnode_->opcode() == TreeNodeOpcode::kPTSelectStmt,
                   Corruption, "QueryPagingState is setup for non-select statement");
    const auto* select_stmt = static_cast<const PTSelectStmt *>(tnode_);

    // Save the last offset status before loading new status from DocDB.
    const bool reached_offset = query_state_->reached_select_offset();
    const bool has_nested_query = select_stmt->child_select() != nullptr;
    RETURN_NOT_OK(
      query_state_->LoadPagingStateFromDocdb(rows_result, number_of_new_rows, has_nested_query));

    if (!reached_offset && has_nested_query714) {
      // Parent query needs to discard the new row when (row_count < SELECT::OFFSET).
      if (!rows_result_) {
        rows_result_ = std::make_shared<RowsResult>(select_stmt);
      }
      rows_result_->SetPagingState(std::move(*rows_result));
      return Status::OK();
    }
  }

  // Append the new rows to result.
  row_count_ += number_of_new_rows;
  if (rows_result_ == nullptr) {
    rows_result_ = std::move(rows_result);
    return Status::OK();
  }
  return rows_result_->Append(std::move(*rows_result));
}

void TnodeContext::InitializePartition(QLReadRequestPB *req, bool continue_user_request) {
  uint64_t start_partition = continue_user_request ? query_state_->next_partition_index()16 : 0850;

  current_partition_index_ = start_partition;
  // Hash values before the first 'IN' condition will be already set.
  // hash_values_options_ vector starts from the first column with an 'IN' restriction.
  // E.g. for a query "h1 = 1 and h2 in (2,3) and h3 in (4,5) and h4 = 6":
  // hashed_column_values() will be [1] and hash_values_options_ will be [[2,3],[4,5],[6]].
  int set_cols_size = req->hashed_column_values().size();
  auto unset_cols_size = hash_values_options_->size();

  // Initialize the missing columns with default values (e.g. h2, h3, h4 in example above).
  req->mutable_hashed_column_values()->Reserve(narrow_cast<int>(set_cols_size + unset_cols_size));
  for (size_t i = 0; i < unset_cols_size; i++885) {
    req->add_hashed_column_values();
  }

  // Set the right values for the missing/unset columns by converting partition index into positions
  // for each hash column and using the corresponding values from the hash values options vector.
  // E.g. In example above, with start_partition = 0:
  //    h4 = 6 since pos is "0 % 1 = 0", (start_position becomes 0 / 1 = 0).
  //    h3 = 4 since pos is "0 % 2 = 0", (start_position becomes 0 / 2 = 0).
  //    h2 = 2 since pos is "0 % 2 = 0", (start_position becomes 0 / 2 = 0).
  for (auto i = unset_cols_size; i > 0;) {
    --i;
    const auto& options = (*hash_values_options_)[i];
    auto pos = start_partition % options.size();
    *req->mutable_hashed_column_values(narrow_cast<int>(i + set_cols_size)) = options[pos];
    start_partition /= options.size();
  }
}

bool TnodeContext::FinishedReadingPartition() {
  return rows_result_->paging_state().empty() ||
      (34.4kquery_state_->next_partition_key().empty()34.4k && query_state_->next_row_key().empty()305);
}

void TnodeContext::AdvanceToNextPartition(QLReadRequestPB *req) {
  // E.g. for a query "h1 = 1 and h2 in (2,3) and h3 in (4,5) and h4 = 6" partition index 2:
  // this will do, index: 2 -> 3 and hashed_column_values(): [1, 3, 4, 6] -> [1, 3, 5, 6].
  current_partition_index_++;
  uint64_t partition_counter = current_partition_index_;
  // Hash_values_options_ vector starts from the first column with an 'IN' restriction.
  const int hash_key_size = req->hashed_column_values().size();
  const auto fixed_cols_size = hash_key_size - hash_values_options_->size();

  // Set the right values for the missing/unset columns by converting partition index into positions
  // for each hash column and using the corresponding values from the hash values options vector.
  // E.g. In example above, with start_partition = 3:
  //    h4 = 6 since pos is "3 % 1 = 0", new partition counter is "3 / 1 = 3".
  //    h3 = 5 since pos is "3 % 2 = 1", pos is non-zero which guarantees previous cols don't need
  //    to be changed (i.e. are the same as for previous partition index) so we break.
  for (size_t i = hash_key_size; i > fixed_cols_size;) {
    --i;
    const auto& options = (*hash_values_options_)[i - fixed_cols_size];
    auto pos = partition_counter % options.size();
    *req->mutable_hashed_column_values(narrow_cast<int>(i)) = options[pos];
    if (pos != 0) break5.44k; // The previous position hash values must be unchanged.
    partition_counter /= options.size();
  }

  req->clear_hash_code();
  req->clear_max_hash_code();

  if (hash_code_from_partition_key_ops_.is_initialized())
    req->set_hash_code(*hash_code_from_partition_key_ops_);
  if (max_hash_code_from_partition_key_ops_.is_initialized())
    req->set_max_hash_code(*max_hash_code_from_partition_key_ops_);
}

bool TnodeContext::HasPendingOperations() const {
  for (const auto& op : ops_) {
    if (!op->response().has_status()) {
      return true;
    }
  }
  if (child_context_) {
    return child_context_->HasPendingOperations();
  }
  return false;
}

void TnodeContext::SetUncoveredSelectOp(const YBqlReadOpPtr& select_op) {
  uncovered_select_op_ = select_op;
  const Schema& schema = static_cast<const PTSelectStmt*>(tnode_)->table()->InternalSchema();
  std::vector<ColumnId> key_column_ids;
  key_column_ids.reserve(schema.num_key_columns());
  for (size_t idx = 0; idx < schema.num_key_columns(); idx++1.08k) {
    key_column_ids.emplace_back(schema.column_id(idx));
  }
  keys_ = std::make_unique<QLRowBlock>(schema, key_column_ids);
}

QueryPagingState *TnodeContext::CreateQueryState(const StatementParameters& user_params,
                                                 bool is_top_level_select) {
  query_state_ = std::make_unique<QueryPagingState>(user_params, is_top_level_select);
  return query_state_.get();
}

Status TnodeContext::ClearQueryState() {
  RSTATUS_DCHECK(query_state_, Corruption, "Query state should not be null for SELECT");
  rows_result_->ClearPagingState();
  query_state_->ClearPagingState();

  return Status::OK();
}

Status TnodeContext::ComposeRowsResultForUser(const TreeNode* child_select_node,
                                              bool for_new_batches) {
  RSTATUS_DCHECK_EQ(tnode_->opcode(), TreeNodeOpcode::kPTSelectStmt,
                    Corruption, "Only SELECT node can have nested query");
  const auto* select_stmt = static_cast<const PTSelectStmt *>(tnode_);

  // Case 1:
  //   SELECT * FROM <table>;
  if (!child_select_node) {
    if (rows_result_->has_paging_state() || for_new_batches2) {
      // Paging state must be provided for two cases. Otherwise, we've reached end of result set.
      // - Docdb sent back rows_result with paging state.
      // - Seting up paging_state for user's next batches.
      RETURN_NOT_OK(query_state_->ComposePagingStateForUser());
      rows_result_->SetPagingState(query_state_->query_pb());
    }
    return Status::OK();
  }

  // Check for nested condition.
  RSTATUS_DCHECK(child_context_ && child_select_node->opcode() == TreeNodeOpcode::kPTSelectStmt,
                 Corruption, "Expecting nested context with a SELECT node");

  // Case 2:
  //   SELECT <fully_covered_columns> FROM <index>;
  // Move result from index query (child) to the table query (this parent node).
  const auto* child_select = static_cast<const PTSelectStmt *>(child_select_node);
  if (child_select->covers_fully()) {
    return AppendRowsResult(std::move(child_context_->rows_result()));
  }

  // Case 3:
  //   SELECT <any columns> FROM <table> WHERE primary_keys IN (SELECT primary_keys FROM <index>);
  // Compose result of the following fields.
  // - The rows_result should be from this node (rows_result_).
  // - The counter_state should be from this node (query_state_::counter_pb_).
  // - The read paging_state should be from the CHILD node (query_state_::query_pb_)
  if (!rows_result_) {
    // Allocate an empty rows_result that will be filled with paging state.
    rows_result_ = std::make_shared<RowsResult>(select_stmt);
  }

  if (child_context_->rows_result()->has_paging_state() &&
      !query_state_->reached_select_limit()60) {
    // If child node has paging state and LIMIT is not yet reached, provide paging state to users
    // to continue reading.
    RETURN_NOT_OK(
        query_state_->ComposePagingStateForUser(child_context_->query_state()->query_pb()));
    rows_result_->SetPagingState(query_state_->query_pb());
  } else {
    // Clear paging state once all requested rows were retrieved.
    rows_result_->ClearPagingState();
  }

  return Status::OK();
}

//--------------------------------------------------------------------------------------------------

QueryPagingState::QueryPagingState(const StatementParameters& user_params,
                                   bool is_top_level_read_node)
    : max_fetch_size_(user_params.page_size()) {
  LoadPagingStateFromUser(user_params, is_top_level_read_node);

  // Just default it to max_int.
  if (max_fetch_size_ <= 0) {
    max_fetch_size_ = INT_MAX;
  }
}

void QueryPagingState::AdjustMaxFetchSizeToSelectLimit() {
  int64_t limit = select_limit();
  if (limit < 0) {
    return;
  }

  int64_t count = read_count();
  if (count < limit7.18M) {
    int64_t wanted = limit - count;
    if (wanted < max_fetch_size_7.20M) {
      max_fetch_size_ = wanted;
    }
  } else {
    max_fetch_size_ = 0;
  }
}

void QueryPagingState::ClearPagingState() {
  // Clear only the paging state.
  // Keep the counter so that we knows how many rows have been processed.
  query_pb_.Clear();
}

void QueryPagingState::LoadPagingStateFromUser(const StatementParameters& user_params,
                                               bool is_top_level_read_node) {
  user_params.WritePagingState(&query_pb_);

  // Calculate "skip_count" and "read_count".
  // (1) Top level read node.
  //  - Either top-level SELECT or fully-covering INDEX query.
  //  - User "params::couter_pb_" should have the valid "counter_pb_"
  //
  // (2) Nested read node.
  //  - Zero out its counters. We don't use "counter_pb_" for this node because LIMIT and OFFSET
  //    restrictions are not applied to nested nodes.
  //  - Because nested node might be running different READ operators (with different hash values)
  //    for different calls from users, the counters from users' message are discarded here.
  if (is_top_level_read_node) {
    counter_pb_.CopyFrom(query_pb_.row_counter());
  } else {
    // These values are not used, set them to zero.
    set_skip_count(0);
    set_read_count(0);
  }
}

Status QueryPagingState::ComposePagingStateForUser() {
  // Write the counters into the paging_state.
  query_pb_.mutable_row_counter()->CopyFrom(counter_pb_);
  return Status::OK();
}

Status QueryPagingState::ComposePagingStateForUser(const QLPagingStatePB& child_state) {
  // Write child_state.
  query_pb_.CopyFrom(child_state);

  // Write the counters into the paging_state.
  query_pb_.mutable_row_counter()->CopyFrom(counter_pb_);

  return Status::OK();
}

Status QueryPagingState::LoadPagingStateFromDocdb(const RowsResult::SharedPtr& rows_result,
                                                  int64_t number_of_new_rows,
                                                  bool has_nested_query) {
  // Load "query_pb_" with the latest result from DocDB.
  query_pb_.ParseFromString(rows_result->paging_state());

  // If DocDB processed the skipping rows, record it here.
  if (total_rows_skipped() > 0) {
    set_skip_count(total_rows_skipped());
  }

  if (!has_nested_query7.52M) {
    // SELECT <select_list> FROM <table or index>
    //   WHERE <filter_cond>
    //   LIMIT <limit> OFFSET <offset>
    // - DocDB processed the <limit> and <offset> restrictions.
    //   Either "reached_select_offset() == TRUE" OR number_of_new_rows == 0.
    // - Two DocDB::counters are used to compute here.
    //   . QLPagingStatePB::total_rows_skipped - Skip count in DocDB.
    //   . number_of_new_rows - Rows of data from DocDB after skipping.
    set_read_count(read_count() + number_of_new_rows);

  } else {
    // SELECT <select_list> FROM <table>
    //   WHERE
    //     primary_key IN (SELECT primary_key FROM <index> WHERE <index_cond>)
    //     AND
    //     <filter_cond>
    //   LIMIT <limit> OFFSET <offset>
    //
    // NOTE:
    // 1. Case INDEX query fully-covers the SELECT command.
    //    - DocDB counters are transfered from nested node to this node.
    //    - "reached_select_offset() == TRUE" OR number_of_new_rows == 0.
    //
    // 2. Case INDEX query does NOT fully-cover the SELECT command.
    //    - Values of <limit> and <offset> are NOT sent together with proto request to DocDB. They
    //      are computed and processed here in CQL layer.
    //    - For this case, "number_of_new_rows" is either 1 or 0.
    //      CQL assumes that outer SELECT reads at most one row at a time as it uses values of
    //      PRIMARY KEY (always unique) to read the rest of the columns of the <table>.
    if (!reached_select_offset()) {
      // Since OFFSET is processed here, this must be case 2.
      RSTATUS_DCHECK_LE(number_of_new_rows, 1, Corruption, "Incorrect counter calculation");
      set_skip_count(skip_count() + number_of_new_rows);
    } else {
      set_read_count(read_count() + number_of_new_rows);
    }
  }

  return Status::OK();
}

const std::string& ExecContext::stmt() const {
  return parse_tree_.stmt();
}

}  // namespace ql
}  // namespace yb

YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

Coverage Report

Created: 2022-03-22 16:43