ClickHouse/dbms/include/DB/DataStreams/RemoteBlockInputStream.h

341 lines
12 KiB
C
Raw Normal View History

2012-05-17 19:15:53 +00:00
#pragma once
#include <Yandex/logger_useful.h>
2012-05-17 19:15:53 +00:00
#include <DB/DataStreams/IProfilingBlockInputStream.h>
#include <DB/DataStreams/OneBlockInputStream.h>
#include <DB/Common/VirtualColumnUtils.h>
#include <DB/Common/Throttler.h>
#include <DB/Interpreters/Context.h>
2012-05-17 19:15:53 +00:00
#include <DB/Client/ConnectionPool.h>
#include <DB/Client/ParallelReplicas.h>
2012-05-17 19:15:53 +00:00
namespace DB
{
/** Позволяет выполнить запрос (SELECT) на удалённых репликах одного шарда и получить результат.
2012-05-17 19:15:53 +00:00
*/
class RemoteBlockInputStream : public IProfilingBlockInputStream
{
private:
void init(const Settings * settings_)
{
if (settings_)
{
send_settings = true;
settings = *settings_;
}
else
send_settings = false;
}
public:
/// Принимает готовое соединение.
RemoteBlockInputStream(Connection & connection_, const String & query_, const Settings * settings_, ThrottlerPtr throttler_ = nullptr,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete,
const Context & context = getDefaultContext())
: connection(&connection_), query(query_), throttler(throttler_), external_tables(external_tables_), stage(stage_), context(context)
{
init(settings_);
}
/// Принимает готовое соединение. Захватывает владение соединением из пула.
RemoteBlockInputStream(ConnectionPool::Entry & pool_entry_, const String & query_, const Settings * settings_, ThrottlerPtr throttler_ = nullptr,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete,
const Context & context = getDefaultContext())
: pool_entry(pool_entry_), connection(&*pool_entry_), query(query_), throttler(throttler_),
external_tables(external_tables_), stage(stage_), context(context)
{
init(settings_);
}
/// Принимает пул, из которого нужно будет достать одно или несколько соединений.
RemoteBlockInputStream(IConnectionPool * pool_, const String & query_, const Settings * settings_, ThrottlerPtr throttler_ = nullptr,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete,
const Context & context = getDefaultContext())
: pool(pool_), query(query_), throttler(throttler_), external_tables(external_tables_), stage(stage_), context(context)
{
init(settings_);
}
2012-10-18 19:24:46 +00:00
String getName() const override { return "RemoteBlockInputStream"; }
2012-05-17 19:15:53 +00:00
String getID() const override
{
2013-05-04 04:05:15 +00:00
std::stringstream res;
res << this;
return res.str();
}
2012-10-18 19:24:46 +00:00
2012-05-17 19:15:53 +00:00
/** Отменяем умолчальное уведомление о прогрессе,
* так как колбэк прогресса вызывается самостоятельно.
*/
void progress(const Progress & value) override {}
2012-05-17 19:15:53 +00:00
2012-10-18 19:24:46 +00:00
void cancel() override
2012-10-18 19:24:46 +00:00
{
2012-11-10 05:13:46 +00:00
if (!__sync_bool_compare_and_swap(&is_cancelled, false, true))
return;
if (hasNoQueryInProgress() || hasThrownException())
return;
2015-03-06 01:06:11 +00:00
tryCancel("Cancelling query");
2012-10-18 19:24:46 +00:00
}
~RemoteBlockInputStream() override
{
/** Если прервались в середине цикла общения с репликами, то прервываем
* все соединения, затем читаем и пропускаем оставшиеся пакеты чтобы
* эти соединения не остались висеть в рассихронизированном состоянии.
*/
if (established || isQueryInProgress())
parallel_replicas->disconnect();
}
2012-10-20 02:10:47 +00:00
protected:
/// Отправить на удаленные реплики все временные таблицы
void sendExternalTables()
{
size_t count = parallel_replicas->size();
std::vector<ExternalTablesData> instances;
instances.reserve(count);
for (size_t i = 0; i < count; ++i)
{
ExternalTablesData res;
for (const auto & table : external_tables)
{
StoragePtr cur = table.second;
QueryProcessingStage::Enum stage = QueryProcessingStage::Complete;
DB::BlockInputStreams input = cur->read(cur->getColumnNamesList(), ASTPtr(), context, settings,
stage, DEFAULT_BLOCK_SIZE, 1);
if (input.size() == 0)
res.push_back(std::make_pair(new OneBlockInputStream(cur->getSampleBlock()), table.first));
else
res.push_back(std::make_pair(input[0], table.first));
}
instances.push_back(std::move(res));
}
parallel_replicas->sendExternalTablesData(instances);
}
Block readImpl() override
2012-10-20 02:10:47 +00:00
{
if (!sent_query)
2012-10-20 02:10:47 +00:00
{
createParallelReplicas();
established = true;
parallel_replicas->sendQuery(query, "", stage, true);
established = false;
sent_query = true;
sendExternalTables();
2012-10-20 02:10:47 +00:00
}
while (true)
{
Connection::Packet packet = parallel_replicas->receivePacket();
2012-10-20 02:10:47 +00:00
switch (packet.type)
{
case Protocol::Server::Data:
/// Если блок не пуст и не является заголовочным блоком
if (packet.block && packet.block.rows() > 0)
2012-10-20 02:10:47 +00:00
return packet.block;
break; /// Если блок пустой - получим другие пакеты до EndOfStream.
case Protocol::Server::Exception:
got_exception_from_replica = true;
2012-10-20 02:10:47 +00:00
packet.exception->rethrow();
break;
case Protocol::Server::EndOfStream:
if (!parallel_replicas->hasActiveReplicas())
{
finished = true;
return Block();
}
break;
2012-10-20 02:10:47 +00:00
case Protocol::Server::Progress:
/** Используем прогресс с удалённого сервера.
* В том числе, запишем его в ProcessList,
* и будем использовать его для проверки
* ограничений (например, минимальная скорость выполнения запроса)
* и квот (например, на количество строчек для чтения).
*/
progressImpl(packet.progress);
2012-10-20 02:10:47 +00:00
if (isQueryInProgress() && isCancelled())
2012-10-20 02:10:47 +00:00
cancel();
break;
2013-06-01 18:57:04 +00:00
case Protocol::Server::ProfileInfo:
info = packet.profile_info;
2013-06-01 18:57:04 +00:00
break;
case Protocol::Server::Totals:
totals = packet.block;
break;
case Protocol::Server::Extremes:
extremes = packet.block;
break;
2012-10-20 02:10:47 +00:00
default:
got_unknown_packet_from_replica = true;
2012-10-20 02:10:47 +00:00
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
}
}
}
void readSuffixImpl() override
2013-09-13 20:33:09 +00:00
{
/** Если одно из:
* - ничего не начинали делать;
* - получили все пакеты до EndOfStream;
* - получили с одной реплики эксепшен;
* - получили с одной реплики неизвестный пакет;
2013-09-13 20:33:09 +00:00
* - то больше читать ничего не нужно.
*/
if (hasNoQueryInProgress() || hasThrownException())
2013-09-13 20:33:09 +00:00
return;
/** Если ещё прочитали не все данные, но они больше не нужны.
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT).
*/
/// Отправим просьбу прервать выполнение запроса, если ещё не отправляли.
2015-03-06 01:06:11 +00:00
tryCancel("Cancelling query because enough data has been read");
2013-09-13 20:33:09 +00:00
/// Получим оставшиеся пакеты, чтобы не было рассинхронизации в соединениях с репликами.
Connection::Packet packet = parallel_replicas->drain();
switch (packet.type)
{
case Protocol::Server::EndOfStream:
finished = true;
break;
case Protocol::Server::Exception:
got_exception_from_replica = true;
packet.exception->rethrow();
break;
default:
got_unknown_packet_from_replica = true;
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
2013-09-13 20:33:09 +00:00
}
}
/// Создать объект для общения с репликами одного шарда, на которых должен выполниться запрос.
void createParallelReplicas()
{
Settings * parallel_replicas_settings = send_settings ? &settings : nullptr;
if (connection != nullptr)
2015-02-10 21:10:58 +00:00
parallel_replicas = std::make_unique<ParallelReplicas>(connection, parallel_replicas_settings, throttler);
else
2015-02-10 21:10:58 +00:00
parallel_replicas = std::make_unique<ParallelReplicas>(pool, parallel_replicas_settings, throttler);
}
/// Возвращает true, если запрос отправлен, а ещё не выполнен.
bool isQueryInProgress() const
{
return sent_query && !finished && !was_cancelled;
}
/// Возвращает true, если никакой запрос не отправлен или один запрос уже выполнен.
bool hasNoQueryInProgress() const
{
return !sent_query || finished;
}
/// Возвращает true, если исключение было выкинуто.
bool hasThrownException() const
{
return got_exception_from_replica || got_unknown_packet_from_replica;
}
private:
/// ITable::read requires a Context, therefore we should create one if the user can't supply it
static Context & getDefaultContext()
{
static Context instance;
return instance;
}
/// Отправить запрос на отмену всех соединений к репликам, если такой запрос ещё не был отправлен.
2015-03-06 01:06:11 +00:00
void tryCancel(const char * reason)
{
bool old_val = false;
2015-03-06 01:06:11 +00:00
if (!was_cancelled.compare_exchange_strong(old_val, true, std::memory_order_seq_cst, std::memory_order_seq_cst))
return;
2015-03-06 01:06:11 +00:00
LOG_TRACE(log, "(" << parallel_replicas->dumpAddresses() << ") " << reason);
parallel_replicas->sendCancel();
}
2012-05-17 19:15:53 +00:00
private:
IConnectionPool * pool = nullptr;
ConnectionPool::Entry pool_entry;
Connection * connection = nullptr;
std::unique_ptr<ParallelReplicas> parallel_replicas;
2012-05-23 19:51:30 +00:00
const String query;
bool send_settings;
Settings settings;
/// Если не nullptr, то используется, чтобы ограничить сетевой трафик.
ThrottlerPtr throttler;
/// Временные таблицы, которые необходимо переслать на удаленные сервера.
Tables external_tables;
2012-05-17 19:15:53 +00:00
QueryProcessingStage::Enum stage;
Context context;
2012-05-17 19:15:53 +00:00
/// Установили соединения с репликами, но ещё не отправили запрос.
std::atomic<bool> established { false };
2012-10-18 19:00:49 +00:00
/// Отправили запрос (это делается перед получением первого блока).
std::atomic<bool> sent_query { false };
/** Получили все данные от всех реплик, до пакета EndOfStream.
2012-10-18 19:00:49 +00:00
* Если при уничтожении объекта, ещё не все данные считаны,
* то для того, чтобы не было рассинхронизации, на реплики отправляются просьбы прервать выполнение запроса,
2012-10-18 19:00:49 +00:00
* и после этого считываются все пакеты до EndOfStream.
*/
std::atomic<bool> finished { false };
/** На каждую реплику была отправлена просьба прервать выполнение запроса, так как данные больше не нужны.
2012-10-18 19:00:49 +00:00
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT),
* или если на стороне клиента произошло исключение.
*/
std::atomic<bool> was_cancelled { false };
/** С одной репилки было получено исключение. В этом случае получать больше пакетов или
* просить прервать запрос на этой реплике не нужно.
*/
std::atomic<bool> got_exception_from_replica { false };
/** С одной реплики был получен неизвестный пакет. В этом случае получать больше пакетов или
* просить прервать запрос на этой реплике не нужно.
*/
std::atomic<bool> got_unknown_packet_from_replica { false };
2012-10-18 19:00:49 +00:00
Logger * log = &Logger::get("RemoteBlockInputStream");
2012-05-17 19:15:53 +00:00
};
}