From d6a6d280dd7ac5ca047d7b7d6f6f97f16acd4e95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?= Date: Tue, 22 Jun 2021 08:43:19 +0200 Subject: [PATCH] Add Streams on top of Kafka Consumer (#172) * Stop the Consumer grafefully when it is destroyed * Add Streams * Add Streams to InterpreterContext * Remove options to limit processed batches in Consumer * Add Streams unit tests * Stop waiting for a full batch if the Consumer stopped * Add ReadLock functionality to Synchronized * Use per Consumer-based locking * Replace shared_mutex with RWLock --- .clang-tidy | 2 + src/integrations/kafka/CMakeLists.txt | 2 +- src/integrations/kafka/consumer.cpp | 38 +-- src/integrations/kafka/consumer.hpp | 18 +- src/memgraph.cpp | 10 +- src/query/CMakeLists.txt | 4 +- src/query/interpreter.cpp | 7 +- src/query/interpreter.hpp | 5 +- src/query/streams.cpp | 287 +++++++++++++++++++++ src/query/streams.hpp | 149 +++++++++++ src/utils/rw_lock.hpp | 5 + src/utils/synchronized.hpp | 42 ++- tests/benchmark/expansion.cpp | 2 +- tests/manual/single_query.cpp | 3 +- tests/unit/CMakeLists.txt | 29 ++- tests/unit/integrations_kafka_consumer.cpp | 129 ++++++--- tests/unit/interpreter.cpp | 3 +- tests/unit/kafka_mock.cpp | 18 +- tests/unit/kafka_mock.hpp | 1 + tests/unit/query_dump.cpp | 6 +- tests/unit/query_plan_edge_cases.cpp | 2 +- tests/unit/query_streams.cpp | 231 +++++++++++++++++ tests/unit/utils_synchronized.cpp | 126 ++++++++- 23 files changed, 1007 insertions(+), 112 deletions(-) create mode 100644 src/query/streams.cpp create mode 100644 src/query/streams.hpp create mode 100644 tests/unit/query_streams.cpp diff --git a/.clang-tidy b/.clang-tidy index 5e357feba..3abb221ad 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -35,8 +35,10 @@ Checks: '*, -hicpp-no-assembler, -hicpp-no-malloc, -hicpp-use-equals-default, + -hicpp-use-nullptr, -hicpp-vararg, -llvm-header-guard, + -llvm-include-order, -llvmlibc-callee-namespace, -llvmlibc-implementation-in-namespace, -llvmlibc-restrict-system-libc-headers, diff --git a/src/integrations/kafka/CMakeLists.txt b/src/integrations/kafka/CMakeLists.txt index 7a1bd442f..98d25db19 100644 --- a/src/integrations/kafka/CMakeLists.txt +++ b/src/integrations/kafka/CMakeLists.txt @@ -3,4 +3,4 @@ set(integrations_kafka_src_files ) add_library(mg-integrations-kafka STATIC ${integrations_kafka_src_files}) -target_link_libraries(mg-integrations-kafka mg-utils librdkafka++ librdkafka Threads::Threads zlib) +target_link_libraries(mg-integrations-kafka mg-utils librdkafka++ librdkafka Threads::Threads) diff --git a/src/integrations/kafka/consumer.cpp b/src/integrations/kafka/consumer.cpp index d63a71306..37b5d22d8 100644 --- a/src/integrations/kafka/consumer.cpp +++ b/src/integrations/kafka/consumer.cpp @@ -44,8 +44,9 @@ int64_t Message::Timestamp() const { return rd_kafka_message_timestamp(c_message, nullptr); } -Consumer::Consumer(ConsumerInfo info) : info_{std::move(info)} { - MG_ASSERT(info_.consumer_function, "Empty consumer function for Kafka consumer"); +Consumer::Consumer(const std::string &bootstrap_servers, ConsumerInfo info, ConsumerFunction consumer_function) + : info_{std::move(info)}, consumer_function_(std::move(consumer_function)) { + MG_ASSERT(consumer_function_, "Empty consumer function for Kafka consumer"); std::unique_ptr conf(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL)); if (conf == nullptr) { throw ConsumerFailedToInitializeException(info_.consumer_name, "Couldn't create Kafka configuration!"); @@ -65,7 +66,7 @@ Consumer::Consumer(ConsumerInfo info) : info_{std::move(info)} { throw ConsumerFailedToInitializeException(info_.consumer_name, error); } - if (conf->set("bootstrap.servers", info_.bootstrap_servers, error) != RdKafka::Conf::CONF_OK) { + if (conf->set("bootstrap.servers", bootstrap_servers, error) != RdKafka::Conf::CONF_OK) { throw ConsumerFailedToInitializeException(info_.consumer_name, error); } @@ -107,17 +108,19 @@ Consumer::Consumer(ConsumerInfo info) : info_{std::move(info)} { } } -void Consumer::Start(std::optional limit_batches) { +Consumer::~Consumer() { StopIfRunning(); } + +void Consumer::Start() { if (is_running_) { throw ConsumerRunningException(info_.consumer_name); } - StartConsuming(limit_batches); + StartConsuming(); } void Consumer::StartIfStopped() { if (!is_running_) { - StartConsuming(std::nullopt); + StartConsuming(); } } @@ -133,6 +136,9 @@ void Consumer::StopIfRunning() { if (is_running_) { StopConsuming(); } + if (thread_.joinable()) { + thread_.join(); + } } void Consumer::Test(std::optional limit_batches, const ConsumerFunction &test_consumer_function) { @@ -199,6 +205,8 @@ void Consumer::Test(std::optional limit_batches, const ConsumerFunction bool Consumer::IsRunning() const { return is_running_; } +const ConsumerInfo &Consumer::Info() const { return info_; } + void Consumer::event_cb(RdKafka::Event &event) { switch (event.type()) { case RdKafka::Event::Type::EVENT_ERROR: @@ -210,7 +218,7 @@ void Consumer::event_cb(RdKafka::Event &event) { break; } } -void Consumer::StartConsuming(std::optional limit_batches) { +void Consumer::StartConsuming() { MG_ASSERT(!is_running_, "Cannot start already running consumer!"); if (thread_.joinable()) { @@ -221,20 +229,18 @@ void Consumer::StartConsuming(std::optional limit_batches) { is_running_.store(true); - thread_ = std::thread([this, limit_batches]() { + thread_ = std::thread([this] { constexpr auto kMaxThreadNameSize = utils::GetMaxThreadNameSize(); const auto full_thread_name = "Cons#" + info_.consumer_name; utils::ThreadSetName(full_thread_name.substr(0, kMaxThreadNameSize)); - int64_t batch_count = 0; - while (is_running_) { auto maybe_batch = this->GetBatch(); if (maybe_batch.HasError()) { spdlog::warn("Error happened in consumer {} while fetching messages: {}!", info_.consumer_name, maybe_batch.GetError()); - is_running_.store(false); + break; } const auto &batch = maybe_batch.GetValue(); @@ -244,18 +250,14 @@ void Consumer::StartConsuming(std::optional limit_batches) { // TODO (mferencevic): Figure out what to do with all other exceptions. try { - info_.consumer_function(batch); + consumer_function_(batch); consumer_->commitSync(); } catch (const utils::BasicException &e) { spdlog::warn("Error happened in consumer {} while processing a batch: {}!", info_.consumer_name, e.what()); break; } - - if (limit_batches != std::nullopt && limit_batches <= ++batch_count) { - is_running_.store(false); - break; - } } + is_running_.store(false); }); } @@ -274,7 +276,7 @@ utils::BasicResult> Consumer::GetBatch() { auto start = std::chrono::steady_clock::now(); bool run_batch = true; - for (int64_t i = 0; remaining_timeout_in_ms > 0 && i < batch_size; ++i) { + for (int64_t i = 0; remaining_timeout_in_ms > 0 && i < batch_size && is_running_.load(); ++i) { std::unique_ptr msg(consumer_->consume(remaining_timeout_in_ms)); switch (msg->err()) { case RdKafka::ERR__TIMED_OUT: diff --git a/src/integrations/kafka/consumer.hpp b/src/integrations/kafka/consumer.hpp index 3874725e4..79ae8518d 100644 --- a/src/integrations/kafka/consumer.hpp +++ b/src/integrations/kafka/consumer.hpp @@ -65,9 +65,7 @@ using ConsumerFunction = std::function &)>; /// ConsumerInfo holds all the information necessary to create a Consumer. struct ConsumerInfo { - ConsumerFunction consumer_function; std::string consumer_name; - std::string bootstrap_servers; std::vector topics; std::string consumer_group; std::optional batch_interval; @@ -84,8 +82,8 @@ class Consumer final : public RdKafka::EventCb { /// /// @throws ConsumerFailedToInitializeException if the consumer can't connect /// to the Kafka endpoint. - explicit Consumer(ConsumerInfo info); - ~Consumer() override = default; + explicit Consumer(const std::string &bootstrap_servers, ConsumerInfo info, ConsumerFunction consumer_function); + ~Consumer() override; Consumer(const Consumer &other) = delete; Consumer(Consumer &&other) noexcept = delete; @@ -96,10 +94,8 @@ class Consumer final : public RdKafka::EventCb { /// /// This method will start a new thread which will poll all the topics for messages. /// - /// @param limit_batches if present, the consumer will only consume the given number of batches and stop afterwards. - /// /// @throws ConsumerRunningException if the consumer is already running - void Start(std::optional limit_batches); + void Start(); /// Starts consuming messages if it is not started already. /// @@ -130,20 +126,22 @@ class Consumer final : public RdKafka::EventCb { /// Returns true if the consumer is actively consuming messages. bool IsRunning() const; + const ConsumerInfo &Info() const; + private: void event_cb(RdKafka::Event &event) override; - void StartConsuming(std::optional limit_batches); + void StartConsuming(); void StopConsuming(); utils::BasicResult> GetBatch(); - // TODO(antaljanosbenjamin) Maybe split this to store only the necessary information ConsumerInfo info_; + ConsumerFunction consumer_function_; mutable std::atomic is_running_{false}; std::optional limit_batches_{std::nullopt}; - std::thread thread_; std::unique_ptr> consumer_; + std::thread thread_; }; } // namespace integrations::kafka diff --git a/src/memgraph.cpp b/src/memgraph.cpp index f9ab8de8b..05970eb05 100644 --- a/src/memgraph.cpp +++ b/src/memgraph.cpp @@ -166,6 +166,10 @@ DEFINE_bool(telemetry_enabled, false, "the database runtime (vertex and edge counts and resource usage) " "to allow for easier improvement of the product."); +// NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables) +DEFINE_string(kafka_bootstrap_servers, "", + "List of Kafka brokers as a comma separated list of broker host or host:port."); + // Audit logging flags. #ifdef MG_ENTERPRISE DEFINE_bool(audit_enabled, false, "Set to true to enable audit logging."); @@ -1065,7 +1069,8 @@ int main(int argc, char **argv) { query::InterpreterContext interpreter_context{ &db, {.query = {.allow_load_csv = FLAGS_allow_load_csv}, .execution_timeout_sec = FLAGS_query_execution_timeout_sec}, - FLAGS_data_directory}; + FLAGS_data_directory, + FLAGS_kafka_bootstrap_servers}; #ifdef MG_ENTERPRISE SessionData session_data{&db, &interpreter_context, &auth, &audit_log}; #else @@ -1084,6 +1089,9 @@ int main(int argc, char **argv) { &interpreter_context.ast_cache, &dba, &interpreter_context.antlr_lock, interpreter_context.config.query); } + // As the Stream transformations are using modules, they have to be restored after the query modules are loaded. + interpreter_context.streams.RestoreStreams(); + #ifdef MG_ENTERPRISE AuthQueryHandler auth_handler(&auth, std::regex(FLAGS_auth_user_or_role_name_regex)); #else diff --git a/src/query/CMakeLists.txt b/src/query/CMakeLists.txt index 812c04450..ba3119eea 100644 --- a/src/query/CMakeLists.txt +++ b/src/query/CMakeLists.txt @@ -32,6 +32,7 @@ set(mg_query_sources procedure/module.cpp procedure/py_module.cpp serialization/property_value.cpp + streams.cpp trigger.cpp trigger_context.cpp typed_value.cpp) @@ -39,9 +40,8 @@ set(mg_query_sources add_library(mg-query STATIC ${mg_query_sources}) add_dependencies(mg-query generate_lcp_query) target_include_directories(mg-query PUBLIC ${CMAKE_SOURCE_DIR}/include) -target_link_libraries(mg-query mg-integrations-kafka) target_link_libraries(mg-query dl cppitertools) -target_link_libraries(mg-query mg-storage-v2 mg-utils mg-kvstore) +target_link_libraries(mg-query mg-integrations-kafka mg-storage-v2 mg-utils mg-kvstore) if("${MG_PYTHON_VERSION}" STREQUAL "") find_package(Python3 3.5 REQUIRED COMPONENTS Development) else() diff --git a/src/query/interpreter.cpp b/src/query/interpreter.cpp index 17ca83879..727022368 100644 --- a/src/query/interpreter.cpp +++ b/src/query/interpreter.cpp @@ -604,8 +604,11 @@ using RWType = plan::ReadWriteTypeChecker::RWType; } // namespace InterpreterContext::InterpreterContext(storage::Storage *db, const InterpreterConfig config, - const std::filesystem::path &data_directory) - : db(db), trigger_store(data_directory / "triggers"), config(config) {} + const std::filesystem::path &data_directory, std::string kafka_bootstrap_servers) + : db(db), + trigger_store(data_directory / "triggers"), + config(config), + streams{this, std::move(kafka_bootstrap_servers), data_directory / "streams"} {} Interpreter::Interpreter(InterpreterContext *interpreter_context) : interpreter_context_(interpreter_context) { MG_ASSERT(interpreter_context_, "Interpreter context must not be NULL"); diff --git a/src/query/interpreter.hpp b/src/query/interpreter.hpp index 2eb5eace1..700809fa3 100644 --- a/src/query/interpreter.hpp +++ b/src/query/interpreter.hpp @@ -14,6 +14,7 @@ #include "query/plan/operator.hpp" #include "query/plan/read_write_type_checker.hpp" #include "query/stream.hpp" +#include "query/streams.hpp" #include "query/trigger.hpp" #include "query/typed_value.hpp" #include "storage/v2/isolation_level.hpp" @@ -150,7 +151,7 @@ struct PreparedQuery { */ struct InterpreterContext { explicit InterpreterContext(storage::Storage *db, InterpreterConfig config, - const std::filesystem::path &data_directory); + const std::filesystem::path &data_directory, std::string kafka_bootstrap_servers); storage::Storage *db; @@ -173,6 +174,8 @@ struct InterpreterContext { utils::ThreadPool after_commit_trigger_pool{1}; const InterpreterConfig config; + + query::Streams streams; }; /// Function that is used to tell all active interpreters that they should stop diff --git a/src/query/streams.cpp b/src/query/streams.cpp new file mode 100644 index 000000000..a613ec292 --- /dev/null +++ b/src/query/streams.cpp @@ -0,0 +1,287 @@ +#include "query/streams.hpp" + +#include +#include +#include + +#include +#include "query/interpreter.hpp" + +namespace query { + +namespace { +utils::SkipList::Iterator GetStream(utils::SkipList::Accessor &accessor, + const std::string &stream_name) { + auto it = accessor.find(stream_name); + if (it == accessor.end()) { + throw StreamsException("Couldn't find stream '{}'", stream_name); + } + return it; +} +} // namespace + +using Consumer = integrations::kafka::Consumer; +using ConsumerInfo = integrations::kafka::ConsumerInfo; +using Message = integrations::kafka::Message; + +// nlohmann::json doesn't support string_view access yet +const std::string kStreamName{"name"}; +const std::string kTopicsKey{"topics"}; +const std::string kConsumerGroupKey{"consumer_group"}; +const std::string kBatchIntervalKey{"batch_interval"}; +const std::string kBatchSizeKey{"batch_size"}; +const std::string kIsRunningKey{"is_running"}; + +void to_json(nlohmann::json &data, StreamStatus &&status) { + auto &info = status.info; + data[kStreamName] = std::move(status.name); + data[kTopicsKey] = std::move(info.topics); + data[kConsumerGroupKey] = info.consumer_group; + + if (info.batch_interval) { + data[kBatchIntervalKey] = info.batch_interval->count(); + } else { + data[kBatchIntervalKey] = nullptr; + } + + if (info.batch_size) { + data[kBatchSizeKey] = *info.batch_size; + } else { + data[kBatchSizeKey] = nullptr; + } + + data[kIsRunningKey] = status.is_running; +} + +void from_json(const nlohmann::json &data, StreamStatus &status) { + auto &info = status.info; + data.at(kStreamName).get_to(status.name); + data.at(kTopicsKey).get_to(info.topics); + data.at(kConsumerGroupKey).get_to(info.consumer_group); + + const auto batch_interval = data.at(kBatchIntervalKey); + if (!batch_interval.is_null()) { + using BatchInterval = decltype(info.batch_interval)::value_type; + info.batch_interval = BatchInterval{batch_interval.get()}; + } else { + info.batch_interval = {}; + } + + const auto batch_size = data.at(kBatchSizeKey); + if (!batch_size.is_null()) { + info.batch_size = batch_size.get(); + } else { + info.batch_size = {}; + } + + data.at(kIsRunningKey).get_to(status.is_running); +} + +bool operator==(const StreamData &lhs, const StreamData &rhs) { return lhs.name == rhs.name; } +// NOLINTNEXTLINE(modernize-use-nullptr) +bool operator<(const StreamData &lhs, const StreamData &rhs) { return lhs.name < rhs.name; } + +bool operator==(const StreamData &stream, const std::string &stream_name) { return stream.name == stream_name; } +// NOLINTNEXTLINE(modernize-use-nullptr) +bool operator<(const StreamData &stream, const std::string &stream_name) { return stream.name < stream_name; } + +Streams::Streams(InterpreterContext *interpreter_context, std::string bootstrap_servers, + std::filesystem::path directory) + : interpreter_context_(interpreter_context), + bootstrap_servers_(std::move(bootstrap_servers)), + storage_(std::move(directory)) {} + +void Streams::RestoreStreams() { + spdlog::info("Loading streams..."); + auto accessor = streams_.access(); + MG_ASSERT(accessor.size() == 0, "Cannot restore streams when some streams already exist!"); + + for (const auto &[stream_name, stream_data] : storage_) { + const auto get_failed_message = [](const std::string_view stream_name, const std::string_view message, + const std::string_view nested_message) { + return fmt::format("Failed to load stream '{}', because: {} caused by {}", stream_name, message, nested_message); + }; + + StreamStatus status; + try { + nlohmann::json::parse(stream_data).get_to(status); + } catch (const nlohmann::json::type_error &exception) { + spdlog::warn(get_failed_message(stream_name, "invalid type conversion", exception.what())); + continue; + } catch (const nlohmann::json::out_of_range &exception) { + spdlog::warn(get_failed_message(stream_name, "non existing field", exception.what())); + continue; + } + MG_ASSERT(status.name == stream_name, "Expected stream name is '{}', but got '{}'", stream_name, status.name); + + try { + CreateConsumer(accessor, stream_name, std::move(status.info), status.is_running, false); + } catch (const utils::BasicException &exception) { + spdlog::warn(get_failed_message(stream_name, "unexpected error", exception.what())); + } + } +} + +void Streams::Create(const std::string &stream_name, StreamInfo info) { + auto accessor = streams_.access(); + CreateConsumer(accessor, stream_name, std::move(info), false, true); +} + +void Streams::Drop(const std::string &stream_name) { + auto accessor = streams_.access(); + + if (!accessor.remove(stream_name)) { + throw StreamsException("Couldn't find stream '{}'", stream_name); + } + + if (!storage_.Delete(stream_name)) { + throw StreamsException("Couldn't delete stream '{}' from persistent store!", stream_name); + } + + // TODO(antaljanosbenjamin) Release the transformation +} + +void Streams::Start(const std::string &stream_name) { + auto accessor = streams_.access(); + auto it = GetStream(accessor, stream_name); + + auto locked_consumer = it->consumer->Lock(); + locked_consumer->Start(); + + Persist(CreateStatus(it->name, it->transformation_name, *locked_consumer)); +} + +void Streams::Stop(const std::string &stream_name) { + auto accessor = streams_.access(); + auto it = GetStream(accessor, stream_name); + + auto locked_consumer = it->consumer->Lock(); + locked_consumer->Stop(); + + Persist(CreateStatus(it->name, it->transformation_name, *locked_consumer)); +} + +void Streams::StartAll() { + for (auto &stream_data : streams_.access()) { + stream_data.consumer->WithLock([this, &stream_data](auto &consumer) { + if (!consumer.IsRunning()) { + consumer.Start(); + Persist(CreateStatus(stream_data.name, stream_data.transformation_name, consumer)); + } + }); + } +} + +void Streams::StopAll() { + for (auto &stream_data : streams_.access()) { + stream_data.consumer->WithLock([this, &stream_data](auto &consumer) { + if (consumer.IsRunning()) { + consumer.Stop(); + Persist(CreateStatus(stream_data.name, stream_data.transformation_name, consumer)); + } + }); + } +} + +std::vector Streams::Show() const { + std::vector result; + { + for (const auto &stream_data : streams_.access()) { + // Create string + result.emplace_back( + CreateStatus(stream_data.name, stream_data.transformation_name, *stream_data.consumer->ReadLock())); + } + } + return result; +} + +TransformationResult Streams::Test(const std::string &stream_name, std::optional batch_limit) { + auto accessor = streams_.access(); + auto it = GetStream(accessor, stream_name); + TransformationResult result; + auto consumer_function = [&result](const std::vector &messages) { + for (const auto &message : messages) { + // TODO(antaljanosbenjamin) Update the logic with using the transform from modules + const auto payload = message.Payload(); + const std::string_view payload_as_string_view{payload.data(), payload.size()}; + result[fmt::format("CREATE (n:MESSAGE {{payload: '{}'}})", payload_as_string_view)] = "replace with params"; + } + }; + + it->consumer->Lock()->Test(batch_limit, consumer_function); + + return result; +} + +StreamStatus Streams::CreateStatus(const std::string &name, const std::string &transformation_name, + const integrations::kafka::Consumer &consumer) { + const auto &info = consumer.Info(); + return StreamStatus{name, + StreamInfo{ + info.topics, + info.consumer_group, + info.batch_interval, + info.batch_size, + transformation_name, + }, + consumer.IsRunning()}; +} + +void Streams::CreateConsumer(utils::SkipList::Accessor &accessor, const std::string &stream_name, + StreamInfo info, const bool start_consumer, const bool persist_consumer) { + if (accessor.contains(stream_name)) { + throw StreamsException{"Stream already exists with name '{}'", stream_name}; + } + + auto consumer_function = [interpreter_context = + interpreter_context_](const std::vector &messages) { + Interpreter interpreter = Interpreter{interpreter_context}; + TransformationResult result; + + for (const auto &message : messages) { + // TODO(antaljanosbenjamin) Update the logic with using the transform from modules + const auto payload = message.Payload(); + const std::string_view payload_as_string_view{payload.data(), payload.size()}; + result[fmt::format("CREATE (n:MESSAGE {{payload: '{}'}})", payload_as_string_view)] = "replace with params"; + } + + for (const auto &[query, params] : result) { + // auto prepared_query = interpreter.Prepare(query, {}); + spdlog::info("Executing query '{}'", query); + // TODO(antaljanosbenjamin) run the query in real life, try not to copy paste the whole execution code, but + // extract it to a function that can be called from multiple places (e.g: triggers) + } + }; + + ConsumerInfo consumer_info{ + .consumer_name = stream_name, + .topics = std::move(info.topics), + .consumer_group = std::move(info.consumer_group), + .batch_interval = info.batch_interval, + .batch_size = info.batch_size, + }; + + auto consumer = std::make_unique(bootstrap_servers_, std::move(consumer_info), + std::move(consumer_function)); + auto locked_consumer = consumer->Lock(); + + if (start_consumer) { + locked_consumer->Start(); + } + if (persist_consumer) { + Persist(CreateStatus(stream_name, info.transformation_name, *locked_consumer)); + } + + auto insert_result = + accessor.insert(StreamData{stream_name, std::move(info.transformation_name), std::move(consumer)}); + MG_ASSERT(insert_result.second, "Unexpected error during storing consumer '{}'", stream_name); +} + +void Streams::Persist(StreamStatus &&status) { + const std::string stream_name = status.name; + if (!storage_.Put(stream_name, nlohmann::json(std::move(status)).dump())) { + throw StreamsException{"Couldn't persist steam data for stream '{}'", stream_name}; + } +} + +} // namespace query diff --git a/src/query/streams.hpp b/src/query/streams.hpp new file mode 100644 index 000000000..34f5bed45 --- /dev/null +++ b/src/query/streams.hpp @@ -0,0 +1,149 @@ +/// @file +#pragma once + +#include +#include +#include +#include + +#include "integrations/kafka/consumer.hpp" +#include "kvstore/kvstore.hpp" +#include "utils/exceptions.hpp" +#include "utils/rw_lock.hpp" +#include "utils/skip_list.hpp" +#include "utils/synchronized.hpp" + +namespace query { + +class StreamsException : public utils::BasicException { + public: + using BasicException::BasicException; +}; + +// TODO(antaljanosbenjamin) Replace this with mgp_trans related thing +using TransformationResult = std::map; +using TransformFunction = std::function &)>; + +struct StreamInfo { + std::vector topics; + std::string consumer_group; + std::optional batch_interval; + std::optional batch_size; + // TODO(antaljanosbenjamin) How to reference the transformation in a better way? + std::string transformation_name; +}; + +struct StreamStatus { + std::string name; + StreamInfo info; + bool is_running; +}; + +using SynchronizedConsumer = utils::Synchronized; + +struct StreamData { + std::string name; + // TODO(antaljanosbenjamin) How to reference the transformation in a better way? + std::string transformation_name; + // TODO(antaljanosbenjamin) consider propagate_const + std::unique_ptr consumer; +}; + +struct InterpreterContext; + +/// Manages Kafka consumers. +/// +/// This class is responsible for all query supported actions to happen. +class Streams final { + public: + /// Initializes the streams. + /// + /// @param interpreter_context context to use to run the result of transformations + /// @param bootstrap_servers initial list of brokers as a comma separated list of broker host or host:port + /// @param directory a directory path to store the persisted streams metadata + Streams(InterpreterContext *interpreter_context, std::string bootstrap_servers, std::filesystem::path directory); + + /// Restores the streams from the persisted metadata. + /// The restoration is done in a best effort manner, therefore no exception is thrown on failure, but the error is + /// logged. If a stream was running previously, then after restoration it will be started. + /// This function should only be called when there are no existing streams. + void RestoreStreams(); + + /// Creates a new import stream. + /// The create implies connecting to the server to get metadata necessary to initialize the stream. This + /// method assures there is no other stream with the same name. + /// + /// @param stream_name the name of the stream which can be used to uniquely identify the stream + /// @param stream_info the necessary informations needed to create the Kafka consumer and transform the messages + /// + /// @throws StreamsException if the stream with the same name exists or if the creation of Kafka consumer fails + void Create(const std::string &stream_name, StreamInfo stream_info); + + /// Deletes an existing stream and all the data that was persisted. + /// + /// @param stream_name name of the stream that needs to be deleted. + /// + /// @throws StreamsException if the stream doesn't exist or if the persisted metadata can't be deleted. + void Drop(const std::string &stream_name); + + /// Start consuming from a stream. + /// + /// @param stream_name name of the stream that needs to be started + /// + /// @throws StreamsException if the stream doesn't exist or if the metadata cannot be persisted + /// @throws ConsumerRunningException if the consumer is already running + void Start(const std::string &stream_name); + + /// Stop consuming from a stream. + /// + /// @param stream_name name of the stream that needs to be stopped + /// + /// @throws StreamsException if the stream doesn't exist or if the metadata cannot be persisted + /// @throws ConsumerStoppedException if the consumer is already stopped + void Stop(const std::string &stream_name); + + /// Start consuming from all streams that are stopped. + /// + /// @throws StreamsException if the metadata cannot be persisted + void StartAll(); + + /// Stop consuming from all streams that are running. + /// + /// @throws StreamsException if the metadata cannot be persisted + void StopAll(); + + /// Return current status for all streams. + /// It might happend that the is_running field is out of date if the one of the streams stops during the invocation of + /// this function because of an error. + std::vector Show() const; + + /// Do a dry-run consume from a stream. + /// + /// @param stream_name name of the stream we want to test + /// @param batch_limit number of batches we want to test before stopping + /// + /// TODO(antaljanosbenjamin) add type of parameters + /// @returns A vector of pairs consisting of the query (std::string) and its parameters ... + /// + /// @throws StreamsException if the stream doesn't exist + /// @throws ConsumerRunningException if the consumer is alredy running + /// @throws ConsumerTestFailedException if the transformation function throws any std::exception during processing + TransformationResult Test(const std::string &stream_name, std::optional batch_limit = std::nullopt); + + private: + static StreamStatus CreateStatus(const std::string &name, const std::string &transformation_name, + const integrations::kafka::Consumer &consumer); + + void CreateConsumer(utils::SkipList::Accessor &accessor, const std::string &stream_name, + StreamInfo stream_info, const bool start_consumer, const bool persist_consumer); + + void Persist(StreamStatus &&status); + + InterpreterContext *interpreter_context_; + std::string bootstrap_servers_; + kvstore::KVStore storage_; + + utils::SkipList streams_; +}; + +} // namespace query diff --git a/src/utils/rw_lock.hpp b/src/utils/rw_lock.hpp index d39015d00..997f2069f 100644 --- a/src/utils/rw_lock.hpp +++ b/src/utils/rw_lock.hpp @@ -110,4 +110,9 @@ class RWLock { pthread_rwlock_t lock_ = PTHREAD_RWLOCK_INITIALIZER; }; +class WritePrioritizedRWLock final : public RWLock { + public: + WritePrioritizedRWLock() : RWLock{Priority::WRITE} {}; +}; + } // namespace utils diff --git a/src/utils/synchronized.hpp b/src/utils/synchronized.hpp index 79ace2240..4cb3931dd 100644 --- a/src/utils/synchronized.hpp +++ b/src/utils/synchronized.hpp @@ -1,10 +1,20 @@ #pragma once +#include #include +#include #include namespace utils { +template +concept SharedMutex = requires(TMutex mutex) { + mutex.lock(); + mutex.unlock(); + mutex.lock_shared(); + mutex.unlock_shared(); +}; + /// A simple utility for easier mutex-based concurrency (influenced by /// Facebook's Folly) /// @@ -74,6 +84,21 @@ class Synchronized { std::lock_guard guard_; }; + class ReadLockedPtr { + private: + friend class Synchronized; + + ReadLockedPtr(const T *object_ptr, TMutex *mutex) : object_ptr_(object_ptr), guard_(*mutex) {} + + public: + const T *operator->() const { return object_ptr_; } + const T &operator*() const { return *object_ptr_; } + + private: + const T *object_ptr_; + std::shared_lock guard_; + }; + LockedPtr Lock() { return LockedPtr(&object_, &mutex_); } template @@ -83,9 +108,24 @@ class Synchronized { LockedPtr operator->() { return LockedPtr(&object_, &mutex_); } + template + requires SharedMutex ReadLockedPtr ReadLock() const { + return ReadLockedPtr(&object_, &mutex_); + } + + template + requires SharedMutex decltype(auto) WithReadLock(TCallable &&callable) const { + return callable(*ReadLock()); + } + + template + requires SharedMutex ReadLockedPtr operator->() const { + return ReadLockedPtr(&object_, &mutex_); + } + private: T object_; - TMutex mutex_; + mutable TMutex mutex_; }; } // namespace utils diff --git a/tests/benchmark/expansion.cpp b/tests/benchmark/expansion.cpp index 8bce6c781..7cf7d3246 100644 --- a/tests/benchmark/expansion.cpp +++ b/tests/benchmark/expansion.cpp @@ -37,7 +37,7 @@ class ExpansionBenchFixture : public benchmark::Fixture { MG_ASSERT(db->CreateIndex(label)); - interpreter_context.emplace(&*db, query::InterpreterConfig{}, data_directory); + interpreter_context.emplace(&*db, query::InterpreterConfig{}, data_directory, "non existing bootstrap servers"); interpreter.emplace(&*interpreter_context); } diff --git a/tests/manual/single_query.cpp b/tests/manual/single_query.cpp index e941948f9..3aea249af 100644 --- a/tests/manual/single_query.cpp +++ b/tests/manual/single_query.cpp @@ -17,7 +17,8 @@ int main(int argc, char *argv[]) { storage::Storage db; auto data_directory = std::filesystem::temp_directory_path() / "single_query_test"; utils::OnScopeExit([&data_directory] { std::filesystem::remove_all(data_directory); }); - query::InterpreterContext interpreter_context{&db, query::InterpreterConfig{}, data_directory}; + query::InterpreterContext interpreter_context{&db, query::InterpreterConfig{}, data_directory, + "non existing bootstrap servers"}; query::Interpreter interpreter{&interpreter_context}; ResultStreamFaker stream(&db); diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index 25c743e2e..15bff8230 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -44,6 +44,19 @@ function(_add_unit_test test_cpp custom_main) endfunction(_add_unit_test) +# Test integrations-kafka + +add_library(kafka-mock STATIC kafka_mock.cpp) +target_link_libraries(kafka-mock mg-utils librdkafka++ librdkafka Threads::Threads zlib gtest) +# Include directories are intentionally not set, because kafka-mock isn't meant to be used apart from unit tests + +add_unit_test(integrations_kafka_consumer.cpp kafka_mock.cpp) +target_link_libraries(${test_prefix}integrations_kafka_consumer kafka-mock mg-integrations-kafka) + +add_unit_test(mgp_kafka_c_api.cpp) +target_link_libraries(${test_prefix}mgp_kafka_c_api mg-query mg-integrations-kafka) + + # Test mg-query add_unit_test(bfs_single_node.cpp) @@ -101,6 +114,9 @@ target_link_libraries(${test_prefix}query_trigger mg-query) add_unit_test(query_serialization_property_value.cpp) target_link_libraries(${test_prefix}query_serialization_property_value mg-query) +add_unit_test(query_streams.cpp) +target_link_libraries(${test_prefix}query_streams mg-query kafka-mock) + # Test query/procedure add_unit_test(query_procedure_mgp_type.cpp) @@ -327,16 +343,3 @@ add_custom_command( add_custom_target(test_lcp ALL DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/test_lcp) add_test(test_lcp ${CMAKE_CURRENT_BINARY_DIR}/test_lcp) add_dependencies(memgraph__unit test_lcp) - - -# Test integrations-kafka - -add_library(kafka-mock STATIC kafka_mock.cpp) -target_link_libraries(kafka-mock mg-utils librdkafka++ librdkafka Threads::Threads zlib gtest) -# Include directories are intentionally not set, because kafka-mock isn't meant to be used apart from unit tests - -add_unit_test(integrations_kafka_consumer.cpp kafka_mock.cpp) -target_link_libraries(${test_prefix}integrations_kafka_consumer kafka-mock mg-integrations-kafka) - -add_unit_test(mgp_kafka_c_api.cpp) -target_link_libraries(${test_prefix}mgp_kafka_c_api mg-query mg-integrations-kafka) diff --git a/tests/unit/integrations_kafka_consumer.cpp b/tests/unit/integrations_kafka_consumer.cpp index 260584885..12a8c26d1 100644 --- a/tests/unit/integrations_kafka_consumer.cpp +++ b/tests/unit/integrations_kafka_consumer.cpp @@ -3,20 +3,22 @@ #include #include #include +#include #include +#include +#include #include #include -#include "gmock/gmock.h" -#include "gtest/gtest.h" #include "integrations/kafka/consumer.hpp" #include "integrations/kafka/exceptions.hpp" #include "kafka_mock.hpp" -#include "utils/timer.hpp" +#include "utils/string.hpp" using namespace integrations::kafka; namespace { +const auto kDummyConsumerFunction = [](const auto & /*messages*/) {}; int SpanToInt(std::span span) { int result{0}; if (span.size() != sizeof(int)) { @@ -33,9 +35,7 @@ struct ConsumerTest : public ::testing::Test { ConsumerInfo CreateDefaultConsumerInfo() const { const auto test_name = std::string{::testing::UnitTest::GetInstance()->current_test_info()->name()}; return ConsumerInfo{ - .consumer_function = [](const std::vector &) {}, .consumer_name = "Consumer" + test_name, - .bootstrap_servers = cluster.Bootstraps(), .topics = {kTopicName}, .consumer_group = "ConsumerGroup " + test_name, .batch_interval = std::nullopt, @@ -43,25 +43,32 @@ struct ConsumerTest : public ::testing::Test { }; }; - std::unique_ptr CreateConsumer(ConsumerInfo &&info) { - auto custom_consumer_function = std::move(info.consumer_function); + std::unique_ptr CreateConsumer(ConsumerInfo &&info, ConsumerFunction consumer_function) { + EXPECT_EQ(1, info.topics.size()); + EXPECT_EQ(info.topics.at(0), kTopicName); auto last_received_message = std::make_shared>(0); - info.consumer_function = [weak_last_received_message = std::weak_ptr{last_received_message}, - custom_consumer_function = - std::move(custom_consumer_function)](const std::vector &messages) { + const auto consumer_function_wrapper = [weak_last_received_message = std::weak_ptr{last_received_message}, + consumer_function = + std::move(consumer_function)](const std::vector &messages) { auto last_received_message = weak_last_received_message.lock(); + + EXPECT_FALSE(messages.empty()); + for (const auto &message : messages) { + EXPECT_EQ(message.TopicName(), kTopicName); + } if (last_received_message != nullptr) { *last_received_message = SpanToInt(messages.back().Payload()); } else { - custom_consumer_function(messages); + consumer_function(messages); } }; - auto consumer = std::make_unique(std::move(info)); + auto consumer = + std::make_unique(cluster.Bootstraps(), std::move(info), std::move(consumer_function_wrapper)); int sent_messages{1}; SeedTopicWithInt(kTopicName, sent_messages); - consumer->Start(std::nullopt); + consumer->Start(); if (!consumer->IsRunning()) { return nullptr; } @@ -104,15 +111,15 @@ TEST_F(ConsumerTest, BatchInterval) { std::vector> received_timestamps{}; info.batch_interval = kBatchInterval; auto expected_messages_received = true; - info.consumer_function = [&](const std::vector &messages) mutable { + auto consumer_function = [&](const std::vector &messages) mutable { received_timestamps.push_back({messages.size(), std::chrono::steady_clock::now()}); for (const auto &message : messages) { expected_messages_received &= (kMessage == std::string_view(message.Payload().data(), message.Payload().size())); } }; - auto consumer = CreateConsumer(std::move(info)); - consumer->Start(std::nullopt); + auto consumer = CreateConsumer(std::move(info), std::move(consumer_function)); + consumer->Start(); ASSERT_TRUE(consumer->IsRunning()); constexpr auto kMessageCount = 7; @@ -149,13 +156,13 @@ TEST_F(ConsumerTest, BatchInterval) { } TEST_F(ConsumerTest, StartStop) { - Consumer consumer{CreateDefaultConsumerInfo()}; + Consumer consumer{cluster.Bootstraps(), CreateDefaultConsumerInfo(), kDummyConsumerFunction}; auto start = [&consumer](const bool use_conditional) { if (use_conditional) { consumer.StartIfStopped(); } else { - consumer.Start(std::nullopt); + consumer.Start(); } }; @@ -178,7 +185,7 @@ TEST_F(ConsumerTest, StartStop) { start(use_conditional_start); EXPECT_TRUE(consumer.IsRunning()); - EXPECT_THROW(consumer.Start(std::nullopt), ConsumerRunningException); + EXPECT_THROW(consumer.Start(), ConsumerRunningException); consumer.StartIfStopped(); EXPECT_TRUE(consumer.IsRunning()); @@ -207,15 +214,15 @@ TEST_F(ConsumerTest, BatchSize) { info.batch_size = kBatchSize; constexpr std::string_view kMessage = "BatchSizeTestMessage"; auto expected_messages_received = true; - info.consumer_function = [&](const std::vector &messages) mutable { + auto consumer_function = [&](const std::vector &messages) mutable { received_timestamps.push_back({messages.size(), std::chrono::steady_clock::now()}); for (const auto &message : messages) { expected_messages_received &= (kMessage == std::string_view(message.Payload().data(), message.Payload().size())); } }; - auto consumer = CreateConsumer(std::move(info)); - consumer->Start(std::nullopt); + auto consumer = CreateConsumer(std::move(info), std::move(consumer_function)); + consumer->Start(); ASSERT_TRUE(consumer->IsRunning()); constexpr auto kLastBatchMessageCount = 1; @@ -259,14 +266,15 @@ TEST_F(ConsumerTest, BatchSize) { TEST_F(ConsumerTest, InvalidBootstrapServers) { auto info = CreateDefaultConsumerInfo(); - info.bootstrap_servers = "non.existing.host:9092"; - EXPECT_THROW(Consumer(std::move(info)), ConsumerFailedToInitializeException); + + EXPECT_THROW(Consumer("non.existing.host:9092", std::move(info), kDummyConsumerFunction), + ConsumerFailedToInitializeException); } TEST_F(ConsumerTest, InvalidTopic) { auto info = CreateDefaultConsumerInfo(); info.topics = {"Non existing topic"}; - EXPECT_THROW(Consumer(std::move(info)), TopicNotFoundException); + EXPECT_THROW(Consumer(cluster.Bootstraps(), std::move(info), kDummyConsumerFunction), TopicNotFoundException); } TEST_F(ConsumerTest, StartsFromPreviousOffset) { @@ -276,8 +284,9 @@ TEST_F(ConsumerTest, StartsFromPreviousOffset) { std::atomic received_message_count{0}; const std::string kMessagePrefix{"Message"}; auto expected_messages_received = true; - info.consumer_function = [&](const std::vector &messages) mutable { + auto consumer_function = [&](const std::vector &messages) mutable { auto message_count = received_message_count.load(); + EXPECT_EQ(messages.size(), 1); for (const auto &message : messages) { std::string message_payload = kMessagePrefix + std::to_string(message_count++); expected_messages_received &= @@ -287,33 +296,34 @@ TEST_F(ConsumerTest, StartsFromPreviousOffset) { }; // This test depends on CreateConsumer starts and stops the consumer, so the offset is stored - auto consumer = CreateConsumer(std::move(info)); + auto consumer = CreateConsumer(std::move(info), std::move(consumer_function)); ASSERT_FALSE(consumer->IsRunning()); - constexpr auto kMessageCount = 4; - for (auto sent_messages = 0; sent_messages < kMessageCount; ++sent_messages) { - cluster.SeedTopic(kTopicName, std::string_view{kMessagePrefix + std::to_string(sent_messages)}); - } - - auto do_batches = [&](int64_t batch_count) { + auto send_and_consume_messages = [&](int batch_count) { SCOPED_TRACE(fmt::format("Already received messages: {}", received_message_count.load())); - consumer->Start(batch_count); + auto expected_total_messages = received_message_count + batch_count; + for (auto sent_messages = 0; sent_messages < batch_count; ++sent_messages) { + cluster.SeedTopic(kTopicName, + std::string_view{kMessagePrefix + std::to_string(received_message_count + sent_messages)}); + } + consumer->Start(); const auto start = std::chrono::steady_clock::now(); ASSERT_TRUE(consumer->IsRunning()); constexpr auto kMaxWaitTime = std::chrono::seconds(5); - while (consumer->IsRunning() && (std::chrono::steady_clock::now() - start) < kMaxWaitTime) { + while (expected_total_messages != received_message_count.load() && + (std::chrono::steady_clock::now() - start) < kMaxWaitTime) { std::this_thread::sleep_for(std::chrono::milliseconds(200)); } // it is stopped because of limited batches + EXPECT_EQ(expected_total_messages, received_message_count); + consumer->Stop(); ASSERT_FALSE(consumer->IsRunning()); + EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received"; }; - ASSERT_NO_FATAL_FAILURE(do_batches(kMessageCount / 2)); - ASSERT_NO_FATAL_FAILURE(do_batches(kMessageCount / 2)); - - EXPECT_TRUE(expected_messages_received) << "Some unexpected message have been received"; - EXPECT_EQ(received_message_count, kMessageCount); + ASSERT_NO_FATAL_FAILURE(send_and_consume_messages(2)); + ASSERT_NO_FATAL_FAILURE(send_and_consume_messages(2)); } TEST_F(ConsumerTest, TestMethodWorks) { @@ -321,10 +331,10 @@ TEST_F(ConsumerTest, TestMethodWorks) { auto info = CreateDefaultConsumerInfo(); info.batch_size = kBatchSize; const std::string kMessagePrefix{"Message"}; - info.consumer_function = [](const std::vector &messages) mutable {}; + auto consumer_function = [](const std::vector &messages) mutable {}; // This test depends on CreateConsumer starts and stops the consumer, so the offset is stored - auto consumer = CreateConsumer(std::move(info)); + auto consumer = CreateConsumer(std::move(info), std::move(consumer_function)); constexpr auto kMessageCount = 4; for (auto sent_messages = 0; sent_messages < kMessageCount; ++sent_messages) { @@ -362,3 +372,38 @@ TEST_F(ConsumerTest, TestMethodWorks) { EXPECT_NO_FATAL_FAILURE(check_test_method()); } } + +TEST_F(ConsumerTest, ConsumerStatus) { + const std::string kConsumerName = "ConsumerGroupNameAAAA"; + const std::vector topics = {"Topic1QWER", "Topic2XCVBB"}; + const std::string kConsumerGroupName = "ConsumerGroupTestAsdf"; + constexpr auto kBatchInterval = std::chrono::milliseconds{111}; + constexpr auto kBatchSize = 222; + + for (const auto &topic : topics) { + cluster.CreateTopic(topic); + } + + auto check_info = [&](const ConsumerInfo &info) { + EXPECT_EQ(kConsumerName, info.consumer_name); + EXPECT_EQ(kConsumerGroupName, info.consumer_group); + EXPECT_EQ(kBatchInterval, info.batch_interval); + EXPECT_EQ(kBatchSize, info.batch_size); + EXPECT_EQ(2, info.topics.size()) << utils::Join(info.topics, ","); + ASSERT_LE(2, info.topics.size()); + EXPECT_EQ(topics[0], info.topics[0]); + EXPECT_EQ(topics[1], info.topics[1]); + }; + + Consumer consumer{cluster.Bootstraps(), + ConsumerInfo{kConsumerName, topics, kConsumerGroupName, kBatchInterval, kBatchSize}, + kDummyConsumerFunction}; + + check_info(consumer.Info()); + consumer.Start(); + check_info(consumer.Info()); + consumer.StartIfStopped(); + check_info(consumer.Info()); + consumer.StopIfRunning(); + check_info(consumer.Info()); +} diff --git a/tests/unit/interpreter.cpp b/tests/unit/interpreter.cpp index ef352eaa6..6080ecccb 100644 --- a/tests/unit/interpreter.cpp +++ b/tests/unit/interpreter.cpp @@ -30,7 +30,8 @@ auto ToEdgeList(const communication::bolt::Value &v) { struct InterpreterFaker { explicit InterpreterFaker(storage::Storage *db, const query::InterpreterConfig config, const std::filesystem::path &data_directory) - : interpreter_context(db, config, data_directory), interpreter(&interpreter_context) {} + : interpreter_context(db, config, data_directory, "not used bootstrap servers"), + interpreter(&interpreter_context) {} auto Prepare(const std::string &query, const std::map ¶ms = {}) { ResultStreamFaker stream(interpreter_context.db); diff --git a/tests/unit/kafka_mock.cpp b/tests/unit/kafka_mock.cpp index ecdfdc928..b0e8a818b 100644 --- a/tests/unit/kafka_mock.cpp +++ b/tests/unit/kafka_mock.cpp @@ -51,18 +51,22 @@ KafkaClusterMock::KafkaClusterMock(const std::vector &topics) { } for (const auto &topic : topics) { - constexpr auto partition_count = 1; - constexpr auto replication_factor = 1; - rd_kafka_resp_err_t topic_err = - rd_kafka_mock_topic_create(cluster_.get(), topic.c_str(), partition_count, replication_factor); - if (RD_KAFKA_RESP_ERR_NO_ERROR != topic_err) { - throw std::runtime_error("Failed to create the mock topic (" + topic + "): " + rd_kafka_err2str(topic_err)); - } + CreateTopic(topic); } }; std::string KafkaClusterMock::Bootstraps() const { return rd_kafka_mock_cluster_bootstraps(cluster_.get()); }; +void KafkaClusterMock::CreateTopic(const std::string &topic_name) { + constexpr auto partition_count = 1; + constexpr auto replication_factor = 1; + rd_kafka_resp_err_t topic_err = + rd_kafka_mock_topic_create(cluster_.get(), topic_name.c_str(), partition_count, replication_factor); + if (RD_KAFKA_RESP_ERR_NO_ERROR != topic_err) { + throw std::runtime_error("Failed to create the mock topic (" + topic_name + "): " + rd_kafka_err2str(topic_err)); + } +} + void KafkaClusterMock::SeedTopic(const std::string &topic_name, std::string_view message) { SeedTopic(topic_name, std::span{message.data(), message.size()}); } diff --git a/tests/unit/kafka_mock.hpp b/tests/unit/kafka_mock.hpp index ab0838266..a04c75bfe 100644 --- a/tests/unit/kafka_mock.hpp +++ b/tests/unit/kafka_mock.hpp @@ -28,6 +28,7 @@ class KafkaClusterMock { explicit KafkaClusterMock(const std::vector &topics); std::string Bootstraps() const; + void CreateTopic(const std::string &topic_name); void SeedTopic(const std::string &topic_name, std::span message); void SeedTopic(const std::string &topic_name, std::string_view message); diff --git a/tests/unit/query_dump.cpp b/tests/unit/query_dump.cpp index 26c6c8f8f..58f86491b 100644 --- a/tests/unit/query_dump.cpp +++ b/tests/unit/query_dump.cpp @@ -190,7 +190,7 @@ DatabaseState GetState(storage::Storage *db) { auto Execute(storage::Storage *db, const std::string &query) { auto data_directory = std::filesystem::temp_directory_path() / "MG_tests_unit_query_dump"; - query::InterpreterContext context(db, query::InterpreterConfig{}, data_directory); + query::InterpreterContext context(db, query::InterpreterConfig{}, data_directory, "non existing bootstrap servers"); query::Interpreter interpreter(&context); ResultStreamFaker stream(db); @@ -704,7 +704,9 @@ TEST(DumpTest, ExecuteDumpDatabase) { class StatefulInterpreter { public: explicit StatefulInterpreter(storage::Storage *db) - : db_(db), context_(db_, query::InterpreterConfig{}, data_directory_), interpreter_(&context_) {} + : db_(db), + context_(db_, query::InterpreterConfig{}, data_directory_, "non existing bootstrap servers"), + interpreter_(&context_) {} auto Execute(const std::string &query) { ResultStreamFaker stream(db_); diff --git a/tests/unit/query_plan_edge_cases.cpp b/tests/unit/query_plan_edge_cases.cpp index 1a1df047f..4fb478d5f 100644 --- a/tests/unit/query_plan_edge_cases.cpp +++ b/tests/unit/query_plan_edge_cases.cpp @@ -24,7 +24,7 @@ class QueryExecution : public testing::Test { void SetUp() { db_.emplace(); - interpreter_context_.emplace(&*db_, query::InterpreterConfig{}, data_directory); + interpreter_context_.emplace(&*db_, query::InterpreterConfig{}, data_directory, "non existing bootstrap servers"); interpreter_.emplace(&*interpreter_context_); } diff --git a/tests/unit/query_streams.cpp b/tests/unit/query_streams.cpp new file mode 100644 index 000000000..db5175fe2 --- /dev/null +++ b/tests/unit/query_streams.cpp @@ -0,0 +1,231 @@ +#include +#include +#include +#include + +#include +#include "kafka_mock.hpp" +#include "query/config.hpp" +#include "query/interpreter.hpp" +#include "query/streams.hpp" +#include "storage/v2/storage.hpp" + +using Streams = query::Streams; +using StreamInfo = query::StreamInfo; +using StreamStatus = query::StreamStatus; +namespace { +const static std::string kTopicName{"TrialTopic"}; + +struct StreamCheckData { + std::string name; + StreamInfo info; + bool is_running; +}; + +std::string GetDefaultStreamName() { + return std::string{::testing::UnitTest::GetInstance()->current_test_info()->name()}; +} + +StreamInfo CreateDefaultStreamInfo() { + return StreamInfo{ + .topics = {kTopicName}, + .consumer_group = "ConsumerGroup " + GetDefaultStreamName(), + .batch_interval = std::nullopt, + .batch_size = std::nullopt, + // TODO(antaljanosbenjamin) Add proper reference once Streams supports that + .transformation_name = "not yet used", + }; +} + +StreamCheckData CreateDefaultStreamCheckData() { return {GetDefaultStreamName(), CreateDefaultStreamInfo(), false}; } + +std::filesystem::path GetCleanDataDirectory() { + const auto path = std::filesystem::temp_directory_path() / "query-streams"; + std::filesystem::remove_all(path); + return path; +} +} // namespace + +class StreamsTest : public ::testing::Test { + public: + StreamsTest() { ResetStreamsObject(); } + + protected: + storage::Storage db_; + std::filesystem::path data_directory_{GetCleanDataDirectory()}; + KafkaClusterMock mock_cluster_{std::vector{kTopicName}}; + // Though there is a Streams object in interpreter context, it makes more sense to use a separate object to test, + // because that provides a way to recreate the streams object and also give better control over the arguments of the + // Streams constructor. + query::InterpreterContext interpreter_context_{&db_, query::InterpreterConfig{}, data_directory_, + "dont care bootstrap servers"}; + std::filesystem::path streams_data_directory_{data_directory_ / "separate-dir-for-test"}; + std::optional streams_; + + void ResetStreamsObject() { + streams_.emplace(&interpreter_context_, mock_cluster_.Bootstraps(), streams_data_directory_); + } + + void CheckStreamStatus(const StreamCheckData &check_data) { + SCOPED_TRACE(fmt::format("Checking status of '{}'", check_data.name)); + const auto &stream_statuses = streams_->Show(); + auto it = std::find_if(stream_statuses.begin(), stream_statuses.end(), + [&check_data](const auto &stream_status) { return stream_status.name == check_data.name; }); + ASSERT_NE(it, stream_statuses.end()); + const auto &status = *it; + // the order don't have to be strictly the same, but based on the implementation it shouldn't change + EXPECT_TRUE(std::equal(check_data.info.topics.begin(), check_data.info.topics.end(), status.info.topics.begin(), + status.info.topics.end())); + EXPECT_EQ(check_data.info.consumer_group, status.info.consumer_group); + EXPECT_EQ(check_data.info.batch_interval, status.info.batch_interval); + EXPECT_EQ(check_data.info.batch_size, status.info.batch_size); + // TODO(antaljanosbenjamin) Add proper reference once Streams supports that + // EXPECT_EQ(check_data.info.transformation_name, status.info.transformation_name); + EXPECT_EQ(check_data.is_running, status.is_running); + } + + void StartStream(StreamCheckData &check_data) { + streams_->Start(check_data.name); + check_data.is_running = true; + } + + void StopStream(StreamCheckData &check_data) { + streams_->Stop(check_data.name); + check_data.is_running = false; + } + + void Clear() { + if (!std::filesystem::exists(data_directory_)) return; + std::filesystem::remove_all(data_directory_); + } +}; + +TEST_F(StreamsTest, SimpleStreamManagement) { + auto check_data = CreateDefaultStreamCheckData(); + streams_->Create(check_data.name, check_data.info); + EXPECT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + + streams_->Start(check_data.name); + check_data.is_running = true; + EXPECT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + + streams_->StopAll(); + check_data.is_running = false; + EXPECT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + + streams_->StartAll(); + check_data.is_running = true; + EXPECT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + + streams_->Stop(check_data.name); + check_data.is_running = false; + EXPECT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + + streams_->Drop(check_data.name); + EXPECT_TRUE(streams_->Show().empty()); +} + +TEST_F(StreamsTest, CreateAlreadyExisting) { + auto stream_info = CreateDefaultStreamInfo(); + auto stream_name = GetDefaultStreamName(); + streams_->Create(stream_name, stream_info); + + try { + streams_->Create(stream_name, stream_info); + FAIL() << "Creating already existing stream should throw\n"; + } catch (query::StreamsException &exception) { + EXPECT_EQ(exception.what(), fmt::format("Stream already exists with name '{}'", stream_name)); + } +} + +TEST_F(StreamsTest, DropNotExistingStream) { + const auto stream_info = CreateDefaultStreamInfo(); + const auto stream_name = GetDefaultStreamName(); + const std::string not_existing_stream_name{"ThisDoesn'tExists"}; + streams_->Create(stream_name, stream_info); + + try { + streams_->Drop(not_existing_stream_name); + FAIL() << "Dropping not existing stream should throw\n"; + } catch (query::StreamsException &exception) { + EXPECT_EQ(exception.what(), fmt::format("Couldn't find stream '{}'", not_existing_stream_name)); + } +} + +TEST_F(StreamsTest, RestoreStreams) { + std::array stream_check_datas{ + CreateDefaultStreamCheckData(), + CreateDefaultStreamCheckData(), + CreateDefaultStreamCheckData(), + CreateDefaultStreamCheckData(), + }; + + // make the stream infos unique + for (auto i = 0; i < stream_check_datas.size(); ++i) { + auto &stream_check_data = stream_check_datas[i]; + auto &stream_info = stream_check_data.info; + auto iteration_postfix = std::to_string(i); + + stream_check_data.name += iteration_postfix; + stream_info.topics[0] += iteration_postfix; + stream_info.consumer_group += iteration_postfix; + stream_info.transformation_name += iteration_postfix; + if (i > 0) { + stream_info.batch_interval = std::chrono::milliseconds((i + 1) * 10); + stream_info.batch_size = 1000 + i; + } + + mock_cluster_.CreateTopic(stream_info.topics[0]); + } + stream_check_datas[1].info.batch_interval = {}; + stream_check_datas[2].info.batch_size = {}; + + const auto check_restore_logic = [&stream_check_datas, this]() { + // Reset the Streams object to trigger reloading + ResetStreamsObject(); + EXPECT_TRUE(streams_->Show().empty()); + streams_->RestoreStreams(); + EXPECT_EQ(stream_check_datas.size(), streams_->Show().size()); + for (const auto &check_data : stream_check_datas) { + ASSERT_NO_FATAL_FAILURE(CheckStreamStatus(check_data)); + } + }; + + streams_->RestoreStreams(); + EXPECT_TRUE(streams_->Show().empty()); + + for (auto &check_data : stream_check_datas) { + streams_->Create(check_data.name, check_data.info); + } + { + SCOPED_TRACE("After streams are created"); + check_restore_logic(); + } + + for (auto &check_data : stream_check_datas) { + StartStream(check_data); + } + { + SCOPED_TRACE("After starting streams"); + check_restore_logic(); + } + + // Stop two of the streams + StopStream(stream_check_datas[1]); + StopStream(stream_check_datas[3]); + { + SCOPED_TRACE("After stopping two streams"); + check_restore_logic(); + } + + // Stop the rest of the streams + StopStream(stream_check_datas[0]); + StopStream(stream_check_datas[2]); + check_restore_logic(); + { + SCOPED_TRACE("After stopping all streams"); + check_restore_logic(); + } +} + +// TODO(antaljanosbenjamin) Add tests for Streams::Test method and transformation diff --git a/tests/unit/utils_synchronized.cpp b/tests/unit/utils_synchronized.cpp index 2b51a688e..08d0c60b5 100644 --- a/tests/unit/utils_synchronized.cpp +++ b/tests/unit/utils_synchronized.cpp @@ -2,8 +2,13 @@ #include "gtest/gtest.h" +#include "utils/rw_lock.hpp" #include "utils/synchronized.hpp" +static_assert(utils::SharedMutex, "std::shared_mutex must be considered as shared mutex!"); +static_assert(utils::SharedMutex, "utils::RWLock must be considered as shared mutex!"); +static_assert(!utils::SharedMutex, "std::mutex must not be considered as shared mutex!"); + class NoMoveNoCopy { public: NoMoveNoCopy(int, int) {} @@ -15,7 +20,6 @@ class NoMoveNoCopy { ~NoMoveNoCopy() = default; }; -// NOLINTNEXTLINE(hicpp-special-member-functions) TEST(Synchronized, Constructors) { { utils::Synchronized> vec; @@ -31,7 +35,6 @@ TEST(Synchronized, Constructors) { std::vector data = {1, 2, 3}; utils::Synchronized> vec(std::move(data)); // data is guaranteed by the standard to be empty after move - // NOLINTNEXTLINE(bugprone-use-after-move, hicpp-invalid-access-moved) EXPECT_TRUE(data.empty()); EXPECT_EQ(vec->size(), 3); } @@ -52,11 +55,33 @@ class TestLock { } }; -// NOLINTNEXTLINE(hicpp-special-member-functions) -TEST(Synchronized, Usage) { - utils::Synchronized, TestLock> my_vector; +bool test_shared_lock_locked = false; + +class TestSharedLock { + public: + void lock() { + ASSERT_FALSE(test_lock_locked); + test_lock_locked = true; + } + void unlock() { + ASSERT_TRUE(test_lock_locked); + test_lock_locked = false; + } + void lock_shared() { + ASSERT_FALSE(test_shared_lock_locked); + test_shared_lock_locked = true; + } + void unlock_shared() { + ASSERT_TRUE(test_shared_lock_locked); + test_shared_lock_locked = false; + } +}; + +template +void CheckWriteLock(TSynchronizedVector &my_vector) { + ASSERT_TRUE(my_vector->empty()) << "Cannot use not empty vector"; { - // LockedPtr + SCOPED_TRACE("LockedPtr"); auto ptr = my_vector.Lock(); ASSERT_TRUE(test_lock_locked); ptr->push_back(5); @@ -64,12 +89,12 @@ TEST(Synchronized, Usage) { ASSERT_FALSE(test_lock_locked); { - // Indirection operator + SCOPED_TRACE("Indirection operator"); my_vector->push_back(6); } { - // Lambda + SCOPED_TRACE("Lambda"); my_vector.WithLock([](auto &my_vector) { ASSERT_TRUE(test_lock_locked); EXPECT_EQ(my_vector.size(), 2); @@ -79,3 +104,88 @@ TEST(Synchronized, Usage) { ASSERT_FALSE(test_lock_locked); } } + +template +void CheckReadLock(TSynchronizedVector &my_vector, const size_t expected_size) { + { + SCOPED_TRACE("ReadLockPtr"); + auto ptr = my_vector.ReadLock(); + ASSERT_TRUE(test_shared_lock_locked); + EXPECT_EQ(expected_size, ptr->size()); + } + ASSERT_FALSE(test_shared_lock_locked); + + { + auto ptr1 = my_vector.ReadLock(); + ASSERT_TRUE(test_shared_lock_locked); + + { + test_shared_lock_locked = false; + auto ptr2 = my_vector.ReadLock(); + EXPECT_EQ(expected_size, ptr1->size()); + EXPECT_EQ(expected_size, ptr2->size()); + } + test_shared_lock_locked = true; + } + ASSERT_FALSE(test_shared_lock_locked); + + { + SCOPED_TRACE("Indirection operator"); + EXPECT_EQ(expected_size, my_vector->size()); + } + { + SCOPED_TRACE("Indirection operator with other ReadLock"); + + [[maybe_unused]] auto ptr = my_vector.ReadLock(); + ASSERT_TRUE(test_shared_lock_locked); + test_shared_lock_locked = false; + EXPECT_EQ(expected_size, my_vector->size()); + ASSERT_FALSE(test_shared_lock_locked); + test_shared_lock_locked = true; + } + ASSERT_FALSE(test_shared_lock_locked); + + { + SCOPED_TRACE("Lambda"); + my_vector.WithReadLock([&expected_size](auto &my_vector) { + ASSERT_TRUE(test_shared_lock_locked); + EXPECT_EQ(my_vector.size(), expected_size); + }); + ASSERT_FALSE(test_shared_lock_locked); + } + { + SCOPED_TRACE("Lambda with other ReadLock"); + auto ptr1 = my_vector.ReadLock(); + ASSERT_TRUE(test_shared_lock_locked); + + { + test_shared_lock_locked = false; + my_vector.WithReadLock([&expected_size](const auto &my_vector) { + ASSERT_TRUE(test_shared_lock_locked); + EXPECT_EQ(my_vector.size(), expected_size); + }); + ASSERT_FALSE(test_shared_lock_locked); + } + test_shared_lock_locked = true; + } + ASSERT_FALSE(test_shared_lock_locked); +} + +TEST(Synchronized, Usage) { + utils::Synchronized, TestLock> my_vector; + CheckWriteLock(my_vector); +} + +TEST(Synchronized, SharedUsage) { + utils::Synchronized, TestSharedLock> my_vector; + CheckWriteLock(my_vector); + { + SCOPED_TRACE("Non const reference"); + ASSERT_NO_FATAL_FAILURE(CheckReadLock(my_vector, 2)); + } + { + const utils::Synchronized, TestSharedLock> &my_const_vector = my_vector; + SCOPED_TRACE("Const reference"); + ASSERT_NO_FATAL_FAILURE(CheckReadLock(my_const_vector, 2)); + } +} \ No newline at end of file