blob: 01492f9043736569aa26f213e641a3d6dc9828e2 [file] [log] [blame]
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
#include <atomic>
#include <cstddef>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include "query_execution/QueryExecutionTypedefs.hpp"
#include "storage/StorageBlockInfo.hpp"
#include "storage/StorageConstants.hpp"
#include "threading/SpinSharedMutex.hpp"
#include "threading/Thread.hpp"
#include "utility/Macros.hpp"
#include "glog/logging.h"
#include "tmb/id_typedefs.h"
#include "tmb/message_bus.h"
namespace quickstep {
/** \addtogroup QueryExecution
* @{
* @brief A class for keeping trace of blocks loaded in a Worker's buffer pool
* in the distributed version.
class BlockLocator : public Thread {
* @brief Constructor.
* @param bus A pointer to the TMB.
* @param cpu_id The ID of the CPU to which the BlockLocator thread can be pinned.
* @note If cpu_id is not specified, BlockLocator thread can be possibly moved
* around on different CPUs by the OS.
BlockLocator(tmb::MessageBus *bus,
const int cpu_id = -1)
: bus_(DCHECK_NOTNULL(bus)),
block_domain_(0) {
locator_client_id_ = bus_->Connect();
bus_->RegisterClientAsReceiver(locator_client_id_, kBlockDomainRegistrationMessage);
bus_->RegisterClientAsSender(locator_client_id_, kBlockDomainRegistrationResponseMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kBlockDomainToShiftbossIndexMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kAddBlockLocationMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kDeleteBlockLocationMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kGetAllDomainNetworkAddressesMessage);
bus_->RegisterClientAsSender(locator_client_id_, kGetAllDomainNetworkAddressesResponseMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kBlockDomainUnregistrationMessage);
bus_->RegisterClientAsReceiver(locator_client_id_, kPoisonMessage);
~BlockLocator() override {}
* @brief Get the TMB client ID of BlockLocator thread.
* @return TMB client ID of BlockLocator thread.
tmb::client_id getBusClientID() const {
return locator_client_id_;
* @brief Get the block domain info for a given block.
* @param block The given block.
* @return The block domain info for a given block.
std::unordered_set<block_id_domain> getBlockDomains(const block_id block) const {
// Lock 'block_locations_shared_mutex_' as briefly as possible as a
// reader.
SpinSharedMutexSharedLock<false> read_lock(block_locations_shared_mutex_);
const auto cit = block_locations_.find(block);
if (cit != block_locations_.end()) {
return cit->second;
return std::unordered_set<block_id_domain>();
* @brief Get the block locality info for scheduling in ForemanDistributed.
* @param block The given block.
* @param shiftboss_index_for_block The index of Shiftboss that has loaded the
* block in the buffer pool.
* @return Whether the block locality info has found.
bool getBlockLocalityInfo(const block_id block, std::size_t *shiftboss_index_for_block) const {
if (block == kInvalidBlockId) { return false; }
const std::unordered_set<block_id_domain> block_domains = getBlockDomains(block);
if (!block_domains.empty()) {
// NOTE(zuyu): This lock is held for the rest duration of this call, as the
// exclusive case is rare.
SpinSharedMutexSharedLock<false> read_lock(block_domain_to_shiftboss_index_shared_mutex_);
for (const block_id_domain block_domain : block_domains) {
// TODO(quickstep-team): choose the best node, instead of the first.
const auto cit = block_domain_to_shiftboss_index_.find(block_domain);
if (cit != block_domain_to_shiftboss_index_.end()) {
*shiftboss_index_for_block = cit->second;
return true;
return false;
void run() override;
void processBlockDomainRegistrationMessage(const tmb::client_id receiver, const std::string &network_address);
void processGetAllDomainNetworkAddressesMessage(const tmb::client_id receiver);
tmb::MessageBus *bus_;
// The ID of the CPU that the BlockLocator thread can optionally be pinned to.
const int cpu_id_;
alignas(kCacheLineBytes) std::atomic<block_id_domain> block_domain_;
// From a block domain to its network info in the ip:port format, i.e.,
// "".
std::unordered_map<block_id_domain, const std::string> domain_network_addresses_;
// From a block domain to its Shiftboss index, used by ForemanDistributed
// to schedule based on the data-locality info.
// Note that not every 'block_id_domain' has a Shiftboss index. For example,
// DistributedCli has StorageManager with a 'block_id_domain', which is not
// a part of Shiftboss.
std::unordered_map<block_id_domain, std::size_t> block_domain_to_shiftboss_index_;
alignas(kCacheLineBytes) mutable SpinSharedMutex<false> block_domain_to_shiftboss_index_shared_mutex_;
// From a block to its domains.
std::unordered_map<block_id, std::unordered_set<block_id_domain>> block_locations_;
alignas(kCacheLineBytes) mutable SpinSharedMutex<false> block_locations_shared_mutex_;
// From a block domain to all blocks loaded in its buffer pool.
std::unordered_map<block_id_domain, std::unordered_set<block_id>> domain_blocks_;
tmb::client_id locator_client_id_;
/** @} */
} // namespace quickstep